Spaces:
Running
Running
File size: 3,954 Bytes
09801ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import json
import logging
from typing import Dict
from core.llm import chat
logger = logging.getLogger(__name__)
def route_query(query: str, active_modes: list = None) -> Dict:
"""
Agentic Orchestrator: Uses a fast LLM to decide the exact route for the user's query.
Returns a dictionary with routing decision.
"""
if active_modes is None:
active_modes = ["chat", "rag", "graph", "vision", "prediction", "hybrid", "chart", "image", "ml_pipeline"]
prompt = f"""You are the Master Orchestrator for an AI Data Analysis Platform.
Your job is to read the user's query and route it to the CORRECT specialized agent.
AVAILABLE ROUTES:
- "chat": For simple greetings, thanking, small talk ("hi", "who are you", "thanks").
- "rag": For specific questions about the data documents ("what is the revenue in Q1?", "who are the top customers?").
- "chart": If the user explicitly asks for a chart, graph, visualization, or plot.
- "image": If the user explicitly asks for an image, img, png, or picture of a dashboard.
- "vision": If the user says something like "look at this image" or "what is in this picture".
- "prediction": If the user asks for a forecast, prediction, or future trend.
- "ml_pipeline": If the user asks to "train a model", "build a model", "predict churn", "run machine learning", etc.
- "etl": If the user asks to "clean data", "generate etl", "fix missing values", or "write python script to clean".
USER QUERY: "{query}"
Determine the best route.
Output ONLY a valid JSON object in this exact format:
{{
"route": "one_of_the_routes_above",
"confidence": 95,
"reason": "Brief explanation why"
}}
"""
try:
# Use a fast, reliable model for orchestration (e.g., Llama 3 8B)
# Using temperature=0.1 for deterministic routing
response = chat(prompt, temperature=0.1, max_tokens=150)
# Parse JSON
response = response.strip()
if '```json' in response:
response = response.split('```json')[1].split('```')[0]
elif '```' in response:
response = response.split('```')[1].split('```')[0]
start = response.find('{')
end = response.rfind('}') + 1
if start >= 0 and end > start:
response = response[start:end]
decision = json.loads(response)
# Validate route
valid_routes = ["chat", "rag", "graph", "vision", "prediction", "hybrid", "chart", "image", "etl", "ml_pipeline"]
if decision.get("route") not in valid_routes:
decision["route"] = "rag" # fallback
logger.info(f"Orchestrator routed '{query}' -> {decision['route']} (Confidence: {decision.get('confidence')}%)")
return decision
except Exception as e:
logger.error(f"Orchestrator failed: {e}")
# Fallback logic if LLM routing fails
query_lower = query.lower()
if any(kw in query_lower for kw in ['img', 'image', 'picture', 'dashboard']):
return {"route": "image", "confidence": 50, "reason": "Fallback: Keyword match"}
if any(kw in query_lower for kw in ['chart', 'graph', 'plot', 'visualize', 'trend']):
return {"route": "chart", "confidence": 50, "reason": "Fallback: Keyword match"}
if any(kw in query_lower for kw in ['predict', 'forecast', 'future']):
return {"route": "prediction", "confidence": 50, "reason": "Fallback: Keyword match"}
if any(kw in query_lower for kw in ['clean', 'fix', 'etl', 'python script to clean']):
return {"route": "etl", "confidence": 50, "reason": "Fallback: Keyword match"}
if any(kw in query_lower for kw in ['hi', 'hello', 'hey', 'thanks', 'bye']):
return {"route": "chat", "confidence": 50, "reason": "Fallback: Keyword match"}
return {"route": "rag", "confidence": 50, "reason": "Fallback: Default route"}
|