Update functions/agent_helper_functions.py

functions/agent_helper_functions.py

@@ -1,124 +1,80 @@
-"""Helper functions for the agent(s) in the GAIA question answering system."""
+# functions/agent_helper_functions.py
+
+"""Helper functions for the GAIA agent."""
 
-import os
 import time
 import json
 import logging
-from smolagents import CodeAgent, ActionStep, MessageRole, InferenceClientModel
-from configuration import CHECK_MODEL, TOKEN_LIMITER, STEP_WAIT
+from smolagents import ActionStep, MessageRole
+from configuration import MODEL, CHECK_MODEL, TOKEN_LIMITER, STEP_WAIT
 
-# Logger per questo modulo
 logger = logging.getLogger(__name__)
 
-# Model dedicato alla summarizzazione dei messaggi
-SUMMARIZER_MODEL = InferenceClientModel(
-    "deepseek-ai/DeepSeek-V3",
-    provider="together",
-    api_key=os.getenv("TOGETHER_API_KEY"),
-    temperature=0,
-    max_tokens=8000
-)
-
+# Usa lo stesso MODEL anche per i riassunti
+SUMMARIZER_MODEL = MODEL
 
 def check_reasoning(final_answer: str, agent_memory) -> bool:
-    """Checks the reasoning and plot of the agent's final answer."""
     prompt = (
         f"Here is a user-given task and the agent steps: "
         f"{agent_memory.get_succinct_steps()}. "
         "Please check that the reasoning process and answer are correct. "
         "First list reasons why yes/no, then write your final decision: "
-        "PASS in caps lock if it is satisfactory, FAIL if it is not. "
+        "PASS in caps lock if satisfactory, FAIL if not. "
         f"Final answer: {final_answer}"
     )
-
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {
-                    "type": "text",
-                    "text": prompt,
-                }
-            ],
-        }
-    ]
-
+    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
     feedback = CHECK_MODEL(messages).content
-    print("Feedback: ", feedback)
-
+    print("Feedback:", feedback)
     if "FAIL" in feedback:
         raise Exception(feedback)
     return True
 
-
-def step_memory_cap(memory_step: ActionStep, agent: CodeAgent) -> None:
-    """Removes old steps from agent memory to keep context length under control."""
-    task_step = agent.memory.steps[0]
-    planning_step = agent.memory.steps[1]
-    latest_step = agent.memory.steps[-1]
-
-    # Keep only the first two + latest
+def step_memory_cap(memory_step: ActionStep, agent) -> None:
     if len(agent.memory.steps) > 2:
-        agent.memory.steps = [task_step, planning_step, latest_step]
-
-    logger.info('Agent memory has %d steps', len(agent.memory.steps))
-    logger.info('Latest step is %d', memory_step.step_number)
-    logger.info('Messages in latest step: %d',
-                len(latest_step.model_input_messages))
-    logger.info('Token usage: %d', latest_step.token_usage.total_tokens)
-
-    # If troppo token usage, faccio il summary
-    if latest_step.token_usage.total_tokens > TOKEN_LIMITER:
-        logger.info('Token usage %d > %d, summarizing old messages',
-                    latest_step.token_usage.total_tokens, TOKEN_LIMITER)
-        summary = summarize_old_messages(latest_step.model_input_messages[1:])
+        first, second, last = (
+            agent.memory.steps[0],
+            agent.memory.steps[1],
+            agent.memory.steps[-1]
+        )
+        agent.memory.steps = [first, second, last]
+
+    latest = agent.memory.steps[-1]
+    if latest.token_usage.total_tokens > TOKEN_LIMITER:
+        logger.info(
+            "Token usage %d > %d, summarizing",
+            latest.token_usage.total_tokens, TOKEN_LIMITER
+        )
+        summary = summarize_old_messages(latest.model_input_messages[1:])
         if summary:
-            # Ricostruisco il solo step con il summary
-            new_messages = [
-                latest_step.model_input_messages[0],
+            new_msgs = [
+                latest.model_input_messages[0],
                 {
-                    'role': MessageRole.USER,
-                    'content': [{
-                        'type': 'text',
-                        'text': f'Here is a summary of your investigation so far: {summary}'
+                    "role": MessageRole.USER,
+                    "content": [{
+                        "type": "text",
+                        "text": f"Here is a summary of your investigation so far: {summary}"
                     }]
                 }
             ]
-            agent.memory.steps = [task_step]
-            agent.memory.steps[0].model_input_messages = new_messages
-            logger.debug('Old messages summarized into new context.')
-
+            agent.memory.steps = [agent.memory.steps[0]]
+            agent.memory.steps[0].model_input_messages = new_msgs
 
 def summarize_old_messages(messages: list) -> str | None:
-    """Summarizes old messages to keep context length under control using DeepSeek."""
     if not messages:
         return None
-
     prompt = (
         "Summarize the following interaction between an AI agent and a user "
-        "in plain text (not JSON): "
-        + json.dumps(messages)
+        "in plain text (not JSON): " + json.dumps(messages)
     )
-    chat_input = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": prompt}
-            ]
-        }
-    ]
-
+    chat_input = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
     try:
-        response = SUMMARIZER_MODEL(chat_input)
-        return response.content
+        resp = SUMMARIZER_MODEL(chat_input)
+        return resp.content
     except Exception as e:
         logger.error("Error during summarization: %s", e)
         return None
 
-
-def step_wait(memory_step: ActionStep, agent: CodeAgent) -> bool:
-    """Waits to prevent hitting API rate limits."""
-    logger.info('Waiting %d seconds (step %d)', STEP_WAIT,
-                memory_step.step_number)
+def step_wait(memory_step: ActionStep, agent) -> bool:
+    logger.info("Waiting %d seconds (step %d)", STEP_WAIT, memory_step.step_number)
     time.sleep(STEP_WAIT)
-    return True
+    return True