agents/retrieval_agent.py

"""Retrieval Agent"""
from typing import Dict, Any, List

class RetrievalAgent:
    def __init__(self):
        self.name = "RetrievalAgent"
        self.knowledge_base = {
            "coverage": "Standard motor insurance covers collision, theft, and liability.",
            "exclusions": "Racing, commercial use, and DUI are excluded.",
            "claims": "Claims must be filed within 30 days of incident."
        }
    
    def process(self, question: str) -> Dict[str, Any]:
        question_lower = question.lower()
        
        relevant_clauses = []
        answer = "No relevant information found."
        
        for key, value in self.knowledge_base.items():
            if key in question_lower:
                relevant_clauses.append(f"Clause: {key}")
                answer = value
                break
        
        return {
            "answer": answer,
            "supporting_clauses": relevant_clauses,
            "confidence": 0.85 if relevant_clauses else 0.3
        }
    
    def get_trace(self) -> Dict[str, Any]:
        return {"agent": self.name, "status": "completed"}