Spaces:
Sleeping
Sleeping
File size: 2,860 Bytes
b840b29 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | 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 |