src/router.py

"""
Stage 1: Query Router - Intelligent Server Selection
"""

import json
from typing import Dict, Any
from openai import OpenAI


class QueryRouter:
    """Stage 1: Routes queries to appropriate MCP servers"""

    def __init__(self, client: OpenAI, registry: Dict[str, Any]):
        self.client = client
        self.registry = registry

    def route(self, query: str, location: Dict[str, Any]) -> Dict[str, Any]:
        """
        Analyze query and determine which MCP servers are needed
        
        Returns:
            {
                "intent": str,
                "required_servers": List[str],
                "reasoning": str
            }
        """
        # Create registry summary
        registry_text = "Available MCP Servers:\n"
        for server_id, info in self.registry.items():
            registry_text += f"\n{server_id}:\n"
            registry_text += f"  Description: {info['description']}\n"
            registry_text += f"  Use for: {', '.join(info['use_for'][:5])}\n"

        system_prompt = f"""You are a query router for Farmer.chat agricultural system.

Your task: Analyze the farmer's query and select which MCP servers are needed.

{registry_text}

Location: {location['name']} ({location['lat']}°N, {location['lon']}°E)

CRITICAL RULES:
1. Select ALL servers that provide data relevant to answering the query completely
2. Consider IMPLICIT needs - look for context clues in the query
3. Keywords that trigger elevation: "elevation", "slope", "terrain", "my land", "my field", "drainage", "waterlogged", "frost risk", "wind exposure"
4. For crop decisions: ALWAYS include soil_properties + water + weather (comprehensive assessment)
5. For weather risk questions (wind, frost, flood): Include weather + elevation (terrain affects risk)
6. For pest questions with weather context: Include pests + weather
7. Be generous - better to have extra data than miss critical information
8. When farmer mentions location characteristics (height, slope, elevation), ALWAYS include elevation

FEW-SHOT EXAMPLES:

Example 1:
Query: "Are strong winds expected at my land elevation?"
Required: ["weather", "elevation"]
Reasoning: Wind forecast from weather, but elevation affects wind exposure and risk. Farmer explicitly mentions elevation.

Example 2:
Query: "Should I plant rice today?"
Required: ["weather", "soil_properties", "water"]
Reasoning: Planting decisions need weather conditions, soil suitability, and water availability for comprehensive assessment.

Example 3:
Query: "Is there risk of frost tonight?"
Required: ["weather", "elevation"]
Reasoning: Frost risk depends on temperature from weather AND elevation (cold air sinks to lower areas).

Example 4:
Query: "What's my soil composition?"
Required: ["soil_properties"]
Reasoning: Direct soil query, only soil data needed. No implicit needs.

Example 5:
Query: "Can I grow tomatoes here?"
Required: ["soil_properties", "water", "weather"]
Reasoning: Crop suitability requires soil type, water availability, and climate conditions.

Example 6:
Query: "My field gets waterlogged after rain"
Required: ["elevation", "soil_properties", "weather"]
Reasoning: Waterlogging relates to drainage (elevation/slope), soil permeability, and rainfall patterns.

Example 7:
Query: "Should I spray pesticides now?"
Required: ["pests", "weather"]
Reasoning: Need to know pest presence AND weather conditions for optimal application timing.

Example 8:
Query: "How's the weather?"
Required: ["weather"]
Reasoning: Direct weather query, no implicit needs.

Example 9:
Query: "Give me complete farm status"
Required: ["weather", "soil_properties", "water", "elevation", "pests"]
Reasoning: Comprehensive assessment requires all available data sources.

Example 10:
Query: "Will it be too windy on my elevated farm?"
Required: ["weather", "elevation"]
Reasoning: Wind from weather, elevation affects exposure. "Elevated" is explicit context clue.

Response format (JSON only):
{{
  "intent": "brief description of farmer's need",
  "required_servers": ["server_id1", "server_id2"],
  "reasoning": "why these servers"
}}
"""

        try:
            response = self.client.chat.completions.create(
                model="gpt-4o",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": query}
                ],
                temperature=0.3
            )

            result_text = response.choices[0].message.content.strip()
            result_text = result_text.replace("```json", "").replace("```", "").strip()

            routing_decision = json.loads(result_text)
            return routing_decision

        except Exception as e:
            print(f"❌ Routing error: {e}")
            # Fallback - include common servers
            return {
                "intent": "general_inquiry",
                "required_servers": ["weather", "soil_properties", "water"],
                "reasoning": "Fallback routing due to error"
            }