CoolFace
Apppublic

creativesar/taskflow

sourceHugging Faceupdated 8mo agoView on Hugging Face
0likes
agent.py276 linesDownload Raw Back to root
1"""2OpenAI Agents SDK Integration with OpenRouter Support3Phase III: AI Chatbot - Agent Definition and Runner4 5This module defines the TodoBot agent using OpenAI Agents SDK.6The agent helps users manage their tasks through natural language.7Supports both OpenAI and OpenRouter APIs.8"""9 10from typing import List, Dict, Any, Optional11from agents import Agent, Runner, function_tool12import os13 14# Import MCP tool implementations15from mcp_server import (16    add_task as mcp_add_task,17    list_tasks as mcp_list_tasks,18    complete_task as mcp_complete_task,19    delete_task as mcp_delete_task,20    update_task as mcp_update_task21)22 23# Configure OpenAI SDK environment variables for OpenRouter support24# The Agents SDK will automatically use these environment variables25def configure_openai_environment():26    """27    Configure OpenAI SDK environment variables.28    Supports both OpenAI and OpenRouter APIs.29 30    If OPENROUTER_API_KEY is set, configures the SDK to use OpenRouter.31    Otherwise, uses standard OpenAI configuration.32    """33    openrouter_key = os.getenv("OPENROUTER_API_KEY")34    openrouter_base_url = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")35 36    if openrouter_key:37        # Configure for OpenRouter38        os.environ["OPENAI_API_KEY"] = openrouter_key39        os.environ["OPENAI_BASE_URL"] = openrouter_base_url40 41        # Set max_tokens to stay within OpenRouter free tier (4000 tokens max)42        # This environment variable is read by the OpenAI SDK43        max_tokens = os.getenv("MAX_TOKENS", "1000")44        os.environ["OPENAI_MAX_COMPLETION_TOKENS"] = max_tokens45    else:46        # Use standard OpenAI configuration47        openai_key = os.getenv("OPENAI_API_KEY")48        if not openai_key:49            raise ValueError("Either OPENROUTER_API_KEY or OPENAI_API_KEY must be set")50 51# Configure environment before creating agent52configure_openai_environment()53 54 55# Global variable to store current user_id context56_current_user_id: Optional[str] = None57 58def set_user_context(user_id: str):59    """Set the current user context for tool execution."""60    global _current_user_id61    _current_user_id = user_id62 63def get_user_context() -> str:64    """Get the current user context."""65    global _current_user_id66    if _current_user_id is None:67        raise ValueError("User context not set")68    return _current_user_id69 70 71# Define function tools using @function_tool decorator72# These tools do NOT require user_id - it's automatically injected from context73@function_tool74async def add_task(title: str, description: Optional[str] = None) -> Dict[str, Any]:75    """76    Create a new task for the user.77 78    Args:79        title: Task title (1-200 characters)80        description: Optional task description (0-1000 characters)81 82    Returns:83        Dictionary with task_id, status, and title84    """85    user_id = get_user_context()86    return await mcp_add_task(user_id, title, description)87 88 89@function_tool90async def list_tasks(status: str = "all") -> List[Dict[str, Any]]:91    """92    Retrieve tasks from the user's list.93 94    Args:95        status: Filter by status ("all", "pending", "completed")96 97    Returns:98        List of task dictionaries99    """100    user_id = get_user_context()101    return await mcp_list_tasks(user_id, status)102 103 104@function_tool105async def complete_task(task_id: int) -> Dict[str, Any]:106    """107    Mark a task as complete or toggle completion status.108 109    Args:110        task_id: The ID of the task to complete111 112    Returns:113        Dictionary with task_id, status, and title114    """115    user_id = get_user_context()116    return await mcp_complete_task(user_id, task_id)117 118 119@function_tool120async def delete_task(task_id: int) -> Dict[str, Any]:121    """122    Remove a task from the list.123 124    Args:125        task_id: The ID of the task to delete126 127    Returns:128        Dictionary with task_id, status, and title129    """130    user_id = get_user_context()131    return await mcp_delete_task(user_id, task_id)132 133 134@function_tool135async def update_task(136    task_id: int,137    title: Optional[str] = None,138    description: Optional[str] = None139) -> Dict[str, Any]:140    """141    Modify task title or description.142 143    Args:144        task_id: The ID of the task to update145        title: New task title (1-200 characters)146        description: New task description (0-1000 characters)147 148    Returns:149        Dictionary with task_id, status, and title150    """151    user_id = get_user_context()152    return await mcp_update_task(user_id, task_id, title, description)153 154 155def get_agent_instructions() -> str:156    """157    Returns the agent's instruction set.158 159    The agent is designed to:160    - Understand natural language task management commands161    - Use MCP tools to perform CRUD operations on tasks162    - Provide friendly, helpful responses163    - Handle errors gracefully164    - Always include task IDs in responses165    """166    return """You are TaskFlowBot, a helpful task management assistant. You help users manage their tasks through natural language.167 168IMPORTANT: You do NOT need to ask for or provide user_id. The system automatically handles user authentication.169 170Use the available tools to perform actions:171- When user wants to add/remember something, use add_task tool with just title and description172- When user asks to see/show/list tasks, use list_tasks tool (optionally with status filter)173- When user says done/complete/finished, use complete_task tool with the task_id174- When user says delete/remove/cancel, use delete_task tool with the task_id175- When user says change/update/rename, use update_task tool with task_id and new values176 177CRITICAL - Always include task IDs in your responses:178- When adding a task, say "I've added Task #[ID]: [title]"179- When listing tasks, format each as "Task #[ID]: [title]"180- When completing/deleting/updating, say "Task #[ID] has been [action]"181- Remind users they can reference tasks by ID (e.g., "To delete it, say 'delete task 5'")182 183Always confirm actions with a friendly response that includes the task ID.184Handle errors gracefully and help user rephrase if needed.185Be concise but helpful in your responses."""186 187 188# Get model name from environment or use default189MODEL_NAME = os.getenv("AI_MODEL", "gpt-4o")  # OpenRouter supports gpt-4o190 191# Get max tokens from environment or use default (reduced for OpenRouter credit limits)192# OpenRouter free tier allows up to 4000 tokens193MAX_TOKENS = int(os.getenv("MAX_TOKENS", "1000"))  # Reduced to stay within free tier194 195# Create the TodoBot agent with all tools (T-311)196# OpenRouter support configured via environment variables above197# Token limits are controlled via OPENAI_MAX_COMPLETION_TOKENS environment variable198todobot_agent = Agent(199    name="TodoBot",200    instructions=get_agent_instructions(),201    model=MODEL_NAME,202    tools=[add_task, list_tasks, complete_task, delete_task, update_task]203)204 205 206# Agent Runner implementation (T-312)207async def run_agent(user_message: str, user_id: str, conversation_history: Optional[List[Dict[str, str]]] = None) -> Dict[str, Any]:208    """209    Run the TodoBot agent with a user message and optional conversation history.210 211    Args:212        user_message: The user's input message213        user_id: The authenticated user's ID (automatically injected into tool calls via context)214        conversation_history: Optional list of previous messages in format [{"role": "user"|"assistant", "content": "..."}]215 216    Returns:217        Dictionary containing:218        - response: The agent's text response219        - tool_calls: List of tool calls made (if any)220        - error: Error message if execution failed221 222    Example:223        result = await run_agent("Add a task to buy groceries", "user123")224        print(result["response"])225    """226    try:227        # Set user context for tool execution228        set_user_context(user_id)229 230        # Build messages array231        messages = conversation_history or []232        messages.append({"role": "user", "content": user_message})233 234        # Run the agent235        # Note: OpenAI Agents SDK doesn't support max_tokens parameter in Runner.run()236        # Token limits are controlled by the model configuration237        result = await Runner.run(238            todobot_agent,239            messages240        )241 242        # Extract response and tool calls243        return {244            "response": result.final_output,245            "tool_calls": getattr(result, "tool_calls", []),246            "error": None247        }248 249    except Exception as e:250        return {251            "response": f"I encountered an error: {str(e)}",252            "tool_calls": [],253            "error": str(e)254        }255 256 257# Verify imports work correctly258if __name__ == "__main__":259    print("[OK] OpenAI Agents SDK imported successfully")260 261    # Check which API is being used262    if os.getenv("OPENROUTER_API_KEY"):263        print("[OK] Using OpenRouter API")264        print(f"[OK] Base URL: {os.getenv('OPENROUTER_BASE_URL', 'https://openrouter.ai/api/v1')}")265    else:266        print("[OK] Using OpenAI API")267 268    print(f"[OK] Agent '{todobot_agent.name}' created with {len(todobot_agent.tools)} tools")269    print(f"[OK] Agent model: {todobot_agent.model}")270    print(f"[OK] Agent instructions: {len(get_agent_instructions())} characters")271    print("\n[OK] Function tools registered:")272    for tool in todobot_agent.tools:273        print(f"     - {tool.name}")274    print("\n[OK] Agent Runner implemented")275    print("\nTodoBot agent is ready for chat API integration.")276