| """ |
| 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 |
|
|
| |
| 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}"} |
|
|
|
|
| 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() |
|
|
| |
| self.client = AsyncOpenAI( |
| api_key=settings.openai_api_key, |
| base_url=settings.openai_base_url, |
| ) |
| |
| self.model = model or "gpt-4o-mini" |
|
|
| 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 |
| """ |
| |
| messages = [ |
| {"role": "system", "content": SYSTEM_PROMPT + f"\n\nCurrent User ID: {user_id}"} |
| ] |
|
|
| |
| if context and context.get("conversation_history"): |
| for msg in context["conversation_history"][-5:]: |
| messages.append({ |
| "role": msg["role"], |
| "content": msg["content"] |
| }) |
|
|
| |
| messages.append({"role": "user", "content": message}) |
|
|
| tool_calls_made = [] |
|
|
| |
| response = await self.client.chat.completions.create( |
| model=self.model, |
| messages=messages, |
| tools=TOOLS, |
| tool_choice="auto", |
| ) |
|
|
| assistant_message = response.choices[0].message |
|
|
| |
| while assistant_message.tool_calls: |
| |
| messages.append(assistant_message) |
|
|
| |
| for tool_call in assistant_message.tool_calls: |
| function_name = tool_call.function.name |
| arguments = json.loads(tool_call.function.arguments) |
|
|
| |
| if "user_id" not in arguments: |
| arguments["user_id"] = user_id |
|
|
| |
| result = execute_tool(function_name, arguments) |
|
|
| tool_calls_made.append({ |
| "name": function_name, |
| "arguments": arguments, |
| "result": result |
| }) |
|
|
| |
| messages.append({ |
| "role": "tool", |
| "tool_call_id": tool_call.id, |
| "content": json.dumps(result) |
| }) |
|
|
| |
| response = await self.client.chat.completions.create( |
| model=self.model, |
| messages=messages, |
| tools=TOOLS, |
| tool_choice="auto", |
| ) |
| assistant_message = response.choices[0].message |
|
|
| |
| 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) |
|
|
