update the code

app/agent/task_agent.py

@@ -1,123 +1,125 @@
-"""
-Task Agent implementation using OpenAI API directly with OpenRouter.
-
-Provides an AI agent that maps natural language to MCP tools for task management.
-"""
-import json
-from typing import Any
-
-from openai import AsyncOpenAI
-
-from .prompts import SYSTEM_PROMPT
-from ..mcp import tools as mcp_tools
-
-# Tool definitions for OpenAI function calling
-TOOLS = [
-    {
-        "type": "function",
-        "function": {
-            "name": "add_task",
-            "description": "Create a new task for a user",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "user_id": {"type": "string", "description": "The user's unique identifier"},
-                    "title": {"type": "string", "description": "The task title"},
-                    "description": {"type": "string", "description": "Optional task description"}
-                },
-                "required": ["user_id", "title"]
-            }
-        }
-    },
-    {
-        "type": "function",
-        "function": {
-            "name": "list_tasks",
-            "description": "Get all tasks for a user",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "user_id": {"type": "string", "description": "The user's unique identifier"}
-                },
-                "required": ["user_id"]
-            }
-        }
-    },
-    {
-        "type": "function",
-        "function": {
-            "name": "complete_task",
-            "description": "Mark a task as completed. IMPORTANT: You must first call list_tasks to get the task UUID, then use that UUID here.",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "user_id": {"type": "string", "description": "The user's unique identifier"},
-                    "task_id": {"type": "string", "description": "The task's UUID (e.g., 'd18627f4-73ff-4450-943e-451c24962de0'). NOT the task title!"}
-                },
-                "required": ["user_id", "task_id"]
-            }
-        }
-    },
-    {
-        "type": "function",
-        "function": {
-            "name": "update_task",
-            "description": "Update a task's title or description",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "user_id": {"type": "string", "description": "The user's unique identifier"},
-                    "task_id": {"type": "string", "description": "The task's unique identifier (UUID)"},
-                    "title": {"type": "string", "description": "New task title"},
-                    "description": {"type": "string", "description": "New task description"}
-                },
-                "required": ["user_id", "task_id"]
-            }
-        }
-    },
-    {
-        "type": "function",
-        "function": {
-            "name": "delete_task",
-            "description": "Permanently delete a task",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "user_id": {"type": "string", "description": "The user's unique identifier"},
-                    "task_id": {"type": "string", "description": "The task's unique identifier (UUID)"}
-                },
-                "required": ["user_id", "task_id"]
-            }
-        }
-    }
-]
-
+Task Agent implementation using OpenAI API directly.
+
+  Provides an AI agent that maps natural language to MCP tools for task management.
+  Uses OpenAI's gpt-4o-mini for reliable function calling/tool use support.
+  """
+  import json
+  from typing import Any
+
+  from openai import AsyncOpenAI
+
+  from .prompts import SYSTEM_PROMPT
+  from ..mcp import tools as mcp_tools
+
+  # Tool definitions for OpenAI function calling
+  TOOLS = [
+      {
+          "type": "function",
+          "function": {
+              "name": "add_task",
+              "description": "Create a new task for a user",
+              "parameters": {
+                  "type": "object",
+                  "properties": {
+                      "user_id": {"type": "string", "description": "The user's unique identifier"},
+                      "title": {"type": "string", "description": "The task title"},
+                      "description": {"type": "string", "description": "Optional task description"}
+                  },
+                  "required": ["user_id", "title"]
+              }
+          }
+      },
+      {
+          "type": "function",
+          "function": {
+              "name": "list_tasks",
+              "description": "Get all tasks for a user",
+              "parameters": {
+                  "type": "object",
+                  "properties": {
+                      "user_id": {"type": "string", "description": "The user's unique identifier"}
+                  },
+                  "required": ["user_id"]
+              }
+          }
+      },
+      {
+          "type": "function",
+          "function": {
+              "name": "complete_task",
+              "description": "Mark a task as completed. IMPORTANT: You must first call list_tasks to get the task UUID, then use that UUID here.", 
+              "parameters": {
+                  "type": "object",
+                  "properties": {
+                      "user_id": {"type": "string", "description": "The user's unique identifier"},
+                      "task_id": {"type": "string", "description": "The task's UUID (e.g., 'd18627f4-73ff-4450-943e-451c24962de0'). NOT the task   
+  title!"}
+                  },
+                  "required": ["user_id", "task_id"]
+              }
+          }
+      },
+      {
+          "type": "function",
+          "function": {
+              "name": "update_task",
+              "description": "Update a task's title or description",
+              "parameters": {
+                  "type": "object",
+                  "properties": {
+                      "user_id": {"type": "string", "description": "The user's unique identifier"},
+                      "task_id": {"type": "string", "description": "The task's unique identifier (UUID)"},
+                      "title": {"type": "string", "description": "New task title"},
+                      "description": {"type": "string", "description": "New task description"}
+                  },
+                  "required": ["user_id", "task_id"]
+              }
+          }
+      },
+      {
+          "type": "function",
+          "function": {
+              "name": "delete_task",
+              "description": "Permanently delete a task",
+              "parameters": {
+                  "type": "object",
+                  "properties": {
+                      "user_id": {"type": "string", "description": "The user's unique identifier"},
+                      "task_id": {"type": "string", "description": "The task's unique identifier (UUID)"}
+                  },
+                  "required": ["user_id", "task_id"]
+              }
+          }
+      }
+  ]
+
+
+  def execute_tool(name: str, arguments: dict) -> Any:
+      """Execute a tool by name with given arguments."""
+      if name == "add_task":
+          return mcp_tools.add_task(
+              arguments["user_id"],
+              arguments["title"],
+              arguments.get("description")
+          )
+      elif name == "list_tasks":
+          return mcp_tools.list_tasks(arguments["user_id"])
+      elif name == "complete_task":
+          return mcp_tools.complete_task(arguments["user_id"], arguments["task_id"])
+      elif name == "update_task":
+          return mcp_tools.update_task(
+              arguments["user_id"],
+              arguments["task_id"],
+              arguments.get("title"),
+              arguments.get("description")
+          )
+      elif name == "delete_task":
+          return mcp_tools.delete_task(arguments["user_id"], arguments["task_id"])
+      else:
+          return {"error": f"Unknown tool: {name}"}
 
-def execute_tool(name: str, arguments: dict) -> Any:
-    """Execute a tool by name with given arguments."""
-    if name == "add_task":
-        return mcp_tools.add_task(
-            arguments["user_id"],
-            arguments["title"],
-            arguments.get("description")
-        )
-    elif name == "list_tasks":
-        return mcp_tools.list_tasks(arguments["user_id"])
-    elif name == "complete_task":
-        return mcp_tools.complete_task(arguments["user_id"], arguments["task_id"])
-    elif name == "update_task":
-        return mcp_tools.update_task(
-            arguments["user_id"],
-            arguments["task_id"],
-            arguments.get("title"),
-            arguments.get("description")
-        )
-    elif name == "delete_task":
-        return mcp_tools.delete_task(arguments["user_id"], arguments["task_id"])
-    else:
-        return {"error": f"Unknown tool: {name}"}
 
- class TaskAgent:
+  class TaskAgent:
       """
       AI agent for task management via natural language.
       Uses OpenAI API directly for reliable tool/function calling support.
@@ -140,122 +142,100 @@ def execute_tool(name: str, arguments: dict) -> Any:
           )
           # Use gpt-4o-mini for fast, reliable tool calling
           self.model = model or "gpt-4o-mini"
-          
-# class TaskAgent:
-#     """
-#     AI agent for task management via natural language.
-#     Uses OpenAI-compatible API (OpenRouter) with function calling.
-#     """
-
-#     def __init__(self, model: str | None = None):
-#         """
-#         Initialize the TaskAgent.
-
-#         Args:
-#             model: Model to use (default: from config)
-#         """
-#         from ..config import get_settings
-#         settings = get_settings()
-
-#         self.client = AsyncOpenAI(
-#             api_key=settings.openrouter_api_key,
-#             base_url=settings.openrouter_base_url,
-#         )
-#         self.model = model or settings.openrouter_model
-
-    async def run(
-        self,
-        message: str,
-        user_id: str,
-        context: dict[str, Any] | None = None
-    ) -> dict[str, Any]:
-        """
-        Process a user message and return the agent's response.
-
-        Args:
-            message: User's natural language message
-            user_id: User identifier for tool calls
-            context: Optional additional context
-
-        Returns:
-            dict with response and tool_calls
-        """
-        # Build messages
-        messages = [
-            {"role": "system", "content": SYSTEM_PROMPT + f"\n\nCurrent User ID: {user_id}"}
-        ]
-
-        # Add conversation history if available
-        if context and context.get("conversation_history"):
-            for msg in context["conversation_history"][-5:]:
-                messages.append({
-                    "role": msg["role"],
-                    "content": msg["content"]
-                })
-
-        # Add current user message
-        messages.append({"role": "user", "content": message})
 
-        tool_calls_made = []
-
-        # Call the model
-        response = await self.client.chat.completions.create(
-            model=self.model,
-            messages=messages,
-            tools=TOOLS,
-            tool_choice="auto",
-        )
-
-        assistant_message = response.choices[0].message
-
-        # Handle tool calls if any
-        while assistant_message.tool_calls:
-            # Add assistant message with tool calls
-            messages.append(assistant_message)
-
-            # Execute each tool call
-            for tool_call in assistant_message.tool_calls:
-                function_name = tool_call.function.name
-                arguments = json.loads(tool_call.function.arguments)
-
-                # Inject user_id if not provided
-                if "user_id" not in arguments:
-                    arguments["user_id"] = user_id
-
-                # Execute the tool
-                result = execute_tool(function_name, arguments)
-
-                tool_calls_made.append({
-                    "name": function_name,
-                    "arguments": arguments,
-                    "result": result
-                })
-
-                # Add tool result to messages
-                messages.append({
-                    "role": "tool",
-                    "tool_call_id": tool_call.id,
-                    "content": json.dumps(result)
-                })
-
-            # Get next response
-            response = await self.client.chat.completions.create(
-                model=self.model,
-                messages=messages,
-                tools=TOOLS,
-                tool_choice="auto",
-            )
-            assistant_message = response.choices[0].message
-
-        # Get final text response
-        response_text = assistant_message.content or ""
+      async def run(
+          self,
+          message: str,
+          user_id: str,
+          context: dict[str, Any] | None = None
+      ) -> dict[str, Any]:
+          """
+          Process a user message and return the agent's response.
 
-        return {
-            "response": response_text,
-            "tool_calls": tool_calls_made
-        }
+          Args:
+              message: User's natural language message
+              user_id: User identifier for tool calls
+              context: Optional additional context
 
+          Returns:
+              dict with response and tool_calls
+          """
+          # Build messages
+          messages = [
+              {"role": "system", "content": SYSTEM_PROMPT + f"\n\nCurrent User ID: {user_id}"}
+          ]
+
+          # Add conversation history if available
+          if context and context.get("conversation_history"):
+              for msg in context["conversation_history"][-5:]:
+                  messages.append({
+                      "role": msg["role"],
+                      "content": msg["content"]
+                  })
+
+          # Add current user message
+          messages.append({"role": "user", "content": message})
+
+          tool_calls_made = []
+
+          # Call the model
+          response = await self.client.chat.completions.create(
+              model=self.model,
+              messages=messages,
+              tools=TOOLS,
+              tool_choice="auto",
+          )
 
-def create_task_agent(model: str | None = None) -> TaskAgent:
-    """Factory function to create a TaskAgent instance."""
-    return TaskAgent(model=model)
+          assistant_message = response.choices[0].message
+
+          # Handle tool calls if any
+          while assistant_message.tool_calls:
+              # Add assistant message with tool calls
+              messages.append(assistant_message)
+
+              # Execute each tool call
+              for tool_call in assistant_message.tool_calls:
+                  function_name = tool_call.function.name
+                  arguments = json.loads(tool_call.function.arguments)
+
+                  # Inject user_id if not provided
+                  if "user_id" not in arguments:
+                      arguments["user_id"] = user_id
+
+                  # Execute the tool
+                  result = execute_tool(function_name, arguments)
+
+                  tool_calls_made.append({
+                      "name": function_name,
+                      "arguments": arguments,
+                      "result": result
+                  })
+
+                  # Add tool result to messages
+                  messages.append({
+                      "role": "tool",
+                      "tool_call_id": tool_call.id,
+                      "content": json.dumps(result)
+                  })
+
+              # Get next response
+              response = await self.client.chat.completions.create(
+                  model=self.model,
+                  messages=messages,
+                  tools=TOOLS,
+                  tool_choice="auto",
+              )
+              assistant_message = response.choices[0].message
+
+          # Get final text response
+          response_text = assistant_message.content or ""
+
+          return {
+              "response": response_text,
+              "tool_calls": tool_calls_made
+          }
+
+
+  def create_task_agent(model: str | None = None) -> TaskAgent:
+      """Factory function to create a TaskAgent instance."""
+      return TaskAgent(model=model)