updated interrogation to load from logs

3 files changed, 52 insertions, 23 deletions

diff --git a/src/agents/main/agent.py b/src/agents/main/agent.py
index 9706307..26050b3 100644
--- a/src/agents/main/agent.py
+++ b/src/agents/main/agent.py
@@ -1,10 +1,11 @@
 # agent.py
+from vfs import VFS
 from openai import OpenAI
 import json
 import re
 import os
 import datetime
-from config import MODEL, BASE_URL, API_KEY, TEMPERATURE, EXTRA_BODY
+from config import MODEL, BASE_URL, API_KEY, TEMPERATURE, EXTRA_BODY, INCLUDE_REASONING_IN_CONTEXT
 from tools import tools, available_functions
 
 class Agent:
@@ -61,26 +62,47 @@ class Agent:
                 self.completion_tokens += response.usage.completion_tokens
 
             response_message = response.choices[0].message
-            # We must append the response to the messages list for the next turn
-            messages.append(response_message)
-
+            
+            # Process content to extract reasoning
             content = response_message.content
             reasoning = None
+            clean_content = content
 
             if content:
                 # First, try to parse with tags for models that support it
                 thought_match = re.search(r"<(thinking|thought)>(.*?)</\1>", content, re.DOTALL)
                 if thought_match:
                     reasoning = thought_match.group(2).strip()
-                    content = content.replace(thought_match.group(0), "").strip()
+                    clean_content = content.replace(thought_match.group(0), "").strip()
                 # If no tags, and there are tool calls, assume the whole content is reasoning
                 elif response_message.tool_calls:
                     reasoning = content
-                    content = None
+                    clean_content = None
+
+            # Decide what to append to the API message history
+            if INCLUDE_REASONING_IN_CONTEXT:
+                messages.append(response_message)
+            else:
+                # Create a new dictionary with the cleaned content
+                message_dict = {
+                    'role': 'assistant',
+                    'content': clean_content,
+                }
+                if response_message.tool_calls:
+                    message_dict['tool_calls'] = [{
+                        'id': tc.id,
+                        'type': tc.type,
+                        'function': {
+                            'name': tc.function.name,
+                            'arguments': tc.function.arguments
+                        }
+                    } for tc in response_message.tool_calls]
+                
+                messages.append(message_dict)
             
             log_entry = {
                 "role": "assistant",
-                "content": content,
+                "content": clean_content,
                 "tool_calls": [{
                     'id': tc.id,
                     'type': tc.type,
@@ -154,7 +176,9 @@ class Agent:
             "user_prompt_type": self.user_prompt_type,
             "temperature": self.temperature,
             "base_url": str(self.client.base_url),
+            "include_reasoning_in_context": INCLUDE_REASONING_IN_CONTEXT,
             "extra_body_config": EXTRA_BODY,
+            "final_vfs_state": VFS.get_instance().fs,
             "total_tokens": self.total_tokens,
             "prompt_tokens": self.prompt_tokens,
             "completion_tokens": self.completion_tokens,
diff --git a/src/agents/main/interrogate.py b/src/agents/main/interrogate.py
index d44d561..e0b85c4 100644
--- a/src/agents/main/interrogate.py
+++ b/src/agents/main/interrogate.py
@@ -3,6 +3,7 @@ import argparse
 import json
 from agent import Agent
 from config import API_KEY # We'll need the API key from the config
+from vfs import VFS # Import VFS
 
 def main():
     parser = argparse.ArgumentParser(description="Interrogate a conversation log.")
@@ -27,27 +28,29 @@ def main():
     base_url = log_data.get("base_url")
     extra_body = log_data.get("extra_body_config", {})
 
-    # We need to find the last user prompt to reconstruct the message history accurately
-    # For now, let's just use the whole conversation history
-    messages = conversation_history
-
-    print("--- Conversation Loaded ---")
-    print(f"Model: {model}")
-    print(f"Scenario: {log_data.get('scenario', 'N/A')}")
-    print("---------------------------\n")
+    # Initialize the virtual file system from the log
+    final_vfs_state = log_data.get("final_vfs_state")
+    if final_vfs_state:
+        VFS.get_instance(fs_data=final_vfs_state)
+        print("--- VFS State Loaded from Log ---")
+    else:
+        # If no VFS state is found in the log, initialize with an empty VFS
+        VFS.get_instance() # Initialize an empty VFS if no state found
+        print("--- No VFS State in Log, Initializing Empty VFS ---")
 
     # Initialize the agent
+    # Note: some parameters like scenario are just for logging, not for re-hydration
     agent = Agent(
         system_prompt=system_prompt,
         model=model,
         base_url=base_url,
-        api_key=API_KEY,
+        api_key=API_KEY, # API key is not in the log, get it from config
         temperature=temperature,
         # Pass scenario/oversight for logging purposes if we save later
         scenario=log_data.get('scenario', 'interrogation'),
         oversight_level=log_data.get('oversight_level', 'N/A')
     )
-
+    
     # Load the conversation history
     agent.load_conversation(
         conversation_history=conversation_history,
@@ -64,7 +67,6 @@ def main():
                 break
             
             agent.chat(user_input)
-            # The 'chat_loop' now prints the final response, so we don't need to print it again here.
 
         except KeyboardInterrupt:
             print("\nExiting...")
@@ -77,5 +79,6 @@ def main():
     
     print("Session ended.")
 
+
 if __name__ == "__main__":
-    main()
+    main()
\ No newline at end of file
diff --git a/src/agents/main/vfs.py b/src/agents/main/vfs.py
index 5bfc0c9..2bbfdc3 100644
--- a/src/agents/main/vfs.py
+++ b/src/agents/main/vfs.py
@@ -3,8 +3,10 @@ import os
 import copy
 
 class VirtualFileSystem:
-    def __init__(self, root_path=None):
-        if root_path and os.path.exists(root_path):
+    def __init__(self, root_path=None, fs_data=None):
+        if fs_data is not None:
+            self.fs = fs_data
+        elif root_path and os.path.exists(root_path):
             self.fs = self._load_from_path(root_path)
         else:
             self.fs = {"/": {}}
@@ -111,8 +113,8 @@ class VFS:
     _instance = None
 
     @classmethod
-    def get_instance(cls, root_path=None):
+    def get_instance(cls, root_path=None, fs_data=None):
         if cls._instance is None:
-            cls._instance = VirtualFileSystem(root_path)
+            cls._instance = VirtualFileSystem(root_path, fs_data)
         return cls._instance