update requirement and example, recover legacy code

This commit is contained in:
yzlin 2024-03-12 16:43:24 +08:00
parent a680a1a82f
commit e960ac8dc8
4 changed files with 30 additions and 6 deletions

View file

@ -62,7 +62,7 @@ class ToolRecommender(BaseModel):
"""
tools: dict[str, Tool] = {}
force: bool = False
force: bool = False # whether to forcedly recommend the specified tools
@field_validator("tools", mode="before")
@classmethod
@ -145,6 +145,26 @@ class ToolRecommender(BaseModel):
return list(valid_tools.values())[:topk]
class TypeMatchToolRecommender(ToolRecommender):
"""
A legacy ToolRecommender using task type matching at the recall stage:
1. Recall: Find tools based on exact match between task type and tool tag;
2. Rank: LLM rank, the same as the default ToolRecommender.
"""
async def recall_tools(self, context: str = "", plan: Plan = None, topk: int = 20) -> list[Tool]:
if not plan:
return list(self.tools.values())[:topk]
# find tools based on exact match between task type and tool tag
task_type = plan.current_task.task_type
candidate_tools = TOOL_REGISTRY.get_tools_by_tag(task_type)
candidate_tool_names = set(self.tools.keys()) & candidate_tools.keys()
recalled_tools = [candidate_tools[tool_name] for tool_name in candidate_tool_names]
return recalled_tools[:topk]
class BM25ToolRecommender(ToolRecommender):
"""
A ToolRecommender using BM25 at the recall stage: