lawyer_messenger simplification

subagents/lawyer_messenger.py

@@ -4,15 +4,15 @@ Lawyer Messenger Agent - Identifies lawyer from conversation and sends message
 """
 
 import os
-import json
-import aiohttp
-from typing import List, Optional, Literal
+from typing import List
 from dotenv import load_dotenv
 from langgraph.graph import StateGraph, END
-from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
-from pydantic import BaseModel, Field
+from langchain_core.messages import HumanMessage
 from logging import getLogger
 from prompts.lawyer_messenger import LAWYER_MESSENGER_PROMPT
+from utils.lawyer_messenger_utils import (
+    call_api, convert_messages, format_lawyers_for_prompt, create_lawyer_message_model
+)
 
 logger = getLogger(__name__)
 load_dotenv()
@@ -36,145 +36,44 @@ class LawyerMessengerAgent:
         workflow.add_edge("send_message", END)
         return workflow.compile()
     
-    def _format_lawyers(self, lawyers: List[dict]) -> str:
-        return "\n\n".join([
-            f"Lawyer ID: {l['lawyer_id']}\n- Name: {l['full_name']}\n- Specialty: {l['primary_specialty']}\n- Experience Level: {l.get('experience_level', 'N/A')}\n- Years: {l.get('experience_years', 'N/A')}\n- Description: {l.get('lawyer_description', 'N/A')[:200]}..."
-            for l in lawyers
-        ])
-    
-    async def _fetch_lawyers_from_frontend(self) -> List[dict]:
-        """Fetch lawyers from frontend API"""
-        base_url = os.getenv("SUPABASE_BASE_URL")
-        if not base_url:
-            raise Exception("SUPABASE_BASE_URL not configured in environment")
-        API_URL = f"{base_url}/functions/v1/get-lawyer-database"
-        API_KEY = os.getenv("CYBERLGL_API_KEY")
-        
-        try:
-            async with aiohttp.ClientSession() as session:
-                async with session.get(
-                    API_URL,
-                    headers={"X-API-Key": API_KEY}
-                ) as response:
-                    if response.status != 200:
-                        error_text = await response.text()
-                        raise Exception(f"Frontend API error: {response.status} - {error_text}")
-                    data = await response.json()
-                    if data.get('success') and data.get('data'):
-                        return data['data']
-                    else:
-                        raise Exception(f"Frontend API returned unexpected format: {data}")
-        except aiohttp.ClientError as e:
-            raise Exception(f"Failed to connect to frontend API: {str(e)}")
-        except Exception as e:
-            raise Exception(f"Error fetching lawyers from frontend: {str(e)}")
-    
     async def _fetch_lawyers(self, state: dict) -> dict:
         """Fetch lawyers from frontend API"""
-        try:
-            lawyers = await self._fetch_lawyers_from_frontend()
-            state["lawyers"] = lawyers
-            logger.info(f"Fetched {len(lawyers)} lawyers")
-            return state
-        except Exception as e:
-            logger.error(f"Failed to fetch lawyers: {str(e)}")
-            state["error"] = str(e)
-            state["result"] = f"❌ Error: {str(e)}"
-            state["message_sent"] = False
-            # Skip to end
-            return state
+        base_url = os.getenv("SUPABASE_BASE_URL")
+        data = await call_api(
+            "GET",
+            f"{base_url}/functions/v1/get-lawyer-database",
+            os.getenv("CYBERLGL_API_KEY")
+        )
+        
+        state["lawyers"] = data['data']
+        return state
     
     async def _identify_lawyer(self, state: dict) -> dict:
         """Use LLM to identify lawyer and extract message from conversation"""
-        
-        # Check if we have lawyers
         lawyers = state.get("lawyers", [])
-        if not lawyers or state.get("error"):
-            state["lawyer_selection"] = None
-            state["result"] = "No lawyers available or error fetching lawyers."
-            state["message_sent"] = False
-            return state
-        
-        # Extract valid lawyer IDs
         valid_lawyer_ids = [l['lawyer_id'] for l in lawyers]
-        logger.info(f"Available lawyer IDs: {valid_lawyer_ids}")
-        
-        # Create lawyer map for lookup
         lawyer_map = {l['lawyer_id']: l for l in lawyers}
         
-        # Dynamically create Pydantic model with Literal types for ID validation
-        LawyerID = Literal[tuple(valid_lawyer_ids)] if len(valid_lawyer_ids) > 0 else str
-        
-        class LawyerMessage(BaseModel):
-            lawyer_id: LawyerID = Field(description=f"The unique ID of the lawyer from the retrieved list")
-            subject: str = Field(description="Subject line for the message")
-            message: str = Field(description="Message content to send to the lawyer")
+        LawyerMessage = create_lawyer_message_model(valid_lawyer_ids)
+        lawyers_text = format_lawyers_for_prompt(lawyers)
         
-        # Format lawyers for prompt
-        lawyers_text = self._format_lawyers(lawyers)
-        prompt = LAWYER_MESSENGER_PROMPT.format(lawyers=lawyers_text)
+        messages = convert_messages(state["conversation_history"])
+        messages.append(HumanMessage(content=LAWYER_MESSENGER_PROMPT.format(lawyers=lawyers_text)))
         
-        # Convert message dicts to Message objects
-        messages = []
-        for msg in state["conversation_history"]:
-            role = msg.get("role") if isinstance(msg, dict) else msg.role
-            content = msg.get("content") if isinstance(msg, dict) else msg.content
-            
-            if role == "system":
-                messages.append(SystemMessage(content=content))
-            elif role == "user":
-                messages.append(HumanMessage(content=content))
-            elif role == "assistant":
-                messages.append(AIMessage(content=content))
+        result = await self.llm.with_structured_output(LawyerMessage).ainvoke(messages)
         
-        # Append messenger prompt
-        messages.append(HumanMessage(content=prompt))
-        
-        # Use dynamic structured output
-        try:
-            result = await self.llm.with_structured_output(LawyerMessage).ainvoke(messages)
-            logger.info(f"Identified lawyer: {result.lawyer_id}")
-            logger.info(f"Subject: {result.subject}")
-            logger.info(f"Message length: {len(result.message)}")
-            
-            state["lawyer_selection"] = {
-                "lawyer_id": result.lawyer_id,
-                "subject": result.subject,
-                "message": result.message,
-                "lawyer": lawyer_map.get(result.lawyer_id)
-            }
-            return state
-        except Exception as e:
-            logger.error(f"Structured output error: {str(e)}")
-            state["lawyer_selection"] = None
-            state["error"] = str(e)
-            state["result"] = f"❌ Error identifying lawyer: {str(e)}"
-            state["message_sent"] = False
-            return state
+        state["lawyer_selection"] = {
+            "lawyer_id": result.lawyer_id,
+            "subject": result.subject,
+            "message": result.message,
+            "lawyer": lawyer_map.get(result.lawyer_id)
+        }
+        return state
     
     async def _send_message(self, state: dict) -> dict:
         """Send message to lawyer via frontend API"""
-        
-        # Check if we have a lawyer selection
         lawyer_selection = state.get("lawyer_selection")
-        
-        if not lawyer_selection or not lawyer_selection.get("lawyer_id"):
-            state["result"] = "❌ No lawyer identified. The client hasn't clearly specified which lawyer to contact."
-            state["message_sent"] = False
-            return state
-        
-        # Check for empty message (LLM decided not to send)
-        if not lawyer_selection.get("message") or lawyer_selection.get("message").strip() == "":
-            state["result"] = "❌ No message to send. The client hasn't provided a clear message."
-            state["message_sent"] = False
-            return state
-        
-        # Prepare API request
         base_url = os.getenv("SUPABASE_BASE_URL")
-        if not base_url:
-            raise Exception("SUPABASE_BASE_URL not configured in environment")
-        API_URL = f"{base_url}/functions/v1/send-message-agent"
-        API_KEY = os.getenv("CYBERLGL_API_KEY")
         
         payload = {
             "userId": state["user_id"],
@@ -183,49 +82,24 @@ class LawyerMessengerAgent:
             "message": lawyer_selection["message"]
         }
         
-        try:
-            async with aiohttp.ClientSession() as session:
-                async with session.post(
-                    API_URL,
-                    headers={
-                        "x-api-key": API_KEY,
-                        "Content-Type": "application/json"
-                    },
-                    json=payload
-                ) as response:
-                    if response.status != 200:
-                        error_text = await response.text()
-                        raise Exception(f"Send message API error: {response.status} - {error_text}")
-                    
-                    result_data = await response.json()
-                    logger.info(f"Message sent successfully: {result_data}")
-                    
-                    # Format success response
-                    lawyer = lawyer_selection.get("lawyer", {})
-                    lawyer_name = lawyer.get("full_name", "Unknown")
-                    
-                    state["result"] = (
-                        f"✅ Message sent successfully!\n\n"
-                        f"📨 To: {lawyer_name}\n"
-                        f"📌 Lawyer ID: {lawyer_selection['lawyer_id']}\n"
-                        f"📝 Subject: {lawyer_selection['subject']}\n\n"
-                        f"The lawyer will receive your message and respond to you shortly."
-                    )
-                    state["message_sent"] = True
-                    return state
-                    
-        except aiohttp.ClientError as e:
-            logger.error(f"Failed to send message: {str(e)}")
-            state["result"] = f"❌ Failed to send message: {str(e)}"
-            state["error"] = str(e)
-            state["message_sent"] = False
-            return state
-        except Exception as e:
-            logger.error(f"Error sending message: {str(e)}")
-            state["result"] = f"❌ Error sending message: {str(e)}"
-            state["error"] = str(e)
-            state["message_sent"] = False
-            return state
+        await call_api(
+            "POST",
+            f"{base_url}/functions/v1/send-message-agent",
+            os.getenv("CYBERLGL_API_KEY"),
+            payload
+        )
+        
+        lawyer_name = lawyer_selection.get("lawyer", {}).get("full_name", "Unknown")
+        state["result"] = (
+            f"✅ Message sent successfully!\n\n"
+            f"📨 To: {lawyer_name}\n"
+            f"📌 Lawyer ID: {lawyer_selection['lawyer_id']}\n"
+            f"📝 Subject: {lawyer_selection['subject']}\n\n"
+            f"The lawyer will receive your message and respond shortly."
+        )
+        state["message_sent"] = True
+        
+        return state
     
     async def send_lawyer_message(self, conversation_history: List[dict], user_id: str) -> str:
         """Main entry point: identify lawyer and send message"""
@@ -238,4 +112,4 @@ class LawyerMessengerAgent:
             "result": "",
             "error": None
         })
-        return result["result"]
+        return result["result"]

utils/lawyer_messenger_utils.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+"""
+Utility functions for Lawyer Messenger Agent
+"""
+
+import os
+import aiohttp
+from typing import List, Optional, Literal
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+from pydantic import BaseModel, Field
+
+
+async def call_api(
+    method: str,
+    url: str,
+    api_key: str,
+    payload: Optional[dict] = None,
+    headers: Optional[dict] = None
+) -> dict:
+    """
+    Generic API call handler for GET and POST requests
+    
+    Args:
+        method: HTTP method ('GET' or 'POST')
+        url: API endpoint URL
+        api_key: API key for authentication
+        payload: Optional JSON payload for POST requests
+        headers: Optional additional headers
+    
+    Returns:
+        Response data as dict
+    
+    Raises:
+        Exception: If API call fails or returns non-200 status
+    """
+    default_headers = {
+        "X-API-Key": api_key,
+        "Content-Type": "application/json"
+    }
+    
+    if headers:
+        default_headers.update(headers)
+    
+    try:
+        async with aiohttp.ClientSession() as session:
+            if method.upper() == "GET":
+                async with session.get(url, headers=default_headers) as response:
+                    if response.status != 200:
+                        error_text = await response.text()
+                        raise Exception(f"API error: {response.status} - {error_text}")
+                    return await response.json()
+            
+            elif method.upper() == "POST":
+                async with session.post(url, headers=default_headers, json=payload) as response:
+                    if response.status != 200:
+                        error_text = await response.text()
+                        raise Exception(f"API error: {response.status} - {error_text}")
+                    return await response.json()
+    
+    except aiohttp.ClientError as e:
+        raise Exception(f"Connection failed: {str(e)}")
+
+
+def convert_messages(conversation_history: List[dict]) -> list:
+    """
+    Convert conversation history dict to LangChain Message objects
+    
+    Args:
+        conversation_history: List of message dicts with 'role' and 'content'
+    
+    Returns:
+        List of LangChain Message objects
+    """
+    messages = []
+    for msg in conversation_history:
+        role = msg.get("role")
+        content = msg.get("content")
+        
+        if role == "system":
+            messages.append(SystemMessage(content=content))
+        elif role == "user":
+            messages.append(HumanMessage(content=content))
+        elif role == "assistant":
+            messages.append(AIMessage(content=content))
+    
+    return messages
+
+
+def format_lawyers_for_prompt(lawyers: List[dict]) -> str:
+    """
+    Format lawyer list for LLM prompt
+    
+    Args:
+        lawyers: List of lawyer dicts
+    
+    Returns:
+        Formatted string with lawyer information
+    """
+    return "\n\n".join([
+        f"Lawyer ID: {l['lawyer_id']}\n"
+        f"- Name: {l['full_name']}\n"
+        f"- Specialty: {l['primary_specialty']}\n"
+        f"- Experience: {l.get('experience_level', 'N/A')} ({l.get('experience_years', 'N/A')} years)\n"
+        f"- Description: {l.get('lawyer_description', 'N/A')[:200]}..."
+        for l in lawyers
+    ])
+
+
+def create_lawyer_message_model(valid_lawyer_ids: List[str]) -> BaseModel:
+    """
+    Dynamically create Pydantic model for lawyer message selection
+    
+    Args:
+        valid_lawyer_ids: List of valid lawyer IDs
+    
+    Returns:
+        Pydantic BaseModel with Literal types for ID validation
+    """
+    LawyerID = Literal[tuple(valid_lawyer_ids)] if len(valid_lawyer_ids) > 0 else str
+    
+    class LawyerMessage(BaseModel):
+        lawyer_id: LawyerID = Field(description="Unique ID of the selected lawyer")
+        subject: str = Field(description="Subject line for the message")
+        message: str = Field(description="Message content to send to the lawyer")
+    
+    return LawyerMessage