Spaces:
Paused
Paused
| 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}") |