Update agent.py

agent.py

@@ -63,7 +63,7 @@ class GaiaAgent:
         
         **Thought:** First, carefully consider the task. What information do you need to answer the question? Do you need to use a tool?
         **Tool:** If you need to search, use the search_tavily tool. The format is: <TOOL_CODE>search_tavily("your search query")</TOOL_CODE>
-        **Observation:** After a tool call, you will receive an observation (the tool's output).
+        **Observation:** After a tool call, you will receive an observation (the tool's output). This is factual information.
         **Answer:** Once you have gathered all necessary information, provide your final, concise answer directly.
 
         Your available tools:
@@ -82,11 +82,11 @@ class GaiaAgent:
         """
 
         max_iterations = 3
-        current_response = ""
+        current_response_history = "" # Ny variabel för att bygga upp historiken
 
         for i in range(max_iterations):
             # Lägg till "Thought:" här för att uppmuntra modellen att starta sin tankeprocess
-            full_prompt = prompt + current_response + "\n\nThought:" 
+            full_prompt = prompt + current_response_history + "\n\nThought:"
             
             print(f"[{i+1}/{max_iterations}] Generating response with prompt length: {len(full_prompt)}")
             
@@ -97,9 +97,10 @@ class GaiaAgent:
                 pad_token_id=self.tokenizer.eos_token_id,
                 do_sample=True,
                 top_k=50, top_p=0.95,
-                temperature=0.7 # Justera temperaturen till 0.7
+                temperature=0.7 
             )[0]['generated_text']
 
+            # Extrahera endast den nya delen av texten (modellens respons efter den sista "Thought:")
             new_content = generated_text[len(full_prompt):].strip()
             print(f"DEBUG - Full generated_text: \n---START---\n{generated_text}\n---END---")
             print(f"DEBUG - Extracted new_content: '{new_content}'")
@@ -108,12 +109,26 @@ class GaiaAgent:
             if "Answer:" in new_content:
                 final_answer = new_content.split("Answer:", 1)[1].strip()
                 print(f"Final answer from model:\n{final_answer}")
-                return final_answer
+                return final_answer # Returnera det slutgiltiga svaret
+            
             elif "<TOOL_CODE>" in new_content and "</TOOL_CODE>" in new_content:
-                start_index = new_content.find("<TOOL_CODE>") + len("<TOOL_CODE>")
-                end_index = new_content.find("</TOOL_CODE>")
-                tool_call_str = new_content[start_index:end_index].strip()
+                # Modellen genererade ett verktygskall.
+                # Vi vill inte inkludera modellens egna "Observation:" eller "Tool:"-text i historiken
+                # innan verktyget faktiskt körts. Vi tar bara själva tool_code strängen.
+                
+                tool_call_start = new_content.find("<TOOL_CODE>")
+                tool_call_end = new_content.find("</TOOL_CODE>") + len("</TOOL_CODE>")
                 
+                # Försök att extrahera tanken som ledde till verktygskallet
+                thought_part = ""
+                if "Thought:" in new_content[:tool_call_start]:
+                    thought_part = new_content.split("Thought:", 1)[1].split("Tool:", 1)[0].strip()
+                elif tool_call_start > 0: # Om det finns text före tool code
+                    thought_part = new_content[:tool_call_start].strip()
+
+                tool_code_section = new_content[tool_call_start:tool_call_end]
+                tool_call_str = tool_code_section.replace("<TOOL_CODE>", "").replace("</TOOL_CODE>", "").strip()
+
                 print(f"Tool call detected: {tool_call_str}")
                 
                 try:
@@ -121,18 +136,24 @@ class GaiaAgent:
                         query = tool_call_str[len("search_tavily("):-1].strip().strip('"').strip("'")
                         tool_output = search_tavily(query)
                         print(f"Tool result: {tool_output[:200]}...")
-                        current_response += f"\n\nObservation: {tool_output}\n"
+                        
+                        # Lägg till tanken, verktygskallet och det FAKTISKA observationen till historiken
+                        current_response_history += f"\n\nThought: {thought_part}\nTool: {tool_code_section}\nObservation: {tool_output}\n"
                     else:
                         tool_output = f"Unknown tool: {tool_call_str}"
                         print(f"Error: {tool_output}")
-                        current_response += f"\n\nObservation: {tool_output}\n"
+                        current_response_history += f"\n\nThought: {thought_part}\nTool: {tool_code_section}\nObservation: {tool_output}\n"
                 except Exception as tool_e:
                     tool_output = f"Error running tool {tool_call_str}: {tool_e}"
                     print(f"Error: {tool_output}")
-                    current_response += f"\n\nObservation: {tool_output}\n"
+                    current_response_history += f"\n\nThought: {thought_part}\nTool: {tool_code_section}\nObservation: {tool_output}\n"
             else:
-                # Om modellen varken ger svar eller verktygskall men genererar något annat
-                current_response += f"\n\n{new_content}\n"
+                # Modellen genererade varken ett verktygskall eller ett slutgiltigt svar.
+                # Lägg till det den faktiskt genererade till historiken så den kan fortsätta sin tanke.
+                current_response_history += f"\n\nThought: {new_content}\n"
                 print(f"Model generated non-tool/non-answer content. Appending: {new_content[:100]}...")
 
-        return "Agent could not complete the task within the allowed iterations. Latest response: " + new_content.strip() if new_content else "Agent could not complete the task within the allowed iterations. No meaningful content generated."
+        # Om max_iterations nås utan slutgiltigt svar
+        return "Agent could not complete the task within the allowed iterations. Latest relevant content: " + \
+               (current_response_history[-500:] if current_response_history else "No meaningful content generated.")
+