|
from __future__ import annotations |
|
|
|
from typing import Any, List, Optional, Type |
|
|
|
from erniebot_agent.messages import AIMessage, HumanMessage |
|
from erniebot_agent.tools.schema import ToolParameterView |
|
from pydantic import Field |
|
|
|
from .base import Tool |
|
|
|
|
|
class BaizhongSearchToolInputView(ToolParameterView): |
|
query: str = Field(description="Query") |
|
top_k: int = Field(description="Number of results to return") |
|
|
|
|
|
class SearchResponseDocument(ToolParameterView): |
|
id: str = Field(description="text id") |
|
title: str = Field(description="title") |
|
document: str = Field(description="content") |
|
|
|
|
|
class BaizhongSearchToolOutputView(ToolParameterView): |
|
documents: List[SearchResponseDocument] = Field(description="research results") |
|
|
|
|
|
class BaizhongSearchTool(Tool): |
|
description: str = "aurora search tool" |
|
input_type: Type[ToolParameterView] = BaizhongSearchToolInputView |
|
ouptut_type: Type[ToolParameterView] = BaizhongSearchToolOutputView |
|
|
|
def __init__(self, description, db, input_type=None, output_type=None, examples=None) -> None: |
|
super().__init__() |
|
self.db = db |
|
self.description = description |
|
if input_type is not None: |
|
self.input_type = input_type |
|
if output_type is not None: |
|
self.ouptut_type = output_type |
|
if examples is not None: |
|
self.few_shot_examples = examples |
|
|
|
async def __call__(self, query: str, top_k: int = 10, filters: Optional[dict[str, Any]] = None): |
|
res = self.db.search(query, top_k, filters) |
|
return res |
|
|
|
@property |
|
def examples( |
|
self, |
|
) -> List[Any]: |
|
few_shot_objects: List[Any] = [] |
|
for item in self.few_shot_examples: |
|
few_shot_objects.append(HumanMessage(item["user"])) |
|
few_shot_objects.append( |
|
AIMessage( |
|
"", |
|
function_call={ |
|
"name": self.tool_name, |
|
"thoughts": item["thoughts"], |
|
"arguments": item["arguments"], |
|
}, |
|
) |
|
) |
|
|
|
return few_shot_objects |
|
|
