nas-tools/dashboard/backend/services/agents/base_agent.py

"""
Base Agent Class - Foundation for all OPC agents
"""
from typing import Dict, Any, List, Optional
import json
import httpx
import os
from datetime import datetime


class BaseAgent:
    """Base class for all OPC agents"""

    def __init__(self, agent_id: str, name: str, role: str, system_prompt: str, config: Dict[str, Any]):
        self.agent_id = agent_id
        self.name = name
        self.role = role
        self.system_prompt = system_prompt
        self.config = config
        self.litellm_url = os.getenv("LITELLM_URL", "http://litellm:4005")
        self.litellm_api_key = os.getenv("LITELLM_API_KEY", "")

    async def execute(self, task: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Execute agent on a task
        Returns: {
            "actions_proposed": [...],
            "reasoning": "...",
            "requires_approval": bool
        }
        """
        # Build prompt with task context
        user_prompt = self._build_prompt(task, context)

        # Call LLM
        response = await self._call_llm(user_prompt)

        # Parse response and extract actions
        result = self._parse_response(response)

        return result

    def _build_prompt(self, task: Dict[str, Any], context: Dict[str, Any]) -> str:
        """Build the prompt for the LLM"""
        prompt = f"""You are {self.name}, a {self.role} agent.

Task Details:
- Title: {task.get('title')}
- Description: {task.get('description', 'No description')}
- Status: {task.get('status')}
- Priority: {task.get('priority')}
- Tags: {', '.join(task.get('tags', []))}

Context:
- Project: {context.get('project', 'None')}
- Related tasks: {len(context.get('related_tasks', []))}
- Company goals: {context.get('goals', 'None')}

Your role: {self.system_prompt}

Based on this task, propose concrete actions you would take. Format your response as JSON:
{{
  "reasoning": "Your analysis of the task",
  "actions": [
    {{"type": "create_subtask", "title": "...", "description": "..."}},
    {{"type": "update_task_status", "status": "..."}},
    {{"type": "send_notification", "message": "..."}},
    {{"type": "request_approval", "question": "..."}}
  ],
  "requires_approval": false
}}

Only propose actions you can actually execute. Be specific and actionable."""

        return prompt

    async def _call_llm(self, prompt: str) -> str:
        """Call LiteLLM proxy"""
        model = self.config.get("model", "claude-sonnet-4-6")
        temperature = self.config.get("temperature", 0.7)

        try:
            headers = {"Content-Type": "application/json"}
            # Only add auth header if API key is set and not empty
            if self.litellm_api_key and self.litellm_api_key.strip():
                headers["Authorization"] = f"Bearer {self.litellm_api_key}"

            async with httpx.AsyncClient(timeout=60.0) as client:
                response = await client.post(
                    f"{self.litellm_url}/chat/completions",
                    headers=headers,
                    json={
                        "model": model,
                        "messages": [
                            {"role": "system", "content": self.system_prompt},
                            {"role": "user", "content": prompt}
                        ],
                        "temperature": temperature,
                        "max_tokens": 2000
                    }
                )
                response.raise_for_status()
                data = response.json()
                return data["choices"][0]["message"]["content"]
        except Exception as e:
            raise Exception(f"LLM call failed: {str(e)}")

    def _parse_response(self, response: str) -> Dict[str, Any]:
        """Parse LLM response and extract actions"""
        try:
            # Try to extract JSON from response
            if "```json" in response:
                json_str = response.split("```json")[1].split("```")[0].strip()
            elif "```" in response:
                json_str = response.split("```")[1].split("```")[0].strip()
            else:
                json_str = response.strip()

            result = json.loads(json_str)

            # Validate structure
            if "actions" not in result:
                result["actions"] = []
            if "reasoning" not in result:
                result["reasoning"] = "No reasoning provided"
            if "requires_approval" not in result:
                # Check if any action requires approval
                result["requires_approval"] = any(
                    action.get("type") == "request_approval" for action in result["actions"]
                )

            return result

        except Exception as e:
            # Fallback: return error action
            return {
                "reasoning": f"Failed to parse response: {str(e)}",
                "actions": [
                    {
                        "type": "error",
                        "message": f"Agent response parsing failed: {str(e)}"
                    }
                ],
                "requires_approval": False
            }

    def get_capabilities(self) -> List[str]:
        """Get agent capabilities"""
        return self.config.get("capabilities", [])

    def requires_approval(self) -> bool:
        """Check if agent requires approval for actions"""
        return self.config.get("approval_level") == "cxo"