Create agents/legal_agent.py

agents/legal_agent.py

@@ -0,0 +1,65 @@
+# agents/legal_agent.py
+from typing import Dict, List, Optional
+from utils.llm_gateway import LLMGateway
+from utils.vector_store import VectorStore
+
+class LegalAgent:
+    def __init__(self, llm: LLMGateway, vector_store: VectorStore):
+        self.llm = llm
+        self.vector_store = vector_store
+
+    def process_query(self, query: str) -> str:
+        """Process a legal research query"""
+        # Search for relevant documents
+        relevant_docs = self.vector_store.search(query, n_results=3)
+        
+        # Construct prompt
+        prompt = self._construct_prompt(query, relevant_docs)
+        
+        # Generate response
+        system_prompt = """You are a legal research assistant. 
+        Provide clear, well-structured responses with citations to relevant documents.
+        Focus on factual analysis and avoid legal advice."""
+        
+        response = self.llm.generate(
+            prompt=prompt,
+            system_prompt=system_prompt,
+            temperature=0.7
+        )
+        
+        return response
+
+    def _construct_prompt(self, query: str, relevant_docs: List[Dict]) -> str:
+        """Construct prompt with context"""
+        prompt = f"""Query: {query}\n\nRelevant documents:\n"""
+        
+        for i, doc in enumerate(relevant_docs, 1):
+            prompt += f"\nDocument {i}:\n{doc['text']}\n"
+            
+        prompt += """\nPlease provide:
+1. Summary of relevant facts
+2. Analysis of key legal points
+3. Citations to relevant documents
+4. Suggested next steps for research"""
+        
+        return prompt
+
+    def generate_analysis(self, template: str, docs: List[Dict], case_type: str) -> str:
+        """Generate analysis based on template"""
+        prompt = f"""Template: {template}
+Case Type: {case_type}
+
+Relevant documents:
+"""
+        
+        for i, doc in enumerate(docs, 1):
+            prompt += f"\nDocument {i}:\n{doc['text']}\n"
+            
+        system_prompt = """Generate a structured analysis following the template.
+Include specific citations and maintain professional legal terminology."""
+        
+        return self.llm.generate(
+            prompt=prompt,
+            system_prompt=system_prompt,
+            temperature=0.7
+        )