feat : zero shot router for other lang

This commit is contained in:
martin legrand 2025-03-24 20:22:30 +01:00
parent a289ddf1fd
commit af168d2c57
7 changed files with 62 additions and 16 deletions

View File

@ -17,7 +17,12 @@ class BrowserAgent(Agent):
self.tools = { self.tools = {
"web_search": searxSearch(), "web_search": searxSearch(),
} }
self.role = "Web search and navigation" self.role = {
"en": "Web search and navigation",
"fr": "Recherche et navigation web",
"zh": "网络搜索和导航",
"es": "Búsqueda y navegación web"
}
self.type = "browser_agent" self.type = "browser_agent"
self.browser = Browser() self.browser = Browser()
self.current_page = "" self.current_page = ""

View File

@ -18,7 +18,12 @@ class CasualAgent(Agent):
"file_finder": FileFinder(), "file_finder": FileFinder(),
"bash": BashInterpreter() "bash": BashInterpreter()
} }
self.role = "talk" self.role = {
"en": "talking",
"fr": "discuter",
"zh": "聊天",
"es": "discutir"
}
self.type = "casual_agent" self.type = "casual_agent"
def process(self, prompt, speech_module) -> str: def process(self, prompt, speech_module) -> str:

View File

@ -20,7 +20,12 @@ class CoderAgent(Agent):
"go": GoInterpreter(), "go": GoInterpreter(),
"file_finder": FileFinder() "file_finder": FileFinder()
} }
self.role = "Coding task" self.role = {
"en": "Coding task",
"fr": "Tâche de codage",
"zh": "编码任务",
"es": "Tarea de codificación"
}
self.type = "code_agent" self.type = "code_agent"
def process(self, prompt, speech_module) -> str: def process(self, prompt, speech_module) -> str:

View File

@ -14,7 +14,12 @@ class FileAgent(Agent):
"file_finder": FileFinder(), "file_finder": FileFinder(),
"bash": BashInterpreter() "bash": BashInterpreter()
} }
self.role = "find and read files" self.role = {
"en": "find and read files",
"fr": "trouver et lire des fichiers",
"zh": "查找和读取文件",
"es": "encontrar y leer archivos"
}
self.type = "file_agent" self.type = "file_agent"
def process(self, prompt, speech_module) -> str: def process(self, prompt, speech_module) -> str:

View File

@ -21,7 +21,12 @@ class PlannerAgent(Agent):
"file": FileAgent(model, name, prompt_path, provider), "file": FileAgent(model, name, prompt_path, provider),
"web": BrowserAgent(model, name, prompt_path, provider) "web": BrowserAgent(model, name, prompt_path, provider)
} }
self.role = "Research, setup and code" self.role = {
"en": "Research, setup and code",
"fr": "Recherche, configuration et codage",
"zh": "研究,设置和编码",
"es": "Investigación, configuración y code"
}
self.type = "planner_agent" self.type = "planner_agent"
def parse_agent_tasks(self, text): def parse_agent_tasks(self, text):

View File

@ -19,7 +19,7 @@ class LanguageUtility:
text: string to analyze text: string to analyze
Returns: ISO639-1 language code Returns: ISO639-1 language code
""" """
langid.set_languages(['fr', 'en', 'zh', 'es']) # ISO 639-1 codes langid.set_languages(['fr', 'en', 'zh', 'es'])
lang, score = langid.classify(text) lang, score = langid.classify(text)
return lang return lang

View File

@ -10,18 +10,23 @@ from sources.agents.code_agent import CoderAgent
from sources.agents.casual_agent import CasualAgent from sources.agents.casual_agent import CasualAgent
from sources.agents.planner_agent import PlannerAgent from sources.agents.planner_agent import PlannerAgent
from sources.agents.browser_agent import BrowserAgent from sources.agents.browser_agent import BrowserAgent
from sources.language import LanguageUtility
from sources.utility import pretty_print from sources.utility import pretty_print
class AgentRouter: class AgentRouter:
""" """
AgentRouter is a class that selects the appropriate agent based on the user query. AgentRouter is a class that selects the appropriate agent based on the user query.
""" """
def __init__(self, agents: list, model_name: str = "facebook/bart-large-mnli"): def __init__(self, agents: list):
self.model = model_name self.model = model_name
self.pipeline = pipeline("zero-shot-classification", self.lang_analysis = LanguageUtility()
model=self.model) 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")
}
self.agents = agents self.agents = agents
self.labels = [agent.role for agent in agents]
def get_device(self) -> str: def get_device(self) -> str:
if torch.backends.mps.is_available(): if torch.backends.mps.is_available():
@ -46,8 +51,23 @@ class AgentRouter:
break break
if first_sentence is None: if first_sentence is None:
first_sentence = text first_sentence = text
result = self.pipeline(first_sentence, self.labels, threshold=threshold) try:
return result 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
except Exception as e:
return None, lang
return result, lang
def select_agent(self, text: str) -> Agent: def select_agent(self, text: str) -> Agent:
""" """
@ -57,13 +77,14 @@ class AgentRouter:
Returns: Returns:
Agent: The selected agent Agent: The selected agent
""" """
if len(self.agents) == 0 or len(self.labels) == 0: if len(self.agents) == 0:
return self.agents[0] return self.agents[0]
result = self.classify_text(text) result, lang = self.classify_text(text)
for agent in self.agents: for agent in self.agents:
if result["labels"][0] == agent.role: if result["labels"][0] == agent.role[lang]:
pretty_print(f"Selected agent: {agent.agent_name} (roles: {agent.role})", color="warning") pretty_print(f"Selected agent: {agent.agent_name} (roles: {agent.role[lang]})", color="warning")
return agent return agent
pretty_print(f"Error choosing agent.", color="error")
return None return None
if __name__ == "__main__": if __name__ == "__main__":