add agent.py

src/agent.py

@@ -0,0 +1,121 @@
+"""Deep Agent with internet search and virtual filesystem."""
+
+import os
+from pathlib import Path
+from typing import Dict, Any, Optional
+
+from langchain_ollama import ChatOllama
+from deepagents import create_deep_agent
+
+from .search_tool import internet_search
+
+
+def create_my_agent():
+    """
+    Create a Deep Agent with internet search capability.
+    
+    The agent automatically has:
+    - File system tools (ls, read_file, write_file, edit_file, grep, glob)
+    - TODO planning tools (write_todos, read_todos)
+    - Internet search tool (custom)
+    
+    Returns:
+        Compiled agent ready for invocation
+    """
+    
+    # Initialize the model
+    model = ChatOllama(
+        model="llama3.2",
+        temperature=0.3,
+        # Optional: adjust based on your setup
+        base_url=os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
+    )
+    
+    # System prompt for the agent
+    system_prompt = """You are a research agent with the following capabilities:
+
+1. **Internet Search**: Use the 'internet_search' tool to find information online
+2. **File System**: You have a virtual filesystem for storing information
+   - Use 'write_file' to save important findings
+   - Use 'read_file' to retrieve saved information
+   - Use 'ls' to list files
+3. **Task Planning**: Use 'write_todos' to break down complex tasks into steps
+
+Guidelines:
+- Always break down complex tasks using write_todos first
+- Save important information to files using write_file
+- Keep files organized with meaningful names
+- After completing research, save a summary to 'summary.md' or similar
+
+Remember: All files you create will be automatically exported to the real filesystem when the task completes."""
+    
+    # Create the agent - filesystem and planning are included by default!
+    agent = create_deep_agent(
+        model=model,
+        tools=[internet_search],
+        system_prompt=system_prompt,
+        # Optional: enable debugging
+        debug=False,
+    )
+    
+    return agent
+
+
+def extract_virtual_files(agent, thread_id: str = "default") -> Dict[str, str]:
+    """
+    Extract virtual filesystem contents from agent state.
+    
+    Deep Agents stores files in the agent state. This function retrieves
+    all files created during the agent's execution.
+    
+    Args:
+        agent: The compiled deep agent
+        thread_id: Thread ID for the conversation (used for checkpointing)
+        
+    Returns:
+        Dictionary mapping file paths to content
+    """
+    try:
+        # Get the state from the agent's checkpoint
+        config = {"configurable": {"thread_id": thread_id}}
+        state = agent.get_state(config)
+        
+        # Access files stored in state
+        # The files are stored in state.values.get('files', {})
+        if state and hasattr(state, 'values'):
+            files = state.values.get('files', {})
+            return files
+    except Exception as e:
+        print(f"Warning: Could not extract files from agent state: {e}")
+    
+    return {}
+
+
+def export_files_to_disk(files: Dict[str, str], export_dir: str = "./exported_files"):
+    """
+    Export virtual files to the real filesystem.
+    
+    Args:
+        files: Dictionary of file paths to content
+        export_dir: Directory to export files to
+    """
+    root = Path(export_dir)
+    root.mkdir(parents=True, exist_ok=True)
+    
+    if not files:
+        print("No virtual files found to export.")
+        return
+    
+    for file_path, content in files.items():
+        # Remove leading slash if present
+        clean_path = file_path.lstrip('/')
+        full_path = root / clean_path
+        
+        # Create parent directories if needed
+        full_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        # Write the file
+        full_path.write_text(content, encoding="utf-8")
+        print(f"  ✓ Exported: {full_path}")
+    
+    print(f"\n✅ Exported {len(files)} file(s) to '{export_dir}/'")