Update app_core.py

app_core.py

@@ -1,90 +1,91 @@
 import os
-import requests
 import json
+import requests
 from openai import OpenAI
 from agent_tools import TOOLS_MAPPING
 
 
 class CustomAgent:
     def __init__(self):
-        print("[Init] CustomAgent initialized.")
+        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=os.getenv("OR_TOKEN")
+            api_key=token
         )
 
-    def run(self, question: str) -> str:
-        print(f"[Run] Question received: {question}")
+    def run(self, prompt: str) -> str:
+        print(f"[Agent] Received prompt: {prompt}")
 
-        # Prepare system message
-        system_message = (
-            "You are taking a quiz. Follow instructions precisely. "
-            "Only return the FINAL ANSWER using digits or plain words, as requested. "
-            "Avoid explanations, symbols, or extra text."
+        instructions = (
+            "You're solving a quiz. Only return the final result. "
+            "Use numbers or minimal text as requested. Avoid extra formatting or explanations."
         )
 
-        # Call internal LLM to expand reasoning
-        try:
-            llm_response = requests.post(
-                "http://localhost:8004/ask",
-                headers={"Content-Type": "application/json"},
-                json={"question": f"Question: {question}\n\n1. Break into sub-questions.\n2. Answer them.\n3. Return final answer."},
-                timeout=60
-            )
-            llm_response.raise_for_status()
-            pre_answer = llm_response.json().get("answer", "")
-        except Exception as e:
-            print(f"[Error] Failed to get reasoning: {e}")
-            pre_answer = ""
-
-        system_message += f"\n\nLLM helper result: {pre_answer}"
-
-        messages = [
-            {"role": "system", "content": system_message},
-            {"role": "user", "content": [{"type": "text", "text": question}]}
+        # 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}]}
         ]
 
-        # Retry loop
-        for _ in range(3):
+        for attempt in range(3):
             try:
-                print("[Run] Querying OpenRouter model...")
-                completion = self.client.chat.completions.create(
+                print(f"[Agent] Attempt {attempt + 1}: requesting response...")
+                reply = self.client.chat.completions.create(
                     model="microsoft/mai-ds-r1:free",
-                    messages=messages,
+                    messages=dialogue,
                     temperature=0.0,
                     max_tokens=2048,
                     extra_headers={
-                        "HTTP-Referer": "<YOUR_SITE_URL>",
-                        "X-Title": "<YOUR_SITE_NAME>",
-                    }
-                )
-
-                response_msg = completion.choices[0].message
-                messages.append(response_msg)
-
-                # If final content is available, return it
-                if response_msg.tool_calls is None:
-                    answer = response_msg.content or "Unable to generate response."
-                    print(f"[Answer] {answer}")
-                    return answer
-
-                # Handle any tool calls
-                for tool_call in response_msg.tool_calls:
-                    func_name = tool_call.function.name
-                    func_args = json.loads(tool_call.function.arguments)
-                    tool_response = TOOLS_MAPPING[func_name](**func_args)
-
-                    tool_msg = {
-                        "role": "tool",
-                        "tool_call_id": tool_call.id,
-                        "name": func_name,
-                        "content": json.dumps(tool_response),
+                        "HTTP-Referer": "https://your-app.example.com",
+                        "X-Title": "quiz-agent"
                     }
+                ).choices[0].message
 
-                    messages.append(tool_msg)
+                dialogue.append(reply)
 
-            except Exception as e:
-                print(f"[Error] Generation failed: {e}")
-                return "I encountered an error while trying to answer the question."
+                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 "Failed to generate a valid response after retries."
+            return response.json().get("answer", "")
+        except Exception as ex:
+            print(f"[Helper] Failed to retrieve context: {ex}")
+            return ""