solana-agent 20.1.2__py3-none-any.whl → 31.4.0__py3-none-any.whl

solana_agent/services/routing.py

@@ -1,15 +1,15 @@
-"""
-Routing service implementation.
-
-This service manages query routing to appropriate agents based on
-specializations and query analysis.
-"""
+import logging
 from typing import Dict, List, Optional, Any
-from solana_agent.interfaces.services.routing import RoutingService as RoutingServiceInterface
+from solana_agent.interfaces.services.routing import (
+    RoutingService as RoutingServiceInterface,
+)
 from solana_agent.interfaces.services.agent import AgentService
 from solana_agent.interfaces.providers.llm import LLMProvider
 from solana_agent.domains.routing import QueryAnalysis
 
+# Setup logger for this module
+logger = logging.getLogger(__name__)
+
 
 class RoutingService(RoutingServiceInterface):
     """Service for routing queries to appropriate agents."""
@@ -18,15 +18,34 @@ class RoutingService(RoutingServiceInterface):
         self,
         llm_provider: LLMProvider,
         agent_service: AgentService,
-    ):
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        model: Optional[str] = None,
+    ) -> None:
         """Initialize the routing service.
 
         Args:
             llm_provider: Provider for language model interactions
             agent_service: Service for agent management
+            api_key: Optional API key for custom LLM provider
+            base_url: Optional base URL for custom LLM provider (e.g., Grok)
+            model: Optional model name to use for routing
         """
         self.llm_provider = llm_provider
         self.agent_service = agent_service
+        self.api_key = api_key
+        self.base_url = base_url
+        # Use provided model, or default based on whether using custom provider
+        if model:
+            self.model = model
+        elif base_url:
+            # Using custom provider (e.g., Grok) but no model specified - use provider's default
+            self.model = None  # Will use adapter's default
+        else:
+            # Using OpenAI - default to small, cheap model for routing
+            self.model = "gpt-4.1-mini"
+        # Simple sticky session: remember last routed agent in-process
+        self._last_agent = None
 
     async def _analyze_query(self, query: str) -> Dict[str, Any]:
         """Analyze a query to determine routing information.
@@ -42,57 +61,74 @@ class RoutingService(RoutingServiceInterface):
         available_specializations = []
 
         for agent_id, agent in agents.items():
-            available_specializations.append({
-                "agent_name": agent_id,
-                "specialization": agent.specialization,
-            })
+            available_specializations.append(
+                {
+                    "agent_name": agent_id,
+                    "specialization": agent.specialization,
+                }
+            )
 
-        specializations_text = "\n".join([
-            f"- {spec['agent_name']}: {spec['specialization']}"
-            for spec in available_specializations
-        ])
+        specializations_text = "\n".join(
+            [
+                f"- {spec['agent_name']}: {spec['specialization']}"
+                for spec in available_specializations
+            ]
+        )
 
         prompt = f"""
         Analyze this user query and determine which agent would be best suited to answer it.
-        
+
         AVAILABLE AGENTS AND THEIR SPECIALIZATIONS:
         {specializations_text}
-        
+
         USER QUERY: {query}
-        
-        Please determine:
-        1. Which agent is the primary best match for this query (must be one of the listed agents)
-        2. Any secondary agents that might be helpful (must be from the listed agents)
-        3. The complexity level (1-5, where 5 is most complex)
-        4. Any key topics or technologies mentioned
-        
-        Think carefully about whether the query is more technical/development-focused or more 
-        financial/market-focused to match with the appropriate agent.
+
+        ROUTING RULES:
+        - Match the user query to the agent whose specialization best fits the user's intent
+        - Return the EXACT agent name that matches best
+
+        INSTRUCTIONS:
+        - primary_agent: The exact name of the best matching agent (e.g., "onboarding", "event_feedback")
+        - secondary_agents: Other agents that might help (usually empty)
+        - complexity_level: 1-5 (5 being most complex)
+        - topics: Key topics mentioned
+        - confidence: 0.0-1.0 (how confident you are in this routing decision)
+
+        For the query "{query}", which agent should handle it?
         """
 
         try:
             analysis = await self.llm_provider.parse_structured_output(
                 prompt=prompt,
-                system_prompt="Match user queries to the most appropriate agent based on specializations.",
+                system_prompt="You are an expert at routing user queries to the most appropriate AI agent. Always return the exact agent name that best matches the user's needs based on the specializations provided. If the user mentions a specific topic, prioritize agents whose specialization matches that topic.",
                 model_class=QueryAnalysis,
+                api_key=self.api_key,
+                base_url=self.base_url,
+                model=self.model,
             )
 
+            logger.debug(f"LLM analysis result: {analysis}")
+
             return {
-                "primary_specialization": analysis.primary_specialization,
-                "secondary_specializations": analysis.secondary_specializations,
+                "primary_specialization": analysis.primary_agent,
+                "secondary_specializations": analysis.secondary_agents,
                 "complexity_level": analysis.complexity_level,
                 "topics": analysis.topics,
-                "confidence": analysis.confidence
+                "confidence": analysis.confidence,
             }
         except Exception as e:
-            print(f"Error analyzing query: {e}")
+            logger.error(f"Error analyzing query: {e}")
+            logger.debug(f"Query that failed: {query}")
+            logger.debug(f"Available agents: {list(agents.keys())}")
             # Return default analysis on error
+            first_agent = list(agents.keys())[0] if agents else "general"
+            logger.debug(f"Defaulting to first agent: {first_agent}")
             return {
-                "primary_specialization": list(agents.keys())[0] if agents else "general",
+                "primary_specialization": first_agent,
                 "secondary_specializations": [],
                 "complexity_level": 1,
                 "topics": [],
-                "confidence": 0.0
+                "confidence": 0.0,
             }
 
     async def route_query(self, query: str) -> str:  # pragma: no cover
@@ -107,20 +143,27 @@ class RoutingService(RoutingServiceInterface):
         # If only one agent - use that agent
         agents = self.agent_service.get_all_ai_agents()
         if len(agents) == 1:
-            print(f"Only one agent available: {next(iter(agents.keys()))}")
-            return next(iter(agents.keys()))
-
-        # Analyze query
+            agent_name = next(iter(agents.keys()))
+            logger.info(f"Only one agent available: {agent_name}")  # Use logger.info
+            self._last_agent = agent_name
+            return agent_name
+
+        # Short reply bypass and default stickiness
+        short = query.strip().lower()
+        short_replies = {"", "yes", "no", "ok", "k", "y", "n", "1", "0"}
+        if short in short_replies and self._last_agent:
+            return self._last_agent
+
+        # Always analyze with a small model to select the best agent
         analysis = await self._analyze_query(query)
-
-        # Find best agent based on analysis
+        logger.debug(f"Routing analysis for query '{query}': {analysis}")
         best_agent = await self._find_best_ai_agent(
-            analysis["primary_specialization"],
-            analysis["secondary_specializations"]
+            analysis["primary_specialization"], analysis["secondary_specializations"]
         )
-
-        # Return best agent
-        return best_agent
+        logger.debug(f"Selected agent: {best_agent}")
+        chosen = best_agent or next(iter(agents.keys()))
+        self._last_agent = chosen
+        return chosen
 
     async def _find_best_ai_agent(
         self,
@@ -143,6 +186,7 @@ class RoutingService(RoutingServiceInterface):
 
         # First, check if primary_specialization is directly an agent name
         if primary_specialization in ai_agents:
+            logger.debug(f"Direct agent match: {primary_specialization}")
             return primary_specialization
 
         # If not a direct agent name match, use specialization matching
@@ -152,24 +196,31 @@ class RoutingService(RoutingServiceInterface):
             score = 0
 
             # Check for specialization match
-            if agent.specialization.lower() in primary_specialization.lower() or \
-                    primary_specialization.lower() in agent.specialization.lower():
+            if (
+                agent.specialization.lower() in primary_specialization.lower()
+                or primary_specialization.lower() in agent.specialization.lower()
+            ):
                 score += 10
+                logger.debug(
+                    f"Specialization match for {agent_id}: '{agent.specialization}' matches '{primary_specialization}'"
+                )
 
             # Check secondary specializations
             for sec_spec in secondary_specializations:
                 if sec_spec in ai_agents:  # Direct agent name match
                     if sec_spec == agent_id:
                         score += 5
-                elif agent.specialization.lower() in sec_spec.lower() or \
-                        sec_spec.lower() in agent.specialization.lower():
+                elif (
+                    agent.specialization.lower() in sec_spec.lower()
+                    or sec_spec.lower() in agent.specialization.lower()
+                ):
                     score += 3
 
             agent_scores.append((agent_id, score))
 
         # Sort by score
         agent_scores.sort(key=lambda x: x[1], reverse=True)
-        print(f"Agent scores: {agent_scores}")
+        logger.debug(f"Agent scores: {agent_scores}")  # Use logger.debug
 
         # Return the highest scoring agent, if any
         if agent_scores and agent_scores[0][1] > 0:
@@ -177,6 +228,8 @@ class RoutingService(RoutingServiceInterface):
 
         # If no match found, return first agent as fallback
         if ai_agents:
-            return next(iter(ai_agents.keys()))
+            fallback_agent = next(iter(ai_agents.keys()))
+            logger.debug(f"No match found, using fallback agent: {fallback_agent}")
+            return fallback_agent
 
         return None