89 lines
3.3 KiB
Python
89 lines
3.3 KiB
Python
import json
|
|
import os
|
|
import re
|
|
from litellm import completion
|
|
from src.config import Config
|
|
from src.crews.definitions import CrewDefinitions
|
|
|
|
class SmartRouter:
|
|
"""
|
|
Decides which Crew should handle a user request using a Fast LLM.
|
|
"""
|
|
|
|
@staticmethod
|
|
def route(user_input: str) -> str:
|
|
"""
|
|
Returns the name of the Crew that should handle the input.
|
|
"""
|
|
|
|
# Get available crews
|
|
crews = CrewDefinitions.get_available_crews()
|
|
crews_str = ", ".join([f"'{c}'" for c in crews])
|
|
|
|
prompt = f"""
|
|
You are an intelligent dispatcher for an IT Enterprise system.
|
|
|
|
AVAILABLE CREWS:
|
|
1. 'Infra Engineering (Zabbix)': For monitoring, server health, Zabbix templates, UUIDs, YAML validation.
|
|
2. 'Security Audit': For hacking, vulnerability checks, penetration testing.
|
|
3. 'HR & Evolution': For creating new agents, learning new rules, system updates.
|
|
4. 'Sales Growth': For sales pitches, CRM, deals, negotiation.
|
|
5. 'Business Strategy': For ROI, legal, compliance, high-level decisions.
|
|
|
|
USER INPUT: "{user_input}"
|
|
|
|
TASK:
|
|
Classify the input into exactly one of the available crews.
|
|
Return ONLY valid JSON like this: {{"crew": "Infra Engineering (Zabbix)"}}
|
|
If unsure, default to "Business Strategy" (Generalist).
|
|
Do NOT include any text before or after the JSON.
|
|
"""
|
|
|
|
config = Config.get_llm_config(mode="fast")
|
|
provider = os.getenv("LLM_PROVIDER", "openai").lower()
|
|
|
|
try:
|
|
# Build completion kwargs - only add response_format for OpenAI
|
|
completion_kwargs = {
|
|
"model": config['model'],
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"temperature": 0.0,
|
|
}
|
|
|
|
# Add optional params based on provider
|
|
if config.get('api_key'):
|
|
completion_kwargs["api_key"] = config['api_key']
|
|
if config.get('base_url'):
|
|
completion_kwargs["base_url"] = config['base_url']
|
|
|
|
# response_format only works with OpenAI-compatible APIs
|
|
if provider == "openai":
|
|
completion_kwargs["response_format"] = {"type": "json_object"}
|
|
|
|
response = completion(**completion_kwargs)
|
|
|
|
content = response.choices[0].message.content.strip()
|
|
|
|
# Robust JSON parsing (handling markdown blocks and extra text)
|
|
# Try to extract JSON from the response
|
|
json_match = re.search(r'\{[^}]+\}', content)
|
|
if json_match:
|
|
content = json_match.group(0)
|
|
elif "```json" in content:
|
|
content = content.split("```json")[1].split("```")[0].strip()
|
|
elif "```" in content:
|
|
content = content.split("```")[1].split("```")[0].strip()
|
|
|
|
data = json.loads(content)
|
|
selected = data.get("crew", "Business Strategy")
|
|
|
|
# Validation
|
|
if selected not in crews:
|
|
return "Business Strategy"
|
|
|
|
return selected
|
|
|
|
except Exception as e:
|
|
print(f"Routing Error: {e}")
|
|
return "Business Strategy" # Fallback
|