update the code

app/agent/task_agent.py

@@ -1,7 +1,8 @@
-Task Agent implementation using OpenAI API directly.
+ """                                                                                                                                              
+  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.
+  Uses OpenAI gpt-4o-mini for reliable function calling/tool use support.
   """
   import json
   from typing import Any
@@ -11,7 +12,6 @@ Task Agent implementation using OpenAI API directly.
   from .prompts import SYSTEM_PROMPT
   from ..mcp import tools as mcp_tools
 
-  # Tool definitions for OpenAI function calling
   TOOLS = [
       {
           "type": "function",
@@ -21,7 +21,7 @@ Task Agent implementation using OpenAI API directly.
               "parameters": {
                   "type": "object",
                   "properties": {
-                      "user_id": {"type": "string", "description": "The user's unique identifier"},
+                      "user_id": {"type": "string", "description": "The user unique identifier"},
                       "title": {"type": "string", "description": "The task title"},
                       "description": {"type": "string", "description": "Optional task description"}
                   },
@@ -37,7 +37,7 @@ Task Agent implementation using OpenAI API directly.
               "parameters": {
                   "type": "object",
                   "properties": {
-                      "user_id": {"type": "string", "description": "The user's unique identifier"}
+                      "user_id": {"type": "string", "description": "The user unique identifier"}
                   },
                   "required": ["user_id"]
               }
@@ -47,13 +47,12 @@ Task Agent implementation using OpenAI API directly.
           "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.", 
+              "description": "Mark a task as completed. You must first call list_tasks to get the task UUID.",
               "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!"}
+                      "user_id": {"type": "string", "description": "The user unique identifier"},
+                      "task_id": {"type": "string", "description": "The task UUID. NOT the task title!"}
                   },
                   "required": ["user_id", "task_id"]
               }
@@ -63,12 +62,12 @@ Task Agent implementation using OpenAI API directly.
           "type": "function",
           "function": {
               "name": "update_task",
-              "description": "Update a task's title or description",
+              "description": "Update a task 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)"},
+                      "user_id": {"type": "string", "description": "The user unique identifier"},
+                      "task_id": {"type": "string", "description": "The task unique identifier (UUID)"},
                       "title": {"type": "string", "description": "New task title"},
                       "description": {"type": "string", "description": "New task description"}
                   },
@@ -84,8 +83,8 @@ Task Agent implementation using OpenAI API directly.
               "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)"}
+                      "user_id": {"type": "string", "description": "The user unique identifier"},
+                      "task_id": {"type": "string", "description": "The task unique identifier (UUID)"}
                   },
                   "required": ["user_id", "task_id"]
               }
@@ -95,7 +94,6 @@ Task Agent implementation using OpenAI API directly.
 
 
   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"],
@@ -120,27 +118,13 @@ Task Agent implementation using OpenAI API directly.
 
 
   class TaskAgent:
-      """
-      AI agent for task management via natural language.
-      Uses OpenAI API directly for reliable tool/function calling support.
-      """
-
       def __init__(self, model: str | None = None):
-          """
-          Initialize the TaskAgent.
-
-          Args:
-              model: Model to use (default: gpt-4o-mini)
-          """
           from ..config import get_settings
           settings = get_settings()
-
-          # Use OpenAI directly for reliable tool use support
           self.client = AsyncOpenAI(
               api_key=settings.openai_api_key,
               base_url=settings.openai_base_url,
           )
-          # Use gpt-4o-mini for fast, reliable tool calling
           self.model = model or "gpt-4o-mini"
 
       async def run(
@@ -149,76 +133,42 @@ Task Agent implementation using OpenAI API directly.
           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,
@@ -226,10 +176,7 @@ Task Agent implementation using OpenAI API directly.
                   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
@@ -237,5 +184,4 @@ Task Agent implementation using OpenAI API directly.
 
 
   def create_task_agent(model: str | None = None) -> TaskAgent:
-      """Factory function to create a TaskAgent instance."""
       return TaskAgent(model=model)