Mayank / router_client.py
Mayank2027's picture
Create router_client.py
286865b verified
Raw
History Blame Contribute Delete
3.31 kB
"""Tool router client - selects appropriate tools for a given task"""
import json
import requests
from typing import List, Dict, Any
class ToolRouter:
"""Router that selects the best tools for a task"""
def __init__(self, use_hosted_api: bool = True):
self.use_hosted_api = use_hosted_api
# If not using hosted API, we'll use a keyword-based fallback
self.tool_keywords = {
"web_search": ["search", "find", "look up", "google", "internet", "online", "browse"],
"fetch_url": ["url", "link", "webpage", "website", "page", "read article"],
"scholar_search": ["paper", "research", "academic", "study", "journal", "arxiv"],
"get_news": ["news", "latest", "headlines", "current events", "breaking"],
"get_trends": ["trending", "popular", "viral", "hot topics", "what's popular"],
"extract_dates": ["date", "timeline", "when", "schedule", "calendar"],
"extract_emails": ["email", "contact", "address"],
"extract_urls": ["link", "url", "reference"],
"summarize": ["summary", "summarize", "shorten", "key points", "tl;dr"],
"wikipedia": ["wiki", "wikipedia", "encyclopedia"],
"weather": ["weather", "temperature", "forecast", "rain", "sunny"],
"calculate": ["calculate", "math", "equation", "solve", "compute"],
"current_time": ["time", "date", "today", "now", "current"],
"translate": ["translate", "language", "translation"],
"stock_price": ["stock", "share", "price", "market", "trading"]
}
def route(self, task: str, k: int = 3) -> List[str]:
"""Return top-k tools for the given task"""
if self.use_hosted_api:
try:
# Try hosted API first (free)
response = requests.post(
"https://dalek-ai-router-api.hf.space/route",
json={"task": task, "k": k},
timeout=5
)
if response.status_code == 200:
return response.json().get("tools", [])
except:
pass # Fall back to keyword matching
# Fallback: keyword-based routing
return self._keyword_route(task, k)
def _keyword_route(self, task: str, k: int) -> List[str]:
"""Simple keyword-based tool selection fallback"""
task_lower = task.lower()
scores = {}
for tool, keywords in self.tool_keywords.items():
score = 0
for keyword in keywords:
if keyword.lower() in task_lower:
score += 1
if score > 0:
scores[tool] = score
# Sort by score and return top k
sorted_tools = sorted(scores.items(), key=lambda x: x[1], reverse=True)
return [tool for tool, score in sorted_tools[:k]]
def get_tool_descriptions(self, tools: List[str]) -> List[Dict[str, str]]:
"""Get descriptions for a list of tool names"""
from tools import TOOL_DESCRIPTIONS
return [
{"name": tool, "description": TOOL_DESCRIPTIONS.get(tool, "No description available")}
for tool in tools
]