Natural Language Processing

class GeminiLLM:
    """Gemini LLM client for query parsing and response generation."""
    
    def __init__(self, api_key: str = None):
        """Initialize Gemini client."""
        self.api_key = api_key or os.getenv('GEMINI_API_KEY')
        if not self.api_key:
            raise ValueError("GEMINI_API_KEY environment variable is required")
        
        self.client = genai.Client(api_key=self.api_key)
        logger.info("Initialized Gemini LLM client")

system_prompt = """
You are an expert at parsing natural language queries about engineering infrastructure and converting them to structured queries.

Given a user query, determine the query type and extract relevant parameters.

Query types and their parameters:
1. "get_node" - Get single node by ID
   - node_id: string
   
2. "get_nodes" - List nodes by type
   - node_type: string (service, database, cache, team, deployment)
   - filters: dict (optional)
   
3. "downstream" - What does X depend on?
   - node_id: string
   - edge_types: list (optional: calls, uses, depends_on)
   
4. "upstream" - What depends on X?
   - node_id: string  
   - edge_types: list (optional: calls, uses, depends_on)
   
5. "blast_radius" - Full impact analysis
   - node_id: string
   
6. "path" - How does X connect to Y?
   - from_id: string
   - to_id: string
   
7. "get_owner" - Who owns X?
   - node_id: string
   
8. "get_team_assets" - What does team X own?
   - team_name: string

9. "find_by_property" - Find nodes with specific property
   - property_name: string
   - property_value: any

10. "services_using_database" - What services use database X?
    - database_name: string

Node ID format: "type:name" (e.g., "service:order-service", "database:users-db")

Respond with JSON only:
{
  "query_type": "...",
  "parameters": {...},
  "confidence": 0.0-1.0,
  "reasoning": "brief explanation"
}
"""

def parse_query_intent(self, user_query: str, context: Dict[str, Any] = None) -> Dict[str, Any]:
    """
    Parse user query to determine intent and extract parameters.
    
    Args:
        user_query: Natural language query from user
        context: Optional conversation context
        
    Returns:
        Dictionary with query type, parameters, and confidence
    """
    user_prompt = f"User query: {user_query}"
    if context:
        user_prompt += f"\nContext: {context}"
    
    try:
        response = self.client.models.generate_content(
            model='gemini-2.5-flash',
            contents=system_prompt + "\n\n" + user_prompt
        )
        
        response_text = response.text.strip()
        
        # Extract JSON from response
        if response_text.startswith('```json'):
            response_text = response_text[7:-3]
        elif response_text.startswith('```'):
            response_text = response_text[3:-3]
        
        parsed_response = json.loads(response_text)
        return parsed_response

class QueryParser:
    """Parses natural language queries and executes graph operations."""
    
    def __init__(self, query_engine: QueryEngine, llm: GeminiLLM):
        self.query_engine = query_engine
        self.llm = llm
        self.conversation_context = {}

def _execute_graph_query(self, intent: Dict[str, Any]) -> Any:
    """Execute the appropriate graph query based on parsed intent."""
    query_type = intent['query_type']
    params = intent['parameters']
    
    try:
        if query_type == 'get_node':
            return self.query_engine.get_node(params['node_id'])
        
        elif query_type == 'get_nodes':
            return self.query_engine.get_nodes(
                params.get('node_type'),
                params.get('filters')
            )
        
        elif query_type == 'downstream':
            return self.query_engine.downstream(
                params['node_id'],
                edge_types=params.get('edge_types')
            )
        
        elif query_type == 'upstream':
            return self.query_engine.upstream(
                params['node_id'],
                edge_types=params.get('edge_types')
            )
        
        elif query_type == 'blast_radius':
            return self.query_engine.blast_radius(params['node_id'])
        
        elif query_type == 'path':
            return self.query_engine.path(
                params['from_id'],
                params['to_id']
            )
        
        elif query_type == 'get_owner':
            return self.query_engine.get_owner(params['node_id'])
        
        elif query_type == 'get_team_assets':
            return self.query_engine.get_team_assets(params['team_name'])

def format_response(self, query_result: Any, user_query: str, query_type: str) -> str:
    """
    Format query results into human-readable response.
    
    Args:
        query_result: Result from graph query
        user_query: Original user query
        query_type: Type of query executed
        
    Returns:
        Human-readable response string
    """
    system_prompt = f"""
You are a helpful assistant explaining engineering infrastructure query results.

The user asked: "{user_query}"
Query type: {query_type}

Format the following query result into a clear, helpful response:
- Use bullet points for lists
- Include relevant details like team ownership, ports, etc.
- Be concise but informative
- If no results, explain why and suggest alternatives
- For blast radius queries, emphasize the impact and affected teams

Query result:
{json.dumps(query_result, indent=2)}

Provide a natural, conversational response.
"""
    
    try:
        response = self.client.models.generate_content(
            model='gemini-2.5-flash',
            contents=system_prompt
        )
        return response.text.strip()

def _update_context(self, session_id: str, user_query: str, intent: Dict[str, Any], result: Any):
    """Update conversation context for follow-up queries."""
    if session_id not in self.conversation_context:
        self.conversation_context[session_id] = {}
    
    context = self.conversation_context[session_id]
    context['last_query'] = user_query
    context['last_intent'] = intent
    context['last_result'] = result
    
    # Extract entities for follow-up queries
    if isinstance(result, list) and result:
        context['last_entities'] = [
            item.get('name') for item in result 
            if isinstance(item, dict) and 'name' in item
        ]
    elif isinstance(result, dict) and 'name' in result:
        context['last_entities'] = [result['name']]