Update utils/legal_prompt_generator.py

utils/legal_prompt_generator.py

@@ -2,6 +2,7 @@ import json
 from typing import Dict, List, Optional
 from datetime import datetime
 
+
 class LegalPromptGenerator:
     def __init__(self, ontology_path: str = "data/legal_ontology.json"):
         """Initialize prompt generator with legal ontology."""
@@ -30,25 +31,24 @@ class LegalPromptGenerator:
                     hierarchies[parent].append(jurisdiction_id)
         return hierarchies
 
-    def generate_system_message(self, 
-                              context_chunks: List[Dict],
-                              query: str,
-                              case_metadata: Optional[Dict] = None) -> str:
+    def generate_system_message(
+        self, context_chunks: List[Dict], query: str, case_metadata: Optional[Dict] = None
+    ) -> str:
         """Generate enhanced system message using legal context and ontology."""
-        
+
         # Extract document types and jurisdictions
         doc_types = set(chunk['metadata'].get('type', 'unknown') for chunk in context_chunks)
         jurisdictions = set(chunk['metadata'].get('jurisdiction', 'unknown') for chunk in context_chunks)
-        
+
         # Get relevant legal concepts from ontology
         relevant_concepts = self._get_relevant_concepts(query, context_chunks)
-        
+
         # Get related jurisdictions
         related_jurisdictions = self._get_related_jurisdictions(jurisdictions)
-        
+
         # Build instruction sets based on document types
         instruction_sets = self._get_type_specific_instructions(doc_types)
-        
+
         # Format the system message
         system_message = f"""You are a specialized legal AI assistant analyzing legal documents with the following context:
 
@@ -62,35 +62,35 @@ LEGAL FRAMEWORK:
 
 RESPONSE GUIDELINES:
 1. Legal Analysis Requirements:
-- Maintain formal legal language and terminology
-- Cite specific sections and precedents when relevant
-- Consider jurisdictional hierarchies and relationships
-- Apply appropriate legal principles based on document type
+- Maintain formal legal language and terminology.
+- Cite specific sections and precedents when relevant.
+- Consider jurisdictional hierarchies and relationships.
+- Apply appropriate legal principles based on document type.
 
 2. Document-Specific Instructions:
 {self._format_instructions(instruction_sets)}
 
 3. Citation and Reference Requirements:
-- Reference specific document sections with clear citations
-- Link analysis to relevant legal concepts and principles
-- Maintain proper legal citation format
-- Include paragraph or page numbers when available
+- Reference specific document sections with clear citations.
+- Link analysis to relevant legal concepts and principles.
+- Maintain proper legal citation format.
+- Include paragraph or page numbers when available.
 
 4. Jurisdictional Considerations:
-- Consider jurisdictional hierarchy and precedence
-- Apply jurisdiction-specific interpretations when relevant
-- Note any cross-jurisdictional implications
+- Consider jurisdictional hierarchy and precedence.
+- Apply jurisdiction-specific interpretations when relevant.
+- Note any cross-jurisdictional implications.
 
 5. Quality Standards:
-- Maintain professional legal terminology
-- Provide balanced analysis considering all relevant factors
-- Structure response in a clear, logical manner
-- Include relevant caveats and limitations
+- Maintain professional legal terminology.
+- Provide balanced analysis considering all relevant factors.
+- Structure response in a clear, logical manner.
+- Include relevant caveats and limitations.
 
 6. Ontological Integration:
-- Incorporate relevant legal concepts from the ontology
-- Link analysis to established legal principles
-- Consider conceptual relationships and hierarchies
+- Incorporate relevant legal concepts from the ontology.
+- Link analysis to established legal principles.
+- Consider conceptual relationships and hierarchies.
 {self._get_case_specific_instructions(case_metadata) if case_metadata else ""}
 """
         return system_message
@@ -106,21 +106,21 @@ CONTEXT:
 {context}
 
 Please provide a comprehensive legal analysis that:
-1. Addresses the specific question
-2. References relevant document sections
-3. Applies appropriate legal principles
-4. Considers jurisdictional implications
-5. Provides clear citations and references"""
+1. Addresses the specific question.
+2. References relevant document sections.
+3. Applies appropriate legal principles.
+4. Considers jurisdictional implications.
+5. Provides clear citations and references."""
 
     def _get_relevant_concepts(self, query: str, context_chunks: List[Dict]) -> List[Dict]:
         """Extract relevant legal concepts from ontology based on query and context."""
         relevant_concepts = []
         combined_text = f"{query} {' '.join(chunk['text'] for chunk in context_chunks)}"
-        
+
         for concept in self.ontology.get("@graph", []):
             if "rdfs:label" not in concept:
                 continue
-            
+
             label = concept["rdfs:label"].lower()
             if label in combined_text.lower():
                 relevant_concepts.append({
@@ -129,7 +129,7 @@ Please provide a comprehensive legal analysis that:
                     "description": concept.get("rdfs:comment", ""),
                     "relationships": concept.get("vocab:relatedConcepts", [])
                 })
-        
+
         return relevant_concepts
 
     def _get_related_jurisdictions(self, jurisdictions: set) -> set:
@@ -148,46 +148,41 @@ Please provide a comprehensive legal analysis that:
 
     def _get_type_specific_instructions(self, doc_types: set) -> Dict[str, List[str]]:
         """Get specific instructions based on document types."""
-        instructions = {}
-        
-        type_instructions = {
+        instructions = {
             "judgment": [
-                "Analyze ratio decidendi and obiter dicta",
-                "Consider precedential value",
-                "Examine judicial reasoning and principles",
-                "Note dissenting opinions if present"
+                "Analyze ratio decidendi and obiter dicta.",
+                "Consider precedential value.",
+                "Examine judicial reasoning and principles.",
+                "Note dissenting opinions if present."
             ],
             "legislation": [
-                "Focus on statutory interpretation",
-                "Consider legislative intent",
-                "Note any amendments or repealed sections",
-                "Examine definitions and scope"
+                "Focus on statutory interpretation.",
+                "Consider legislative intent.",
+                "Note any amendments or repealed sections.",
+                "Examine definitions and scope."
             ],
             "contract": [
-                "Analyze contractual terms and conditions",
-                "Consider contract formation elements",
-                "Examine rights and obligations",
-                "Note any breach or performance issues"
+                "Analyze contractual terms and conditions.",
+                "Consider contract formation elements.",
+                "Examine rights and obligations.",
+                "Note any breach or performance issues."
             ],
             "regulatory": [
-                "Focus on compliance requirements",
-                "Consider regulatory framework",
-                "Examine enforcement mechanisms",
-                "Note reporting obligations"
+                "Focus on compliance requirements.",
+                "Consider regulatory framework.",
+                "Examine enforcement mechanisms.",
+                "Note reporting obligations."
             ]
         }
-        
-        for doc_type in doc_types:
-            if doc_type in type_instructions:
-                instructions[doc_type] = type_instructions[doc_type]
-        
-        return instructions
+
+        return {doc_type: instructions.get(doc_type, ["Apply general legal analysis principles."])
+                for doc_type in doc_types}
 
     def _format_legal_concepts(self, concepts: List[Dict]) -> str:
         """Format legal concepts for system message."""
         if not concepts:
             return "No specific legal concepts identified."
-        
+
         formatted = "Key Legal Concepts:\n"
         for concept in concepts:
             formatted += f"- {concept['concept']}\n"
@@ -201,7 +196,7 @@ Please provide a comprehensive legal analysis that:
         """Format type-specific instructions."""
         if not instruction_sets:
             return "Apply general legal analysis principles."
-        
+
         formatted = ""
         for doc_type, instructions in instruction_sets.items():
             formatted += f"\nFor {doc_type.title()} Documents:\n"
@@ -212,9 +207,8 @@ Please provide a comprehensive legal analysis that:
         """Generate case-specific instructions based on metadata."""
         if not case_metadata:
             return ""
-            
-        return f"""
 
+        return f"""
 CASE-SPECIFIC CONSIDERATIONS:
 - Case Type: {case_metadata.get('case_type', 'Unknown')}
 - Priority: {case_metadata.get('priority', 'Normal')}
@@ -223,11 +217,9 @@ CASE-SPECIFIC CONSIDERATIONS:
 - Tags: {', '.join(case_metadata.get('tags', []))}
 """
 
-    def generate_follow_up_prompt(self, 
-                                original_query: str,
-                                follow_up_query: str,
-                                previous_response: str,
-                                context_chunks: List[Dict]) -> str:
+    def generate_follow_up_prompt(
+        self, original_query: str, follow_up_query: str, previous_response: str, context_chunks: List[Dict]
+    ) -> str:
         """Generate prompt for follow-up questions."""
         return f"""This is a follow-up question to a previous legal inquiry.
 
@@ -241,8 +233,9 @@ Follow-up Question:
 {follow_up_query}
 
 Please provide a response that:
-1. Maintains consistency with the previous analysis
-2. Addresses the specific follow-up inquiry
-3. Builds upon the established legal framework
-4. Provides additional relevant context
-5. References any new relevant documents or principles"""
+1. Maintains consistency with the previous analysis.
+2. Addresses the specific follow-up inquiry.
+3. Builds upon the established legal framework.
+4. Provides additional relevant context.
+5. References any new relevant documents or principles.
+"""