agent.py

import os
output_dir = "/mnt/agents/output"

# 4. agent.py - Updated for local llama.cpp with model download
agent_content = r'''import json
import time
import traceback
import os
from typing import Generator, Tuple, List, Dict, Any
from huggingface_hub import hf_hub_download

from config import (
    MODEL_REPO, MODEL_FILE, MODEL_PATH, 
    MAX_TOKENS, TEMPERATURE, N_CTX, N_BATCH, VERBOSE,
    JAILBREAK_SYSTEM_PROMPT, DEEPTHINK_PROMPT
)
from tools import TOOL_MAP, TOOL_DEFINITIONS
from utils import retry_with_backoff

class ModelLoader:
    """Singleton model loader to avoid reloading on every request."""
    _instance = None
    _llm = None
    _loaded = False
    
    @classmethod
    def get_llm(cls):
        if cls._llm is None:
            cls._load_model()
        return cls._llm
    
    @classmethod
    def _load_model(cls):
        """Download and load the GGUF model."""
        from llama_cpp import Llama
        
        # Download model if not exists
        if not os.path.exists(MODEL_PATH):
            os.makedirs("models", exist_ok=True)
            print(f"Downloading model from {MODEL_REPO}...")
            hf_hub_download(
                repo_id=MODEL_REPO,
                filename=MODEL_FILE,
                local_dir="models",
                local_dir_use_symlinks=False
            )
        
        print(f"Loading model from {MODEL_PATH}...")
        cls._llm = Llama(
            model_path=MODEL_PATH,
            n_ctx=N_CTX,
            n_batch=N_BATCH,
            verbose=VERBOSE
        )
        cls._loaded = True
        print("Model loaded successfully!")
    
    @classmethod
    def is_loaded(cls):
        return cls._loaded


class AutonomousAgent:
    def __init__(self):
        self.llm = None
        self.conversation_history: List[Dict[str, Any]] = []
        self.max_iterations = 10
    
    def _ensure_model(self):
        """Ensure model is loaded."""
        if self.llm is None:
            self.llm = ModelLoader.get_llm()
    
    def reset(self):
        """Reset conversation history."""
        self.conversation_history = []
    
    def _build_messages(self, user_input: str, system_prompt: str, deepthink: bool = False, force_search: bool = False) -> List[Dict[str, Any]]:
        """Build the message list for the API call."""
        messages = []
        
        # System prompt
        sys = system_prompt or JAILBREAK_SYSTEM_PROMPT
        if deepthink:
            sys += "\n\n" + DEEPTHINK_PROMPT
        messages.append({"role": "system", "content": sys})
        
        # Add conversation history (last 10 exchanges)
        for msg in self.conversation_history[-20:]:
            messages.append(msg)
        
        # User input
        if force_search:
            user_input = "Search the web for: " + user_input
        messages.append({"role": "user", "content": user_input})
        
        return messages
    
    def _call_llm(self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Call the local LLM with optional tools."""
        self._ensure_model()
        
        def _do_call():
            kwargs = {
                "messages": messages,
                "max_tokens": MAX_TOKENS,
                "temperature": TEMPERATURE,
            }
            if tools:
                kwargs["tools"] = tools
                kwargs["tool_choice"] = "auto"
            return self.llm.create_chat_completion(**kwargs)
        
        try:
            return retry_with_backoff(
                _do_call,
                max_retries=2,
                initial_delay=1.0,
                backoff_factor=2.0,
                exceptions=(Exception,)
            )
        except Exception as e:
            raise RuntimeError(f"LLM inference error: {str(e)}")
    
    def _execute_tool(self, tool_name: str, tool_args: Dict[str, Any]) -> str:
        """Execute a tool and return the result."""
        if tool_name not in TOOL_MAP:
            return f"Error: Tool '{tool_name}' not found."
        
        try:
            tool_func = TOOL_MAP[tool_name]
            result = tool_func(**tool_args)
            # Truncate very long results
            if isinstance(result, str) and len(result) > 12000:
                result = result[:12000] + "\n...[truncated]"
            return str(result)
        except Exception as e:
            return f"Tool execution error: {str(e)}"
    
    def run_stream(self, user_input: str, system_prompt: str = None, deepthink: bool = False, force_search: bool = False) -> Generator[Tuple[str, bool], None, None]:
        """
        Run the agent in streaming mode.
        Yields (text, is_final) tuples.
        """
        # Yield loading message if model not loaded yet
        if not ModelLoader.is_loaded():
            yield ("\n⏳ **Loading model...** This may take 30-60 seconds on first run.\n", False)
        
        messages = self._build_messages(user_input, system_prompt, deepthink, force_search)
        iteration = 0
        final_response = ""
        
        while iteration < self.max_iterations:
            iteration += 1
            
            try:
                # Call LLM with tools
                response = self._call_llm(messages, TOOL_DEFINITIONS)
                message = response["choices"][0]["message"]
                
                # Check if the model wants to use tools
                if "tool_calls" in message and message["tool_calls"]:
                    tool_calls = message["tool_calls"]
                    tool_names = [tc["function"]["name"] for tc in tool_calls]
                    yield (f"\n🔧 **Using tools:** {', '.join(tool_names)}\n", False)
                    
                    # Add assistant message with tool calls to history
                    assistant_msg = {
                        "role": "assistant",
                        "content": message.get("content", "") or "",
                        "tool_calls": [
                            {
                                "id": tc["id"],
                                "type": "function",
                                "function": {
                                    "name": tc["function"]["name"],
                                    "arguments": tc["function"]["arguments"]
                                }
                            } for tc in tool_calls
                        ]
                    }
                    messages.append(assistant_msg)
                    self.conversation_history.append({"role": "assistant", "content": message.get("content", "") or f"Using tools: {', '.join(tool_names)}"})
                    
                    # Execute each tool call
                    for tool_call in tool_calls:
                        tool_name = tool_call["function"]["name"]
                        try:
                            tool_args = json.loads(tool_call["function"]["arguments"])
                        except:
                            tool_args = {}
                        
                        yield (f"\n⏳ **Executing:** `{tool_name}`...\n", False)
                        
                        # Execute the tool
                        tool_result = self._execute_tool(tool_name, tool_args)
                        
                        # Yield tool result preview
                        preview = tool_result[:200] + "..." if len(tool_result) > 200 else tool_result
                        yield (f"\n✅ **Result:** `{preview}`\n", False)
                        
                        # Add tool result to messages
                        messages.append({
                            "role": "tool",
                            "tool_call_id": tool_call["id"],
                            "content": tool_result
                        })
                
                else:
                    # No tool calls - this is the final answer
                    final_response = message.get("content", "") or ""
                    
                    # Add to conversation history
                    self.conversation_history.append({"role": "user", "content": user_input})
                    self.conversation_history.append({"role": "assistant", "content": final_response})
                    
                    # Trim history if too long
                    if len(self.conversation_history) > 40:
                        self.conversation_history = self.conversation_history[-40:]
                    
                    yield (final_response, True)
                    return
                    
            except Exception as e:
                error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
                yield (error_msg, True)
                return
        
        # Max iterations reached
        yield ("\n⚠️ **Max iterations reached.** The agent was unable to complete the task within the allowed number of steps. Please try a more specific request.", True)


class AgentManager:
    """Manager class for handling multiple agent sessions."""
    
    def __init__(self):
        self.agents: Dict[str, AutonomousAgent] = {}
    
    def get_agent(self, session_id: str) -> AutonomousAgent:
        """Get or create an agent for a session."""
        if session_id not in self.agents:
            self.agents[session_id] = AutonomousAgent()
        return self.agents[session_id]
    
    def reset_agent(self, session_id: str):
        """Reset an agent's conversation."""
        if session_id in self.agents:
            self.agents[session_id].reset()
    
    def delete_agent(self, session_id: str):
        """Delete an agent session."""
        if session_id in self.agents:
            del self.agents[session_id]
'''

with open(f"{output_dir}/agent.py", "w") as f:
    f.write(agent_content)

print("agent.py written")

# Verify
import py_compile
try:
    py_compile.compile(f"{output_dir}/agent.py", doraise=True)
    print("agent.py compiles successfully!")
except Exception as e:
    print(f"Compile error: {e}")