fix graph by jurisdiction

agent_api.py

@@ -19,7 +19,8 @@ import secrets
 
 from structured_outputs.api_models import (
     Message, DocumentAnalysis, ChatRequest, ChatResponse,
-    HealthResponse, AnalyzePDFRequest, AnalyzePDFResponse
+    HealthResponse, AnalyzePDFRequest, AnalyzePDFResponse,
+    LawyerProfile
 )
 from langraph_agent import CyberLegalAgent
 from utils.conversation_manager import ConversationManager
@@ -125,21 +126,45 @@ class CyberLegalAPI:
         self.conversation_manager = ConversationManager()
         logger.info(f"🔧 CyberLegalAPI initialized with {llm_provider.upper()} provider")
     
-    def _build_lawyer_prompt(self, document_analyses: Optional[List[DocumentAnalysis]], jurisdiction: str) -> str:
-        """Build lawyer prompt with optional document context"""
-        if not document_analyses:
-            return SYSTEM_PROMPT_LAWYER.format(jurisdiction=jurisdiction)
+    def _build_lawyer_prompt(self, document_analyses: Optional[List[DocumentAnalysis]], jurisdiction: str, lawyer_profile: Optional[LawyerProfile] = None) -> str:
+        """Build lawyer prompt with optional document context and lawyer profile"""
+        prompt_parts = []
         
-        docs_text = "\n\n### Documents parsed in the lawyer profile\n"
-        for i, doc in enumerate(document_analyses, 1):
-            docs_text += f"[Doc {i}] {doc.file_name}\n"
-            if doc.summary: docs_text += f"Summary: {doc.summary}\n"
-            if doc.actors: docs_text += f"Actors: {doc.actors}\n"
-            if doc.key_details: docs_text += f"Key Details: {doc.key_details}\n"
-            docs_text += "\n"
+        # Add lawyer profile context if available
+        if lawyer_profile:
+            profile_text = "\n\n### Lawyer Profile Context\n"
+            if lawyer_profile.full_name:
+                profile_text += f"Name: {lawyer_profile.full_name}\n"
+            if lawyer_profile.primary_specialty:
+                profile_text += f"Primary Specialty: {lawyer_profile.primary_specialty}\n"
+            if lawyer_profile.legal_specialties:
+                profile_text += f"Specialties: {', '.join(lawyer_profile.legal_specialties)}\n"
+            if lawyer_profile.experience_level:
+                profile_text += f"Experience Level: {lawyer_profile.experience_level}\n"
+            if lawyer_profile.languages:
+                profile_text += f"Languages: {', '.join(lawyer_profile.languages)}\n"
+            if lawyer_profile.lawyer_description:
+                profile_text += f"Description: {lawyer_profile.lawyer_description}\n"
+            profile_text += "\nWhen answering, consider this lawyer's expertise and experience level. Tailor your responses to be appropriate for their seniority and specialization.\n"
+            prompt_parts.append(profile_text)
         
-        docs_text += "Use these documents if the user's question is related to their content.\n"
-        return SYSTEM_PROMPT_LAWYER.format(jurisdiction=jurisdiction) + docs_text
+        # Add document analyses if available
+        if document_analyses:
+            docs_text = "\n### Documents parsed in the lawyer profile\n"
+            for i, doc in enumerate(document_analyses, 1):
+                docs_text += f"[Doc {i}] {doc.file_name}\n"
+                if doc.summary: docs_text += f"Summary: {doc.summary}\n"
+                if doc.actors: docs_text += f"Actors: {doc.actors}\n"
+                if doc.key_details: docs_text += f"Key Details: {doc.key_details}\n"
+                docs_text += "\n"
+            docs_text += "Use these documents if the user's question is related to their content.\n"
+            prompt_parts.append(docs_text)
+        
+        # Combine base prompt with context
+        base_prompt = SYSTEM_PROMPT_LAWYER.format(jurisdiction=jurisdiction)
+        if prompt_parts:
+            return base_prompt + "\n".join(prompt_parts)
+        return base_prompt
     
     async def process_request(self, request: ChatRequest) -> ChatResponse:
         """
@@ -173,13 +198,22 @@ class CyberLegalAPI:
         logger.info(f"💬 User query: {request.message}")
         
         try:
-            # Build dynamic system prompt for lawyers with document analyses
-            if request.userType == "lawyer" and request.documentAnalyses:
-                system_prompt = self._build_lawyer_prompt(request.documentAnalyses, request.jurisdiction)
-                logger.info(f"📚 Using lawyer prompt with {len(request.documentAnalyses)} document analyses")
-            elif request.userType == "lawyer":
-                system_prompt = SYSTEM_PROMPT_LAWYER.format(jurisdiction=request.jurisdiction)
-                logger.info(f"📝 Using default lawyer prompt with jurisdiction: {request.jurisdiction}")
+            # Build dynamic system prompt for lawyers with document analyses and/or lawyer profile
+            if request.userType == "lawyer":
+                system_prompt = self._build_lawyer_prompt(
+                    request.documentAnalyses, 
+                    request.jurisdiction, 
+                    request.lawyerProfile
+                )
+                context_parts = []
+                if request.lawyerProfile:
+                    context_parts.append("lawyer profile")
+                if request.documentAnalyses:
+                    context_parts.append(f"{len(request.documentAnalyses)} document analyses")
+                if context_parts:
+                    logger.info(f"📚 Using lawyer prompt with {', '.join(context_parts)}")
+                else:
+                    logger.info(f"📝 Using default lawyer prompt with jurisdiction: {request.jurisdiction}")
             else:
                 system_prompt = SYSTEM_PROMPT_CLIENT.format(jurisdiction=request.jurisdiction)
                 logger.info(f"👤 Using client prompt with jurisdiction: {request.jurisdiction}")

langraph_agent.py

@@ -92,8 +92,11 @@ class CyberLegalAgent:
                     args["conversation_history"] = state.get("conversation_history", [])
                     logger.info(f"📝 Passing conversation_history to {tool_call['name']}: {len(args['conversation_history'])} messages")
                 
+                # Pass jurisdiction to query_knowledge_graph tool
+                if tool_call['name'] == "query_knowledge_graph":
+                    args["jurisdiction"] = state.get("jurisdiction")
                 result = await tool_func.ainvoke(args)
-                logger.info(f"🔧 Tool {tool_call['name']} returned: {result}")
+                logger.info(f"🔧 Tool {tool_call} returned: {result}")
                 intermediate_steps.append(ToolMessage(content=str(result), tool_call_id=tool_call['id'], name=tool_call['name']))
         
         state["intermediate_steps"] = intermediate_steps

structured_outputs/api_models.py

@@ -13,6 +13,17 @@ class Message(BaseModel):
     content: str = Field(..., description="Message content")
 
 
+class LawyerProfile(BaseModel):
+    """Lawyer's professional profile"""
+    full_name: Optional[str] = Field(None, description="Lawyer's full name")
+    primary_specialty: Optional[str] = Field(None, description="Primary legal specialty")
+    legal_specialties: Optional[List[str]] = Field(None, description="List of legal specialties")
+    jurisdiction: Optional[str] = Field(None, description="Primary jurisdiction of practice")
+    experience_level: Optional[str] = Field(None, description="Experience level (e.g., Junior, Mid-level, Senior, Partner)")
+    languages: Optional[List[str]] = Field(None, description="Languages spoken")
+    lawyer_description: Optional[str] = Field(None, description="Professional description or bio")
+
+
 class DocumentAnalysis(BaseModel):
     """Document analysis result"""
     file_name: str
@@ -28,6 +39,7 @@ class ChatRequest(BaseModel):
     userType: Optional[str] = Field(default="client", description="User type: 'client' for general users or 'lawyer' for legal professionals")
     jurisdiction: Optional[str] = Field(default="Romania", description="Jurisdiction of the user")
     documentAnalyses: Optional[List[DocumentAnalysis]] = Field(default=None, description="Lawyer's document analyses")
+    lawyerProfile: Optional[LawyerProfile] = Field(default=None, description="Lawyer's professional profile")
 
 
 class ChatResponse(BaseModel):

utils/lawyer_profile_formatter.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python3
+"""
+Utility functions to transform external API lawyer data into our LawyerProfile format
+"""
+
+import logging
+from typing import List, Optional, Dict, Any
+from structured_outputs.api_models import LawyerProfile
+
+logger = logging.getLogger(__name__)
+
+
+def format_lawyer_profile_from_api(api_lawyer_data: Dict[str, Any]) -> Optional[LawyerProfile]:
+    """Transform external API lawyer data into LawyerProfile model."""
+    try:
+        return LawyerProfile(
+            full_name=api_lawyer_data.get('full_name'),
+            primary_specialty=api_lawyer_data.get('primary_specialty'),
+            legal_specialties=_process_list_field(api_lawyer_data.get('legal_specialties')),
+            jurisdiction=api_lawyer_data.get('jurisdiction'),
+            experience_level=api_lawyer_data.get('experience_level'),
+            languages=_process_list_field(api_lawyer_data.get('languages')),
+            lawyer_description=_process_string_field(api_lawyer_data.get('lawyer_description'))
+        )
+    except Exception as e:
+        logger.error(f"Error formatting lawyer profile: {e}")
+        return None
+
+
+def format_lawyer_profiles_from_response(api_response: Dict[str, Any]) -> List[LawyerProfile]:
+    """Extract and format all lawyers from external API response."""
+    if not isinstance(api_response, dict) or not api_response.get('success'):
+        return []
+    
+    lawyers_data = api_response.get('data', [])
+    profiles = []
+    
+    for idx, lawyer_data in enumerate(lawyers_data):
+        if isinstance(lawyer_data, dict):
+            profile = format_lawyer_profile_from_api(lawyer_data)
+            if profile:
+                profiles.append(profile)
+    
+    logger.info(f"Formatted {len(profiles)} lawyer profiles out of {len(lawyers_data)}")
+    return profiles
+
+
+def _process_list_field(value: Any) -> Optional[List[str]]:
+    """Process list fields, return None if empty or None."""
+    if value is None:
+        return None
+    if isinstance(value, list):
+        return value if value else None
+    logger.warning(f"Expected list, got {type(value)}: {value}")
+    return None
+
+
+def _process_string_field(value: Any) -> Optional[str]:
+    """Process string fields, return None if empty or None."""
+    if value is None:
+        return None
+    if isinstance(value, str):
+        stripped = value.strip()
+        return stripped if stripped else None
+    try:
+        str_value = str(value).strip()
+        return str_value if str_value else None
+    except Exception as e:
+        logger.warning(f"Could not convert to string: {e}")
+        return None
+
+
+def get_lawyer_profile_summary(profile: LawyerProfile) -> str:
+    """Generate human-readable summary of a lawyer profile."""
+    parts = []
+    if profile.full_name:
+        parts.append(profile.full_name)
+    if profile.primary_specialty:
+        parts.append(f"- {profile.primary_specialty}")
+    if profile.experience_level:
+        parts.append(f"({profile.experience_level})")
+    return " ".join(parts) if parts else "Unnamed Lawyer"
+
+
+def format_lawyers_as_string(profiles: List[LawyerProfile]) -> str:
+    """Format multiple lawyer profiles as a single string."""
+    if not profiles:
+        return "No lawyer profiles available."
+    
+    result = ["Lawyer Profiles:", "-" * 80]
+    
+    for idx, profile in enumerate(profiles, 1):
+        result.append(f"\n{idx}. {get_lawyer_profile_summary(profile)}")
+        
+        if profile.primary_specialty:
+            result.append(f"   Specialty: {profile.primary_specialty}")
+        if profile.legal_specialties:
+            result.append(f"   Specialties: {', '.join(profile.legal_specialties)}")
+        if profile.experience_level:
+            result.append(f"   Experience Level: {profile.experience_level}")
+        if profile.jurisdiction:
+            result.append(f"   Jurisdiction: {profile.jurisdiction}")
+        if profile.languages:
+            result.append(f"   Languages: {', '.join(profile.languages)}")
+        if profile.lawyer_description:
+            result.append(f"   Description: {profile.lawyer_description}")
+        
+        if idx < len(profiles):
+            result.append("")
+    
+    return "\n".join(result)

utils/lightrag_client.py

@@ -18,19 +18,78 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # LightRAG configuration
-LIGHTRAG_PORT = int(os.getenv("LIGHTRAG_PORT", "9621"))
 LIGHTRAG_HOST = os.getenv("LIGHTRAG_HOST", "127.0.0.1")
-SERVER_URL = f"http://{LIGHTRAG_HOST}:{LIGHTRAG_PORT}"
+DEFAULT_LIGHTRAG_PORT = int(os.getenv("LIGHTRAG_PORT", "9621"))
 API_KEY = os.getenv("LIGHTRAG_API_KEY")
 
+def parse_jurisdiction_graphs() -> Dict[str, int]:
+    """
+    Parse LIGHTRAG_GRAPHS environment variable to get jurisdiction-to-port mappings.
+    
+    Expected format: "jurisdiction1:port1,jurisdiction2:port2,..."
+    Example: "romania:9621,bahrain:9622"
+    
+    Returns:
+        Dictionary mapping jurisdiction names to their respective ports
+    """
+    graphs_config = os.getenv("LIGHTRAG_GRAPHS", "")
+    jurisdiction_map = {}
+    
+    if not graphs_config:
+        logger.info("No LIGHTRAG_GRAPHS configured, using default port")
+        return jurisdiction_map
+    
+    try:
+        # Parse the comma-separated list
+        for mapping in graphs_config.split(","):
+            mapping = mapping.strip()
+            if ":" in mapping:
+                jurisdiction, port = mapping.split(":", 1)
+                jurisdiction = jurisdiction.strip().lower()
+                port = int(port.strip())
+                jurisdiction_map[jurisdiction] = port
+                logger.info(f"Loaded jurisdiction mapping: {jurisdiction} → port {port}")
+        
+        logger.info(f"Total jurisdictions loaded: {len(jurisdiction_map)}")
+        return jurisdiction_map
+    
+    except Exception as e:
+        logger.error(f"Error parsing LIGHTRAG_GRAPHS: {e}")
+        return jurisdiction_map
+
+
+# Parse jurisdiction mappings at module load time
+JURISDICTION_PORTS = parse_jurisdiction_graphs()
+
+def get_server_url_for_jurisdiction(jurisdiction: Optional[str] = None) -> str:
+    """
+    Get the appropriate server URL based on jurisdiction.
+    
+    Args:
+        jurisdiction: The jurisdiction name (e.g., "romania", "bahrain")
+    
+    Returns:
+        Server URL string
+    """
+    if jurisdiction and jurisdiction.lower() in JURISDICTION_PORTS:
+        port = JURISDICTION_PORTS[jurisdiction.lower()]
+        logger.info(f"Using jurisdiction-specific server: {jurisdiction} → port {port}")
+        return f"http://{LIGHTRAG_HOST}:{port}"
+    else:
+        if jurisdiction:
+            logger.warning(f"Jurisdiction '{jurisdiction}' not found in mappings, using default port {DEFAULT_LIGHTRAG_PORT}")
+        else:
+            logger.info(f"No jurisdiction specified, using default port {DEFAULT_LIGHTRAG_PORT}")
+        return f"http://{LIGHTRAG_HOST}:{DEFAULT_LIGHTRAG_PORT}"
+
 
 class LightRAGClient:
     """
     Client for interacting with LightRAG server
     """
     
-    def __init__(self, server_url: str = SERVER_URL, api_key: Optional[str] = API_KEY):
-        self.server_url = server_url
+    def __init__(self, server_url: Optional[str] = None, api_key: Optional[str] = API_KEY):
+        self.server_url = server_url or get_server_url_for_jurisdiction(None)
         self.api_key = api_key
         self.timeout = 300
     
@@ -149,3 +208,59 @@ class ResponseProcessor:
                 legal_entities.append(reg)
         
         return list(set(legal_entities))  # Remove duplicates
+
+
+_lightrag_client_cache: Dict[str, LightRAGClient] = {}
+
+
+def get_lightrag_client(jurisdiction: Optional[str] = None) -> LightRAGClient:
+    """
+    Get or create a LightRAG client for the specified jurisdiction.
+    Clients are cached to reuse connections.
+    
+    Args:
+        jurisdiction: The jurisdiction name (e.g., "romania", "bahrain")
+                    If None, uses default port
+    
+    Returns:
+        LightRAGClient instance configured for the jurisdiction
+    """
+    cache_key = jurisdiction.lower() if jurisdiction else "default"
+    
+    # Return cached client if available
+    if cache_key in _lightrag_client_cache:
+        logger.debug(f"Using cached LightRAG client for jurisdiction: {cache_key}")
+        return _lightrag_client_cache[cache_key]
+    
+    # Create new client
+    server_url = get_server_url_for_jurisdiction(jurisdiction)
+    client = LightRAGClient(server_url=server_url, api_key=API_KEY)
+    
+    # Cache the client
+    _lightrag_client_cache[cache_key] = client
+    logger.info(f"Created and cached LightRAG client for jurisdiction: {cache_key} → {server_url}")
+    
+    return client
+
+
+def validate_jurisdiction(jurisdiction: str) -> bool:
+    """
+    Validate if a jurisdiction is supported.
+    
+    Args:
+        jurisdiction: The jurisdiction name to validate
+    
+    Returns:
+        True if jurisdiction is configured, False otherwise
+    """
+    return jurisdiction.lower() in JURISDICTION_PORTS
+
+
+def get_available_jurisdictions() -> List[str]:
+    """
+    Get list of available jurisdictions.
+    
+    Returns:
+        List of configured jurisdiction names
+    """
+    return list(JURISDICTION_PORTS.keys())

utils/tools.py

@@ -8,7 +8,7 @@ from typing import List, Dict, Any, Optional
 from langchain_core.tools import tool
 from langchain_tavily import TavilySearch
 from subagents.lawyer_selector import LawyerSelectorAgent
-from utils.lightrag_client import LightRAGClient
+from utils.lightrag_client import LightRAGClient, get_lightrag_client, validate_jurisdiction, get_available_jurisdictions
 import resend
 
 # Global instances - will be initialized in agent_api.py
@@ -18,34 +18,54 @@ tavily_search = None
 resend_api_key: Optional[str] = None
 
 @tool
-async def query_knowledge_graph(query: str, conversation_history: List[Dict[str, str]]) -> str:
+async def query_knowledge_graph(
+    query: str, 
+    conversation_history: List[Dict[str, str]],
+    jurisdiction: Optional[str] = None
+) -> str:
     """
-    Query the legal knowledge graph for relevant information about EU cyber regulations and directives.
+    Query the legal knowledge graph for relevant information about cyber regulations and directives.
     
     This tool searches through a comprehensive knowledge graph containing legal documents,
-    regulations, and directives related to cyber law, GDPR, NIS2, DORA, Cyber Resilience Act, eIDAS 2.0, etc.
+    regulations, and directives related to law
+    
+    The knowledge graph is dynamically selected based on jurisdiction:
+    - Romania: Romanian law documents
+    - Bahrain: Bahraini law documents
+    - Default: Falls back to default port if jurisdiction not specified
     
     Use this tool when answering legal questions to provide accurate, up-to-date information
-    from official EU legal sources.
+    from official legal sources specific to the user's jurisdiction.
     
     Args:
         query: The legal question or topic to search for in the knowledge graph
         conversation_history: Optional conversation history for context (automatically provided by the agent)
+        jurisdiction: The jurisdiction name (e.g., "romania", "bahrain") to query the appropriate graph
     
     Returns:
         Relevant legal information from the knowledge graph with context and references
     """
     try:
-        # Use the globally initialized LightRAG client
-        if lightrag_client is None:
-            raise ValueError("LightRAGClient not initialized. Please initialize it in agent_api.py")
+        # Validate jurisdiction if provided
+        if jurisdiction:
+            jurisdiction = jurisdiction.strip().lower()
+            if not validate_jurisdiction(jurisdiction):
+                available = ", ".join(get_available_jurisdictions())
+                return f"Error: Jurisdiction '{jurisdiction}' is not supported. Available jurisdictions: {available}"
+            
+        # Get the appropriate LightRAG client for the jurisdiction
+        client = get_lightrag_client(jurisdiction)
         
         # Query the knowledge graph
-        result = lightrag_client.query(
+        result = client.query(
             query=query,
             conversation_history=conversation_history
         )
         
+        # Check for errors
+        if "error" in result:
+            return f"Error querying knowledge graph: {result['error']}"
+        
         # Extract the response content
         response = result.get("response", "")