replicate.py

import os
import sys
from typing import Dict, Any
import cog
from dotenv import load_dotenv

# Add the current directory to the path so we can import from src
sys.path.append(".")

# Load environment variables
load_dotenv()

# Import after path setup
from src.agents.agent_director import AgentDirector
from src.utils.config import CHAT_MODEL

class Predictor(cog.Predictor):
    def setup(self):
        """Setup the model and load resources."""
        # Check environment variables
        if not os.environ.get("OPENAI_API_KEY"):
            raise ValueError("OPENAI_API_KEY environment variable is required")
        
        # Check data files exist
        required_files = [
            "embeddings/faiss_index.index",
            "data/doc_chunks.pkl",
            "embeddings/embeddings.pkl"
        ]
        
        for file_path in required_files:
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"Required file not found: {file_path}")
        
        # Initialize the agent director
        self.director = AgentDirector(model=CHAT_MODEL, top_k=50)
        print("Agentic Defensor initialized successfully")

    @cog.input("query", type=str, help="The legal query to process")
    @cog.input("top_k", type=int, default=50, help="Number of chunks to retrieve")
    @cog.input("model", type=str, default=CHAT_MODEL, help="OpenAI model to use")
    @cog.input("debug", type=bool, default=False, help="Show agent reasoning steps")
    def predict(self, query: str, top_k: int = 50, model: str = CHAT_MODEL, debug: bool = False) -> Dict[str, Any]:
        """
        Process a query using the Agentic Defensor system.
        
        Args:
            query: The legal query to process
            top_k: Number of chunks to retrieve
            model: OpenAI model to use
            debug: Whether to show agent reasoning steps
            
        Returns:
            Dictionary containing the answer and metadata
        """
        # Re-initialize with custom parameters if needed
        if model != self.director.model or top_k != self.director.top_k or debug != self.director.debug:
            self.director = AgentDirector(model=model, top_k=top_k, debug=debug)
        
        # Process the query
        result = self.director.process_query(query)
        
        # Format the response
        response = {
            "query": query,
            "answer": result["answer"],
            "model_used": result["model_used"]
        }
        
        # Add debug information if available
        if debug and "reasoning_steps" in result:
            response["reasoning_steps"] = result["reasoning_steps"]
        
        # Add metadata
        if "num_chunks_retrieved" in result:
            response["num_chunks_retrieved"] = result["num_chunks_retrieved"]
        
        return response