@@ -2,13 +2,16 @@ import os
import sys
import torch
from transformers import pipeline
# adaptive-classifier==0.0.10
from adaptive_classifier import AdaptiveClassifier
sys . path . append ( os . path . dirname ( os . path . dirname ( os . path . abspath ( __file__ ) ) ) )
from sources . agents . agent import Agent
from sources . agents . code_agent import CoderAgent
from sources . agents . casual_agent import CasualAgent
from sources . agents . planner_agent import PlannerAgent
from sources . agents . planner_agent import FileAgent
from sources . agents . browser_agent import BrowserAgent
from sources . language import LanguageUtility
from sources . utility import pretty_print
@@ -17,16 +20,30 @@ class AgentRouter:
"""
AgentRouter is a class that selects the appropriate agent based on the user query.
"""
# TODO add adaptive-classifier==0.0.10 to requirements.txt
def __init__ ( self , agents : list ) :
self . model = model_name
self . agents = agents
self . lang_analysis = LanguageUtility ( )
self . pipelines = {
" fr " : pipeline ( " zero-shot-classification " , model = " facebook/bart-large-mnli " ) ,
" zh " : pipeline ( " zero-shot-classification " , model = " morit/chinese_xlm_xnli " ) ,
" es " : pipeline ( " zero-shot-classification " , model = " facebook/bart-large-mnli " ) ,
" en " : pipeline ( " zero-shot-classification " , model = " facebook/bart-large-mnli " )
" bart " : pipeline ( " zero-shot-classification " , model = " facebook/bart-large-mnli " )
}
self . agents = agents
self . classifier = self . load_llm_router ( )
self . learn_few_shots_en ( )
def load_llm_router ( self ) - > AdaptiveClassifier :
"""
Load the LLM router model.
returns:
AdaptiveClassifier: The loaded model
exceptions:
Exception: If the safetensors fails to load
"""
path = " ../llm_router " if __name__ == " __main__ " else " ./llm_router "
try :
classifier = AdaptiveClassifier . from_pretrained ( path )
except Exception as e :
raise Exception ( " Failed to load the routing model. Please run the dl_safetensors.sh script inside llm_router/ directory to download the model. " )
return classifier
def get_device ( self ) - > str :
if torch . backends . mps . is_available ( ) :
@@ -36,37 +53,184 @@ class AgentRouter:
else :
return " cpu "
def learn_few_shots_en ( self ) - > None :
"""
Few shot learning for the LLM router.
Use the build in add_examples method of the AdaptiveClassifier.
"""
few_shots = [
( " Write a python script to check if the device on my network is connected to the internet " , " coding " ) ,
( " Hey could you search the web for the latest news on the tesla stock market ? " , " web " ) ,
( " I would like you to search for weather api " , " web " ) ,
( " Plan a 3-day trip to New York, including flights and hotels. " , " web " ) ,
( " Find on the web the latest research papers on AI. " , " web " ) ,
( " Can you debug this Java code? It’ s not working. " , " code " ) ,
( " Can you browse the web and find me a 4090 for cheap? " , " web " ) ,
( " i would like to setup a new AI project, index as mark2 " , " files " ) ,
( " Hey, can you find the old_project.zip file somewhere on my drive? " , " files " ) ,
( " Tell me a funny story " , " talk " ) ,
( " Can you locate the backup folder I created last month on my system? " , " files " ) ,
( " Share a random fun fact about space. " , " talk " ) ,
( " Write a script to rename all files in a directory to lowercase. " , " files " ) ,
( " Could you check if the presentation.pdf file exists in my downloads? " , " files " ) ,
( " Tell me about the weirdest dream you’ ve ever heard of. " , " talk " ) ,
( " Search my drive for a file called vacation_photos_2023.jpg. " , " files " ) ,
( " Help me organize my desktop files into folders by type. " , " files " ) ,
( " What’ s your favorite movie and why? " , " talk " ) ,
( " Search my drive for a file named budget_2024.xlsx " , " files " ) ,
( " Write a Python function to sort a list of dictionaries by key " , " code " ) ,
( " Find the latest updates on quantum computing on the web " , " web " ) ,
( " Check if the folder ‘ Work_Projects’ exists on my desktop " , " files " ) ,
( " Create a bash script to monitor CPU usage " , " code " ) ,
( " Search online for the best budget smartphones of 2025 " , " web " ) ,
( " What’ s the strangest food combination you’ ve heard of? " , " talk " ) ,
( " Move all .txt files from Downloads to a new folder called Notes " , " files " ) ,
( " Debug this C++ code that keeps crashing " , " code " ) ,
( " can you browse the web to find out who fosowl is ? " , " web " ) ,
( " Find the file ‘ important_notes.txt’ " , " files " ) ,
( " Find out the latest news on the upcoming Mars mission " , " web " ) ,
( " Write a Java program to calculate the area of a circle " , " code " ) ,
( " Search the web for the best ways to learn a new language " , " web " ) ,
( " Locate the file ‘ presentation.pptx’ in my Documents folder " , " files " ) ,
( " Write a Python script to download all images from a webpage " , " code " ) ,
( " Search the web for the latest trends in AI and machine learning " , " web " ) ,
( " Tell me about a time when you had to solve a difficult problem " , " talk " ) ,
( " Organize all image files on my desktop into a folder called ‘ Pictures’ " , " files " ) ,
( " Generate a Ruby script to calculate Fibonacci numbers up to 100 " , " code " ) ,
( " Find out what device are connected to my network " , " code " ) ,
( " Show me how much disk space is left on my drive " , " code " ) ,
( " Look up recent posts on X about climate change " , " web " ) ,
( " Find the photo I took last week named sunset_beach.jpg " , " files " ) ,
( " Write a JavaScript snippet to fetch data from an API " , " code " ) ,
( " Search the web for tutorials on machine learning with Python " , " web " ) ,
( " Locate the file ‘ meeting_notes.docx’ in my Documents folder " , " files " ) ,
( " Write a Python script to scrape a website’ s title and links " , " code " ) ,
( " Search the web for the latest breakthroughs in fusion energy " , " web " ) ,
( " Tell me about a historical event that sounds too wild to be true " , " talk " ) ,
( " Organize all image files on my desktop into a folder called ‘ Pictures’ " , " files " ) ,
( " Generate a Ruby script to calculate Fibonacci numbers up to 100 " , " code " ) ,
( " Find recent X posts about SpaceX’ s next rocket launch " , " web " ) ,
( " What’ s the funniest misunderstanding you’ ve seen between humans and AI? " , " talk " ) ,
( " Check if ‘ backup_032025.zip’ exists anywhere on my drive " , " files " ) ,
( " Create a shell script to automate backups of a directory " , " code " ) ,
( " Look up the top AI conferences happening in 2025 online " , " web " ) ,
( " Write a C# program to simulate a basic calculator " , " code " ) ,
( " Browse the web for open-source alternatives to Photoshop " , " web " ) ,
( " Hey how are you " , " talk " ) ,
( " Write a Python script to ping a website " , " code " ) ,
( " Search the web for the latest iPhone release " , " web " ) ,
( " What’ s the weather like today? " , " web " ) ,
( " Hi, how’ s your day going? " , " talk " ) ,
( " Can you find a file called resume.docx on my drive? " , " files " ) ,
( " Write a simple Java program to print ' Hello World ' " , " code " ) ,
( " Tell me a quick joke " , " talk " ) ,
( " Search online for the best coffee shops in Seattle " , " web " ) ,
( " Check if ‘ project_plan.pdf’ exists in my Downloads folder " , " files " ) ,
( " What’ s your favorite color? " , " talk " ) ,
( " Write a bash script to list all files in a directory " , " code " ) ,
( " Find recent X posts about electric cars " , " web " ) ,
( " Hey, you doing okay? " , " talk " ) ,
( " Locate the file ‘ family_photo.jpg’ on my system " , " files " ) ,
( " Search the web for beginner guitar lessons " , " web " ) ,
( " Write a Python function to reverse a string " , " code " ) ,
( " What’ s the weirdest animal you know of? " , " talk " ) ,
( " Organize all .pdf files on my desktop into a ‘ Documents’ folder " , " files " ) ,
( " Browse the web for the latest space mission updates " , " web " ) ,
( " Hey, what’ s up with you today? " , " talk " ) ,
( " Write a JavaScript function to add two numbers " , " code " ) ,
( " Find the file ‘ notes.txt’ in my Documents folder " , " files " ) ,
( " Tell me something random about the ocean " , " talk " ) ,
( " Search the web for cheap flights to Paris " , " web " ) ,
( " Check if ‘ budget.xlsx’ is on my drive " , " files " ) ,
( " Write a Python script to count words in a text file " , " code " ) ,
( " How’ s it going today? " , " talk " ) ,
( " Find recent X posts about AI advancements " , " web " ) ,
( " Move all .jpg files from Downloads to a ‘ Photos’ folder " , " files " ) ,
( " Search online for the best laptops of 2025 " , " web " ) ,
( " What’ s the funniest thing you’ ve heard lately? " , " talk " ) ,
( " Write a Ruby script to generate random numbers " , " code " ) ,
( " Hey, how’ s everything with you? " , " talk " ) ,
( " Locate ‘ meeting_agenda.docx’ in my system " , " files " ) ,
( " Search the web for tips on growing indoor plants " , " web " ) ,
( " Write a C++ program to calculate the sum of an array " , " code " ) ,
( " Tell me a fun fact about dogs " , " talk " ) ,
( " Check if the folder ‘ Old_Projects’ exists on my desktop " , " files " ) ,
( " Browse the web for the latest gaming console reviews " , " web " ) ,
( " Hi, how are you feeling today? " , " talk " ) ,
( " Write a Python script to check disk space " , " code " ) ,
( " Find the file ‘ vacation_itinerary.pdf’ on my drive " , " files " ) ,
( " Search the web for news on renewable energy " , " web " ) ,
( " What’ s the strangest thing you’ ve learned recently? " , " talk " ) ,
( " Organize all video files into a ‘ Videos’ folder " , " files " ) ,
( " Write a shell script to delete temporary files " , " code " ) ,
( " Hey, how’ s your week been so far? " , " talk " ) ,
( " Search online for the top movies of 2025 " , " web " ) ,
( " Locate ‘ taxes_2024.xlsx’ in my Documents folder " , " files " ) ,
( " Tell me about a cool invention from history " , " talk " ) ,
( " Write a Java program to check if a number is even or odd " , " code " ) ,
( " Find recent X posts about cryptocurrency trends " , " web " ) ,
( " Hey, you good today? " , " talk " ) ,
( " Search the web for easy dinner recipes " , " web " ) ,
( " Check if ‘ photo_backup.zip’ exists on my drive " , " files " ) ,
( " Write a Python script to rename files with a timestamp " , " code " ) ,
( " What’ s your favorite thing about space? " , " talk " ) ,
( " Browse the web for the latest fitness trends " , " web " ) ,
( " Move all .docx files to a ‘ Work’ folder " , " files " ) ,
]
texts = [ text for text , _ in few_shots ]
labels = [ label for _ , label in few_shots ]
self . classifier . clear_memory ( )
self . classifier . add_examples ( texts , labels )
def llm_router ( self , text : str ) - > tuple :
"""
Inference of the LLM router model.
Args:
text: The input text
"""
predictions = self . classifier . predict ( text )
predictions = [ pred for pred in predictions if pred [ 0 ] not in [ " HIGH " , " LOW " ] ]
predictions = sorted ( predictions , key = lambda x : x [ 1 ] , reverse = True )
return predictions [ 0 ]
def router_vote ( self , text : str , labels : list ) - > str :
"""
Vote between the LLM router and BART model.
Args:
text: The input text
labels: The labels to classify
Returns:
str: The selected label
"""
result_bart = self . pipelines [ ' bart ' ] ( text , labels , threshold = 0.3 )
result_llm_router = self . llm_router ( text )
bart , confidence_bart = result_bart [ ' labels ' ] [ 0 ] , result_bart [ ' scores ' ] [ 0 ]
llm_router , confidence_llm_router = result_llm_router [ 0 ] , result_llm_router [ 1 ]
confidence_bart * = 0.8 # was always a bit too confident
print ( " BART: " , bart , " LLM Router: " , llm_router )
print ( " Confidence BART: " , confidence_bart , " Confidence LLM Router: " , confidence_llm_router )
if confidence_bart > confidence_llm_router :
return bart
else :
return llm_router
def classify_text ( self , text : str , threshold : float = 0.4 ) - > list :
"""
Classify the text into labels (agent roles) .
Args:
text (str): The text to classify
threshold (float, optional): The threshold for the classification.
Returns:
list: The list of agents and their scores
Classify the text using the LLM router and BART model .
"""
first_sentence = None
lang = " en "
for line in text . split ( " \n " ) :
first_sentence = line . strip ( )
break
if first_sentence is None :
first_sentence = text
try :
lang = lang_analysis . detect_language ( first_sentence )
assert lang in [ " en " , " fr " , " zh " , " es " ]
labels = [ agent . role [ lang ] for agent in agents ]
if lang == " en " :
result = self . pipelines [ ' en ' ] ( first_sentence , labels , threshold = threshold )
elif lang == " fr " :
result = self . pipelines [ ' fr ' ] ( first_sentence , labels , threshold = threshold )
elif lang == " zh " :
result = self . pipelines [ ' zh ' ] ( first_sentence , labels , threshold = threshold )
elif lang == " es " :
result = self . pipelines [ ' es ' ] ( first_sentence , labels , threshold = threshold )
else :
result = None
lang = self . lang_analysis . detect_language ( first_sentence )
labels = [ agent . role [ lang ] for agent in self . agents ]
result = self . router_vote ( first_sentence , labels )
except Exception as e :
return None , lang
raise e
return result , lang
def select_agent ( self , text : str ) - > Agent :
@@ -81,20 +245,29 @@ class AgentRouter:
return self . agents [ 0 ]
result , lang = self . classify_text ( text )
for agent in self . agents :
if result [ " labels " ] [ 0 ] == agent . role [ lang ] :
if result == agent . role [ lang ] :
pretty_print ( f " Selected agent: { agent . agent_name } (roles: { agent . role [ lang ] } ) " , color = " warning " )
return agent
pretty_print ( f " Error choosing agent. " , color = " error " )
pretty_print ( f " Error choosing agent. Routing system is not multilingual yet. " , color = " failure " )
pretty_print ( f " 选择代理时出错。路由系统尚不支持多语言 " , color = " failure " )
pretty_print ( f " エージェントの選択エラー。ルーティングシステムはまだ多言語に対応していません " , color = " failure " )
pretty_print ( f " Erreur lors du choix de l ' agent. Le système de routage n ' est pas encore multilingue. " , color = " failure " )
pretty_print ( f " Error al elegir agente. El sistema de enrutamiento aún no es multilingüe. " , color = " failure " )
return None
if __name__ == " __main__ " :
agents = [
CasualAgent ( " deepseek-r1:14b " , " jarvis " , " ../prompts/casual_agent.txt " , None ) ,
BrowserAgent ( " deepseek-r1:14b " , " browser " , " ../prompts/planner_agent.txt " , None ) ,
CoderAgent ( " deepseek-r1:14b " , " coder " , " ../prompts/coder_agent.txt " , None )
CasualAgent ( " jarvis " , " ../prompts/casual_agent.txt " , None ) ,
BrowserAgent ( " browser " , " ../prompts/planner_agent.txt " , None ) ,
CoderAgent ( " coder " , " ../prompts/coder_agent.txt " , None ) ,
FileAgent ( " file " , " ../prompts/coder_agent.txt " , None )
]
router = AgentRouter ( agents )
texts = [
" hi " ,
#"你好",
#"Bonjour",
" Write a python script to check if the device on my network is connected to the internet " ,
# "Peut tu écrire un script python qui vérifie si l'appareil sur mon réseau est connecté à internet?",
# "写一个Python脚本,检查我网络上的设备是否连接到互联网",
@@ -107,10 +280,16 @@ if __name__ == "__main__":
" Plan a 3-day trip to New York, including flights and hotels. " ,
# "计划一次为期3天的纽约之旅,包括机票和酒店。",
# "Planifie un trip de 3 jours à Paris, y compris les vols et hotels.",
" Find me the latest research papers on AI. " ,
#"给我找最新的AI 研究论文。",
#"Trouve moi les derniers papier s de recherche en IA ",
" Find on the web the latest research papers on AI. " ,
# "在网上找到最新的人工智能 研究论文。",
# "Trouve moi les derniers article s de recherche sur l'IA sur internet ",
" Help me write a C++ program to sort an array " ,
" Tell me what France been up to lately " ,
# "告诉我法国最近在做什么",
# "Dis moi ce que la France a fait récemment",
" Who is Sergio Pesto ? " ,
# "谁是Sergio Pesto? ",
# "Qui est Sergio Pesto ?",
# "帮我写一个C++程序来排序数组",
# "Aide moi à faire un programme c++ pour trier une array.",
" What’ s the weather like today? Oh, and can you find a good weather app? " ,
@@ -119,18 +298,15 @@ if __name__ == "__main__":
" Can you debug this Java code? It’ s not working. " ,
# "你能调试这段Java代码吗?它不起作用。",
# "Peut tu m'aider à debugger ce code java, ça marche pas",
" What ' s the latest brainrot on the internet ? " ,
#"互联网上最新的“脑残”是什么 ? ",
#"Quel est la dernière connerie sur internet ?",
" i would like to setup a new AI project, index as mark2 " ,
#"我想建立一个新的 AI 项目,索引为 Mark2",
#"Je voudrais configurer un nouveau projet d'IA, index Mark2",
" Hey, can you find the old_project.zip file somewhere on my drive? " ,
#"Can you browse the web and find me a 4090 for cheap?",
#"你能浏览网页,为我找一个便宜的4090吗 ? ",
#"Peut tu chercher sur internet et me trouver une 4090 pas cher ?",
#"Hey, can you find the old_project.zip file somewhere on my drive?",
#"嘿,你能在我驱动器上找到old_project.zip文件吗?",
#"Hé trouve moi le old_project.zip, il est quelque part sur mon disque.",
" Tell me a funny story " ,
#" 给我讲一个有趣的故事",
#" Raconte moi une histoire drole"
" 给我讲一个有趣的故事 " ,
" Raconte moi une histoire drole "
]
for text in texts :
print ( " Input text: " , text )