import os import json import requests from openai import OpenAI from agent_tools import TOOLS_MAPPING class CustomAgent: def __init__(self): print("[CustomAgent] Starting up...") token = os.getenv("OR_TOKEN") if not token: raise RuntimeError("Missing OpenRouter token. Set OR_TOKEN environment variable.") self.client = OpenAI( base_url="https://openrouter.ai/api/v1", api_key=token ) def run(self, prompt: str) -> str: print(f"[Agent] Received prompt: {prompt}") instructions = ( "You're solving a quiz. Only return the final result. " "Use numbers or minimal text as requested. Avoid extra formatting or explanations." ) # Optional reasoning helper context_info = self._retrieve_helper_response(prompt) if context_info: instructions += f"\n\nBackground reference: {context_info}" dialogue = [ {"role": "system", "content": instructions}, {"role": "user", "content": [{"type": "text", "text": prompt}]} ] for attempt in range(3): try: print(f"[Agent] Attempt {attempt + 1}: requesting response...") reply = self.client.chat.completions.create( model="microsoft/mai-ds-r1:free", messages=dialogue, temperature=0.0, max_tokens=2048, extra_headers={ "HTTP-Referer": "https://your-app.example.com", "X-Title": "quiz-agent" } ).choices[0].message dialogue.append(reply) if reply.tool_calls is None: return reply.content.strip() if reply.content else "No response produced." for tool_call in reply.tool_calls: tool_name = tool_call.function.name tool_input = json.loads(tool_call.function.arguments) if tool_name in TOOLS_MAPPING: result = TOOLS_MAPPING[tool_name](**tool_input) dialogue.append({ "role": "tool", "tool_call_id": tool_call.id, "name": tool_name, "content": json.dumps(result) }) except Exception as err: print(f"[Error] Attempt {attempt + 1} failed: {err}") return "Agent could not produce a valid response after multiple attempts." def _retrieve_helper_response(self, query: str) -> str: """Optionally enhance context by querying a helper LLM.""" helper_url = os.getenv("REASONING_API", "http://localhost:8004/ask") try: payload = { "question": f"Question: {query}\n\n1. Break into parts.\n2. Answer.\n3. Give final result." } headers = {"Content-Type": "application/json"} response = requests.post(helper_url, json=payload, headers=headers, timeout=30) response.raise_for_status() return response.json().get("answer", "") except Exception as ex: print(f"[Helper] Failed to retrieve context: {ex}") return ""