utils/tools.py

#!/usr/bin/env python3
"""
Tools for the CyberLegal Agent
"""

import os
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 subagents.lawyer_messenger import LawyerMessengerAgent
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
lawyer_selector_agent: Optional[LawyerSelectorAgent] = None
lawyer_messenger_agent: Optional[LawyerMessengerAgent] = None
lightrag_client: Optional[LightRAGClient] = None
tavily_search = None
resend_api_key: Optional[str] = None

@tool
async def query_knowledge_graph(
    query: str
) -> str:
    """
    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 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 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
    """
    return

@tool
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 cyber regulations and directives.
    
    This tool searches through a comprehensive knowledge graph containing legal documents,
    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 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:
        # 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 = 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", "")
        
        return response
    
    except Exception as e:
        return f"Error querying knowledge graph: {str(e)}"


@tool
async def search_web(query: str) -> str:
    """Search the web for current legal updates and news using Tavily."""
    try:
        if tavily_search is None:
            raise ValueError("TavilySearch not initialized in agent_api.py")
        
        result = await tavily_search.ainvoke({"query": query})
        import json
        data = json.loads(result) if isinstance(result, str) else result
        
        output = ["🌐 WEB SEARCH RESULTS", "=" * 80]
        if data.get('answer'):
            output.append(f"\n💡 AI Answer: {data['answer']}")
        
        for i, r in enumerate(data.get('results', []), 1):
            output.append(f"\n📄 Result {i}")
            output.append(f"   Title: {r.get('title', 'N/A')}")
            output.append(f"   URL: {r.get('url', 'N/A')}")
            output.append(f"   Summary: {r.get('content', '')[:300]}...")
        
        return "\n".join(output)
    except Exception as e:
        return f"Error: {str(e)}"

@tool
async def send_email(to_email: str, subject: str, content: str) -> str:
    """Send an email using Resend."""
    try:
        from_email = os.getenv("RESEND_FROM_EMAIL")
        from_name = os.getenv("RESEND_FROM_NAME", "CyberLegalAI")
        
        params = {
            "from": f"{from_name} <{from_email}>",
            "to": [to_email],
            "subject": subject,
            "text": content
        }
        
        response = resend.Emails.send(params)
        return f"✅ Email sent to {to_email} (ID: {response.get('id', 'N/A')})"
    
    except Exception as e:
        return f"❌ Failed: {str(e)}"

@tool
async def find_lawyers() -> str:
    """
    Find the top 3 most suitable lawyers based on the legal issue.
    
    This tool analyzes the conversation context to understand the user's legal problem
    and recommends the best lawyers from the available pool. It provides
    client-friendly explanations of why each lawyer is suitable.
    
    Use this tool when the user asks for lawyer recommendations, mentions needing legal representation,
    or asks about finding a lawyer for their specific legal issue.
    
    Args:
        No args
    Returns:
        A formatted string with the top 3 lawyer recommendations
    """
    return
@tool
async def _find_lawyers(conversation_history: List[Dict[str, str]]) -> str:
    """
    Find the top 3 most suitable lawyers based on the legal issue.
    
    This tool analyzes the conversation context to understand the user's legal problem
    and recommends the best lawyers from the available pool. It provides
    client-friendly explanations of why each lawyer is suitable.
    
    Use this tool when the user asks for lawyer recommendations, mentions needing legal representation,
    or asks about finding a lawyer for their specific legal issue.
    
    Args:
        query: A brief description of the legal issue or request (e.g., "I need a lawyer for a data breach case")
        conversation_history: The full conversation history with the user (automatically provided by the agent)
    
    Returns:
        A formatted string with the top 3 lawyer recommendations
    """
    try:
        if lawyer_selector_agent is None:
            raise ValueError("LawyerSelectorAgent not initialized. Please initialize it in agent_api.py")
        
        return await lawyer_selector_agent.select_lawyers(conversation_history)
    
    except Exception as e:
        return f"Error finding lawyers: {str(e)}"

@tool
async def message_lawyer() -> str:
    """
    Send a message to a lawyer identified from the conversation.
    
    This tool analyzes the conversation to identify which lawyer the client wants to contact,
    extracts the message they want to send, and transmits it to the lawyer through the
    frontend messaging system.
    
    Use this tool when the client clearly expresses intent to contact a specific lawyer
    and provides a message they want to send.
    
    Returns:
        A confirmation message indicating whether the message was sent successfully
    """
    try:
        if lawyer_messenger_agent is None:
            raise ValueError("LawyerMessengerAgent not initialized. Please initialize it in agent_api.py")
        
        # conversation_history and client_id will be injected by the agent from state
        raise ValueError("conversation_history and client_id not provided - these should be injected by the agent")
    
    except Exception as e:
        return f"Error sending message to lawyer: {str(e)}"

@tool
async def _message_lawyer(conversation_history,client_id) -> str:
    """
    Send a message to a lawyer identified from the conversation.
    
    This tool analyzes the conversation to identify which lawyer the client wants to contact,
    extracts the message they want to send, and transmits it to the lawyer through the
    frontend messaging system.
    
    Use this tool when the client clearly expresses intent to contact a specific lawyer
    and provides a message they want to send.
    
    Returns:
        A confirmation message indicating whether the message was sent successfully
    """
    try:
        if lawyer_messenger_agent is None:
            raise ValueError("LawyerMessengerAgent not initialized. Please initialize it in agent_api.py")
        
        
        return await lawyer_messenger_agent.send_lawyer_message(conversation_history,client_id)
    except Exception as e:
        return f"Error sending message to lawyer: {str(e)}"


# Export tool sets for different user types
tools_for_client_facade=[query_knowledge_graph, find_lawyers, message_lawyer, search_web]
tools_for_client = [_query_knowledge_graph, _find_lawyers, _message_lawyer, search_web]
tools_for_lawyer_facade = [query_knowledge_graph, search_web]
tools_for_lawyer = [_query_knowledge_graph, search_web]
tools = tools_for_client