Added error handling

src/ai/anthropic.py

@@ -1,5 +1,11 @@
 import anthropic
+from anthropic.types import APIError, APITimeoutError, APIConnectionError
+from src.ai.tool import ToolError
+import logger
 
+class AnthropicError(Exception):
+    """Custom exception for anthropic-related errors."""
+    pass
 
 async def chat_completion(
     system_prompt: str,
@@ -8,24 +14,31 @@ async def chat_completion(
 ):
     client = anthropic.AsyncClient()
 
-    message = await client.messages.create(
-        model=model,
-        max_tokens=2000,
-        temperature=0,
-        system=system_prompt,
-        messages=[
-            {
-                "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": prompt,
-                    }
-                ],
-            }
-        ],
-    )
-    return message.content
+    try:
+        message = await client.messages.create(
+            model=model,
+            max_tokens=2000,
+            temperature=0,
+            system=system_prompt,
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": prompt,
+                        }
+                    ],
+                }
+            ],
+        )
+        return message.content
+    except (APIError, APITimeoutError, APIConnectionError) as e:
+        logger.log.error(f"Anthropic API error in chat_completion: {str(e)}")
+        raise AnthropicError(f"Anthropic API error: {str(e)}")
+    except Exception as e:
+        logger.log.error(f"Unexpected error in chat_completion: {str(e)}")
+        raise AnthropicError(f"Unexpected error: {str(e)}")
 
 
 async def tool_completion(
@@ -44,28 +57,35 @@ async def tool_completion(
     else:
         tool_choice = {"type": "auto"}
 
-    message = await client.messages.create(
-        model=model,
-        max_tokens=2000,
-        temperature=0,
-        system=system_prompt,
-        messages=[
-            {
-                "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": prompt,
-                    }
-                ],
-            }
-        ],
-        tools=tools,
-        tool_choice=tool_choice,
-    )
-    if message.stop_reason == "tool_use":
-        for response in message.content:
-            if response.type == "tool_use":
-                return response.input
-    else:
-        return message.content[0].text
+    try:
+        message = await client.messages.create(
+            model=model,
+            max_tokens=2000,
+            temperature=0,
+            system=system_prompt,
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": prompt,
+                        }
+                    ],
+                }
+            ],
+            tools=tools,
+            tool_choice=tool_choice,
+        )
+        if message.stop_reason == "tool_use":
+            for response in message.content:
+                if response.type == "tool_use":
+                    return response.input
+        else:
+            return message.content[0].text
+    except (APIError, APITimeoutError, APIConnectionError) as e:
+        logger.log.error(f"Anthropic API error in tool_completion: {str(e)}")
+        raise AnthropicError(f"Anthropic API error: {str(e)}")
+    except Exception as e:
+        logger.log.error(f"Unexpected error in tool_completion: {str(e)}")
+        raise AnthropicError(f"Unexpected error: {str(e)}")

src/ai/assistant.py

@@ -5,7 +5,7 @@
 import logger
 from . import prompt, schema
 from .anthropic import tool_completion
-from .tool import TOOLS
+from .tool import TOOLS, ToolError
 
 
 class Assistant:
@@ -21,26 +21,33 @@ async def overview(
         self,
         context: code.model.PullRequestContextModel,
     ) -> code.review.model.OverviewModel:
-        completion = await tool_completion(
-            system_prompt=self._builder.render_template(
-                name="overview",
-                prefix="system",
-                persona=self._persona,
-            ),
-            prompt=self._builder.render_template(
-                name="overview",
-                prefix="user",
-            ),
-            model=self._model_name,
-            tools=[TOOLS["review_files"]],
-            tool_override="review_files",
-        )
-        response = schema.ReviewRequestsResponseModel(**completion)
+        try:
+            completion = await tool_completion(
+                system_prompt=self._builder.render_template(
+                    name="overview",
+                    prefix="system",
+                    persona=self._persona,
+                ),
+                prompt=self._builder.render_template(
+                    name="overview",
+                    prefix="user",
+                ),
+                model=self._model_name,
+                tools=[TOOLS["review_files"]],
+                tool_override="review_files",
+            )
+            response = schema.ReviewRequestsResponseModel(**completion)
 
-        return code.review.comment.parse_overview(
-            response=response,
-            context=context,
-        )
+            return code.review.comment.parse_overview(
+                response=response,
+                context=context,
+            )
+        except ToolError as e:
+            logger.log.error(f"Error in overview: {str(e)}")
+            raise
+        except Exception as e:
+            logger.log.error(f"Unexpected error in overview: {str(e)}")
+            raise
 
     async def review_file(
         self,
@@ -53,54 +60,68 @@ async def review_file(
         logger.log.debug(f"Reviewing file: {context.path}")
         logger.log.debug(f"Hunks: {context.patch.hunks}")
 
-        completion = await tool_completion(
-            system_prompt=self._builder.render_template(
-                name="file-review",
-                prefix="system",
-                persona=self._persona,
-            ),
-            prompt=self._builder.render_template(
-                name="file-review",
-                prefix="user",
-                observations=observations,
-                file_context=context,
-            ),
-            model=self._model_name,
-            tools=[TOOLS["post_feedback"]],
-            tool_override="post_feedback",
-        )
-        response = schema.FileReviewResponseModel(**completion)
+        try:
+            completion = await tool_completion(
+                system_prompt=self._builder.render_template(
+                    name="file-review",
+                    prefix="system",
+                    persona=self._persona,
+                ),
+                prompt=self._builder.render_template(
+                    name="file-review",
+                    prefix="user",
+                    observations=observations,
+                    file_context=context,
+                ),
+                model=self._model_name,
+                tools=[TOOLS["post_feedback"]],
+                tool_override="post_feedback",
+            )
+            response = schema.FileReviewResponseModel(**completion)
 
-        return code.review.comment.extract_comments(
-            response=response,
-            file_context=context,
-            severity_limit=severity_limit,
-        )
+            return code.review.comment.extract_comments(
+                response=response,
+                file_context=context,
+                severity_limit=severity_limit,
+            )
+        except ToolError as e:
+            logger.log.error(f"Error in review_file for {context.path}: {str(e)}")
+            raise
+        except Exception as e:
+            logger.log.error(f"Unexpected error in review_file for {context.path}: {str(e)}")
+            raise
 
     async def get_feedback(
         self,
         prioritized_comments: list[code.model.GitHubCommentModel],
         remaining_comments: list[code.model.GitHubCommentModel],
     ) -> code.review.model.Feedback:
-        completion = await tool_completion(
-            system_prompt=self._builder.render_template(
-                name="review-summary",
-                prefix="system",
-                persona=self._persona,
-            ),
-            prompt=self._builder.render_template(
-                name="review-summary",
-                prefix="user",
-                prioritized_comments=prioritized_comments,
-                comments=remaining_comments,
-            ),
-            model=self._model_name,
-            tools=[TOOLS["submit_review"]],
-            tool_override="submit_review",
-        )
-        response = schema.ReviewResponseModel(**completion)
+        try:
+            completion = await tool_completion(
+                system_prompt=self._builder.render_template(
+                    name="review-summary",
+                    prefix="system",
+                    persona=self._persona,
+                ),
+                prompt=self._builder.render_template(
+                    name="review-summary",
+                    prefix="user",
+                    prioritized_comments=prioritized_comments,
+                    comments=remaining_comments,
+                ),
+                model=self._model_name,
+                tools=[TOOLS["submit_review"]],
+                tool_override="submit_review",
+            )
+            response = schema.ReviewResponseModel(**completion)
 
-        return code.review.comment.parse_feedback(
-            response=response,
-            comments=prioritized_comments,
-        )
+            return code.review.comment.parse_feedback(
+                response=response,
+                comments=prioritized_comments,
+            )
+        except ToolError as e:
+            logger.log.error(f"Error in get_feedback: {str(e)}")
+            raise
+        except Exception as e:
+            logger.log.error(f"Unexpected error in get_feedback: {str(e)}")
+            raise

src/ai/prompt.py

@@ -6,6 +6,10 @@
 import code.model
 import code.review.model
 
+class PromptError(Exception):
+    """Custom exception for prompt-related errors."""
+    pass
+
 PROMPT_DIR = pathlib.Path(__file__).parent / "prompts"
 
 
@@ -29,14 +33,26 @@ def _load_template(
         name: str,
         prefix: typing.Literal["system", "user", "persona"],
     ) -> jinja2.Template:
-        return self._templates[prefix].get_template(name)
+        try:
+            return self._templates[prefix].get_template(name)
+        except jinja2.TemplateNotFound:
+            raise PromptError(f"Template not found: {prefix}/{name}")
+        except jinja2.TemplateError as e:
+            raise PromptError(f"Error loading template {prefix}/{name}: {str(e)}")
 
     def render_template(
         self,
         name: str,
         prefix: typing.Literal["system", "user", "persona"],
         **kwargs,
     ) -> str:
-        template = self._load_template(f"{name}.md", prefix=prefix)
-        overview = template.render(context=self._context, prefix=prefix, **kwargs)
-        return overview
+        try:
+            template = self._load_template(f"{name}.md", prefix=prefix)
+            overview = template.render(context=self._context, prefix=prefix, **kwargs)
+            return overview
+        except PromptError as e:
+            raise e
+        except jinja2.TemplateError as e:
+            raise PromptError(f"Error rendering template {prefix}/{name}.md: {str(e)}")
+        except Exception as e:
+            raise PromptError(f"Unexpected error rendering template {prefix}/{name}.md: {str(e)}")

src/ai/tool.py

@@ -1,22 +1,37 @@
 import json
 import pathlib
-from typing import Any
+from typing import Any, Dict
 
-TOOL_DIR = pathlib.Path(__file__).parent / "tools"
+import logger
 
+TOOL_DIR = pathlib.Path(__file__).parent / "tools"
 
-def load_tool(name: str) -> dict[str, Any]:
-    path = TOOL_DIR / f"{name}.json"
-    with open(path, "r") as file:
-        return json.load(file)
+class ToolError(Exception):
+    """Custom exception for tool-related errors."""
+    pass
 
+def load_tool(name: str) -> Dict[str, Any]:
+    try:
+        path = TOOL_DIR / f"{name}.json"
+        with open(path, "r") as file:
+            return json.load(file)
+    except FileNotFoundError:
+        raise ToolError(f"Tool file not found: {name}.json")
+    except json.JSONDecodeError:
+        raise ToolError(f"Invalid JSON in tool file: {name}.json")
 
-def get_all_tools() -> dict[str, dict[str, Any]]:
+def get_all_tools() -> Dict[str, Dict[str, Any]]:
     tools = {}
     for file in TOOL_DIR.glob("*.json"):
         tool_name = file.stem
-        tools[tool_name] = load_tool(tool_name)
+        try:
+            tools[tool_name] = load_tool(tool_name)
+        except ToolError as e:
+            logger.log.error(f"Error loading tool {tool_name}: {str(e)}")
     return tools
 
-
-TOOLS = get_all_tools()
+try:
+    TOOLS = get_all_tools()
+except Exception as e:
+    logger.log.critical(f"Failed to load tools: {str(e)}")
+    TOOLS = {}