add tool resource

This commit is contained in:
jyong 2024-05-09 18:24:51 +08:00
parent 8137d63000
commit 95fae0438d
9 changed files with 129 additions and 32 deletions

View File

@ -77,12 +77,13 @@ class ToolInvokeMessage(BaseModel):
LINK = "link" LINK = "link"
BLOB = "blob" BLOB = "blob"
IMAGE_LINK = "image_link" IMAGE_LINK = "image_link"
CHUNK = "chunk"
type: MessageType = MessageType.TEXT type: MessageType = MessageType.TEXT
""" """
plain text, image url or link url plain text, image url or link url
""" """
message: Union[str, bytes] = None message: Union[str, bytes, list] = None
meta: dict[str, Any] = None meta: dict[str, Any] = None
save_as: str = '' save_as: str = ''

View File

@ -40,7 +40,7 @@ class BingSearchTool(BuiltinTool):
news = response['news']['value'] if 'news' in response else [] news = response['news']['value'] if 'news' in response else []
computation = response['computation']['value'] if 'computation' in response else None computation = response['computation']['value'] if 'computation' in response else None
if result_type == 'link': if result_type == 'link' or result_type == 'chunk':
results = [] results = []
if search_results: if search_results:
for result in search_results: for result in search_results:
@ -72,7 +72,7 @@ class BingSearchTool(BuiltinTool):
)) ))
return results return results
else: if result_type == 'text' or result_type == 'chunk':
# construct text # construct text
text = '' text = ''
if search_results: if search_results:

View File

@ -6,6 +6,7 @@ from serpapi import GoogleSearch
from core.tools.entities.tool_entities import ToolInvokeMessage from core.tools.entities.tool_entities import ToolInvokeMessage
from core.tools.tool.builtin_tool import BuiltinTool from core.tools.tool.builtin_tool import BuiltinTool
from core.workflow.nodes.llm.knowledge_resource import KnowledgeResource
class HiddenPrints: class HiddenPrints:
@ -35,7 +36,7 @@ class SerpAPI:
self.serpapi_api_key = api_key self.serpapi_api_key = api_key
self.search_engine = GoogleSearch self.search_engine = GoogleSearch
def run(self, query: str, **kwargs: Any) -> str: def run(self, query: str, **kwargs: Any) -> str | list[KnowledgeResource]:
"""Run query through SerpAPI and parse result.""" """Run query through SerpAPI and parse result."""
typ = kwargs.get("result_type", "text") typ = kwargs.get("result_type", "text")
return self._process_response(self.results(query), typ=typ) return self._process_response(self.results(query), typ=typ)
@ -64,63 +65,79 @@ class SerpAPI:
return params return params
@staticmethod @staticmethod
def _process_response(res: dict, typ: str) -> str: def _process_response(res: dict, typ: str) -> str | list[KnowledgeResource]:
"""Process response from SerpAPI.""" """Process response from SerpAPI."""
if "error" in res.keys(): if "error" in res.keys():
raise ValueError(f"Got error from SerpAPI: {res['error']}") raise ValueError(f"Got error from SerpAPI: {res['error']}")
chunks = []
if typ == "text": toret = ""
toret = "" if typ == "text" or typ == "chunk":
if "answer_box" in res.keys() and type(res["answer_box"]) == list: if "answer_box" in res.keys() and type(res["answer_box"]) == list:
res["answer_box"] = res["answer_box"][0] + "\n" res["answer_box"] = res["answer_box"][0] + "\n"
if "answer_box" in res.keys() and "answer" in res["answer_box"].keys(): if "answer_box" in res.keys() and "answer" in res["answer_box"].keys():
toret += res["answer_box"]["answer"] + "\n" toret += res["answer_box"]["answer"] + "\n"
chunks.append(KnowledgeResource(content=res["answer_box"]["answer"], title=res["answer_box"]["answer"]))
if "answer_box" in res.keys() and "snippet" in res["answer_box"].keys(): if "answer_box" in res.keys() and "snippet" in res["answer_box"].keys():
toret += res["answer_box"]["snippet"] + "\n" toret += res["answer_box"]["snippet"] + "\n"
chunks.append(
KnowledgeResource(content=res["answer_box"]["snippet"], title=res["answer_box"]["snippet"]))
if ( if (
"answer_box" in res.keys() "answer_box" in res.keys()
and "snippet_highlighted_words" in res["answer_box"].keys() and "snippet_highlighted_words" in res["answer_box"].keys()
): ):
for item in res["answer_box"]["snippet_highlighted_words"]: for item in res["answer_box"]["snippet_highlighted_words"]:
toret += item + "\n" toret += item + "\n"
chunks.append(KnowledgeResource(content=item, title=item))
if ( if (
"sports_results" in res.keys() "sports_results" in res.keys()
and "game_spotlight" in res["sports_results"].keys() and "game_spotlight" in res["sports_results"].keys()
): ):
toret += res["sports_results"]["game_spotlight"] + "\n" toret += res["sports_results"]["game_spotlight"] + "\n"
chunks.append(KnowledgeResource(content=res["sports_results"]["game_spotlight"],
title=res["sports_results"]["game_spotlight"]))
if ( if (
"shopping_results" in res.keys() "shopping_results" in res.keys()
and "title" in res["shopping_results"][0].keys() and "title" in res["shopping_results"][0].keys()
): ):
toret += res["shopping_results"][:3] + "\n" toret += res["shopping_results"][:3] + "\n"
chunks.append(KnowledgeResource(content=res["shopping_results"][:3], title=res["shopping_results"][:3]))
if ( if (
"knowledge_graph" in res.keys() "knowledge_graph" in res.keys()
and "description" in res["knowledge_graph"].keys() and "description" in res["knowledge_graph"].keys()
): ):
toret = res["knowledge_graph"]["description"] + "\n" toret = res["knowledge_graph"]["description"] + "\n"
chunks.append(KnowledgeResource(content=res["knowledge_graph"]["description"],
title=res["knowledge_graph"]["description"]))
if "snippet" in res["organic_results"][0].keys(): if "snippet" in res["organic_results"][0].keys():
for item in res["organic_results"]: for item in res["organic_results"]:
toret += "content: " + item["snippet"] + "\n" + "link: " + item["link"] + "\n" toret += "content: " + item["snippet"] + "\n" + "link: " + item["link"] + "\n"
chunks.append(KnowledgeResource(content=item["snippet"], title=item["title"], url=item["link"]))
if ( if (
"images_results" in res.keys() "images_results" in res.keys()
and "thumbnail" in res["images_results"][0].keys() and "thumbnail" in res["images_results"][0].keys()
): ):
thumbnails = [item["thumbnail"] for item in res["images_results"][:10]] thumbnails = [item["thumbnail"] for item in res["images_results"][:10]]
toret = thumbnails toret = thumbnails
chunks.append(KnowledgeResource(content=thumbnails, title=thumbnails))
if toret == "": if toret == "":
toret = "No good search result found" toret = "No good search result found"
elif typ == "link": if typ == "link" or typ == "chunk":
if "knowledge_graph" in res.keys() and "title" in res["knowledge_graph"].keys() \ if "knowledge_graph" in res.keys() and "title" in res["knowledge_graph"].keys() \
and "description_link" in res["knowledge_graph"].keys(): and "description_link" in res["knowledge_graph"].keys():
toret = res["knowledge_graph"]["description_link"] toret = res["knowledge_graph"]["description_link"]
chunks.append(KnowledgeResource(content=res["knowledge_graph"]["description"],
title=res["knowledge_graph"]["title"],
url=res["knowledge_graph"]["knowledge_graph_search_link"])
)
elif "knowledge_graph" in res.keys() and "see_results_about" in res["knowledge_graph"].keys() \ elif "knowledge_graph" in res.keys() and "see_results_about" in res["knowledge_graph"].keys() \
and len(res["knowledge_graph"]["see_results_about"]) > 0: and len(res["knowledge_graph"]["see_results_about"]) > 0:
see_result_about = res["knowledge_graph"]["see_results_about"] see_result_about = res["knowledge_graph"]["see_results_about"]
toret = "" toret = ""
for item in see_result_about: for item in see_result_about:
if "name" not in item.keys() or "link" not in item.keys(): if "name" not in item.keys() or "link" not in item.keys():
continue continue
toret += f"[{item['name']}]({item['link']})\n" toret += f"[{item['name']}]({item['link']})\n"
chunks.append(KnowledgeResource(content=f"[{item['name']}]({item['link']})\n", title=item['name'], url=item['link']))
elif "organic_results" in res.keys() and len(res["organic_results"]) > 0: elif "organic_results" in res.keys() and len(res["organic_results"]) > 0:
organic_results = res["organic_results"] organic_results = res["organic_results"]
toret = "" toret = ""
@ -128,6 +145,7 @@ class SerpAPI:
if "title" not in item.keys() or "link" not in item.keys(): if "title" not in item.keys() or "link" not in item.keys():
continue continue
toret += f"[{item['title']}]({item['link']})\n" toret += f"[{item['title']}]({item['link']})\n"
chunks.append(KnowledgeResource(content=f"[{item['title']}]({item['link']})\n", title=item['title'], url=item['link']))
elif "related_questions" in res.keys() and len(res["related_questions"]) > 0: elif "related_questions" in res.keys() and len(res["related_questions"]) > 0:
related_questions = res["related_questions"] related_questions = res["related_questions"]
toret = "" toret = ""
@ -135,6 +153,7 @@ class SerpAPI:
if "question" not in item.keys() or "link" not in item.keys(): if "question" not in item.keys() or "link" not in item.keys():
continue continue
toret += f"[{item['question']}]({item['link']})\n" toret += f"[{item['question']}]({item['link']})\n"
chunks.append(KnowledgeResource(content=f"[{item['question']}]({item['link']})\n", title=item['title'], url=item['link']))
elif "related_searches" in res.keys() and len(res["related_searches"]) > 0: elif "related_searches" in res.keys() and len(res["related_searches"]) > 0:
related_searches = res["related_searches"] related_searches = res["related_searches"]
toret = "" toret = ""
@ -142,15 +161,19 @@ class SerpAPI:
if "query" not in item.keys() or "link" not in item.keys(): if "query" not in item.keys() or "link" not in item.keys():
continue continue
toret += f"[{item['query']}]({item['link']})\n" toret += f"[{item['query']}]({item['link']})\n"
chunks.append(KnowledgeResource(content=f"[{item['query']}]({item['link']})\n", title=item['query'], url=item['link']))
else: else:
toret = "No good search result found" toret = "No good search result found"
if typ == "chunk":
return chunks
return toret return toret
class GoogleSearchTool(BuiltinTool): class GoogleSearchTool(BuiltinTool):
def _invoke(self, def _invoke(self,
user_id: str, user_id: str,
tool_parameters: dict[str, Any], tool_parameters: dict[str, Any],
) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
""" """
invoke tools invoke tools
""" """
@ -160,5 +183,9 @@ class GoogleSearchTool(BuiltinTool):
result = SerpAPI(api_key).run(query, result_type=result_type) result = SerpAPI(api_key).run(query, result_type=result_type)
if result_type == 'text': if result_type == 'text':
return self.create_text_message(text=result) return self.create_text_message(text=result)
return self.create_link_message(link=result) elif result_type == 'link':
return self.create_link_message(link=result)
elif result_type == 'chunk':
return self.create_chunk_message(chunks=result)
else:
raise ValueError(f"Invalid result type: {result_type}")

View File

@ -39,6 +39,11 @@ parameters:
en_US: link en_US: link
zh_Hans: 链接 zh_Hans: 链接
pt_BR: link pt_BR: link
- value: chunk
label:
en_US: chunk
zh_Hans: 分段
pt_BR: chunk
default: link default: link
label: label:
en_US: Result type en_US: Result type

View File

@ -1,6 +1,6 @@
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from enum import Enum from enum import Enum
from typing import Any, Optional, Union from typing import Any, Optional, Union, List
from pydantic import BaseModel, validator from pydantic import BaseModel, validator
@ -15,6 +15,7 @@ from core.tools.entities.tool_entities import (
ToolRuntimeVariablePool, ToolRuntimeVariablePool,
) )
from core.tools.tool_file_manager import ToolFileManager from core.tools.tool_file_manager import ToolFileManager
from core.workflow.nodes.llm.knowledge_resource import KnowledgeResource
class Tool(BaseModel, ABC): class Tool(BaseModel, ABC):
@ -337,6 +338,8 @@ class Tool(BaseModel, ABC):
create an image message create an image message
:param image: the url of the image :param image: the url of the image
:param save_as: the save_as
:return: the image message :return: the image message
""" """
return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.IMAGE, return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.IMAGE,
@ -348,6 +351,7 @@ class Tool(BaseModel, ABC):
create a link message create a link message
:param link: the url of the link :param link: the url of the link
:param save_as: the save_as
:return: the link message :return: the link message
""" """
return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.LINK, return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.LINK,
@ -359,21 +363,37 @@ class Tool(BaseModel, ABC):
create a text message create a text message
:param text: the text :param text: the text
:param save_as: the save_as
:return: the text message :return: the text message
""" """
return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.TEXT, return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.TEXT,
message=text, message=text,
save_as=save_as save_as=save_as
) )
def create_chunk_message(self, chunks: List[KnowledgeResource], save_as: str = '') -> ToolInvokeMessage:
"""
create a chunk message
:param chunks: the chunks
:param save_as: the save_as
:return: the text message
"""
return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.CHUNK,
message=chunks,
save_as=save_as
)
def create_blob_message(self, blob: bytes, meta: dict = None, save_as: str = '') -> ToolInvokeMessage: def create_blob_message(self, blob: bytes, meta: dict = None, save_as: str = '') -> ToolInvokeMessage:
""" """
create a blob message create a blob message
:param blob: the blob :param blob: the blob
:param meta: the meta
:param save_as: the save_as
:return: the blob message :return: the blob message
""" """
return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.BLOB, return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.BLOB,
message=blob, meta=meta, message=blob, meta=meta,
save_as=save_as save_as=save_as
) )

View File

@ -131,7 +131,7 @@ class ToolEngine:
# hit the callback handler # hit the callback handler
workflow_tool_callback.on_tool_end( workflow_tool_callback.on_tool_end(
tool_name=tool.identity.name, tool_name=tool.identity.name,
tool_inputs=tool_parameters, tool_inputs=tool_parameters,
tool_outputs=response tool_outputs=response
) )

View File

@ -0,0 +1,16 @@
from typing import Any, Optional
from pydantic import BaseModel
class KnowledgeResource(BaseModel):
"""
Knowledge Resource.
"""
content: str
title: str
url: Optional[str] = None
icon: Optional[str] = None
score: Optional[float] = None
metadata: Optional[dict[str, Any]] = None

View File

@ -311,6 +311,9 @@ class LLMNode(BaseNode):
} }
return source return source
if ('metadata' in context_dict and '_source' in context_dict['metadata']
and context_dict['metadata']['_source'] == 'tool'):
return context_dict
return None return None

View File

@ -70,13 +70,14 @@ class ToolNode(BaseNode):
) )
# convert tool messages # convert tool messages
plain_text, files = self._convert_tool_messages(messages) plain_text, files, chunks = self._convert_tool_messages(messages)
return NodeRunResult( return NodeRunResult(
status=WorkflowNodeExecutionStatus.SUCCEEDED, status=WorkflowNodeExecutionStatus.SUCCEEDED,
outputs={ outputs={
'text': plain_text, 'text': plain_text,
'files': files 'files': files,
'chunks': chunks
}, },
metadata={ metadata={
NodeRunMetadataKey.TOOL_INFO: tool_info NodeRunMetadataKey.TOOL_INFO: tool_info
@ -111,7 +112,7 @@ class ToolNode(BaseNode):
return template_parser.format(inputs) return template_parser.format(inputs)
def _convert_tool_messages(self, messages: list[ToolInvokeMessage]) -> tuple[str, list[FileVar]]: def _convert_tool_messages(self, messages: list[ToolInvokeMessage]) -> tuple[str, list[FileVar], list]:
""" """
Convert ToolInvokeMessages into tuple[plain_text, files] Convert ToolInvokeMessages into tuple[plain_text, files]
""" """
@ -125,8 +126,9 @@ class ToolNode(BaseNode):
# extract plain text and files # extract plain text and files
files = self._extract_tool_response_binary(messages) files = self._extract_tool_response_binary(messages)
plain_text = self._extract_tool_response_text(messages) plain_text = self._extract_tool_response_text(messages)
chunks = self._extract_tool_response_chunk(messages)
return plain_text, files return plain_text, files, chunks
def _extract_tool_response_binary(self, tool_response: list[ToolInvokeMessage]) -> list[FileVar]: def _extract_tool_response_binary(self, tool_response: list[ToolInvokeMessage]) -> list[FileVar]:
""" """
@ -180,6 +182,29 @@ class ToolNode(BaseNode):
for message in tool_response for message in tool_response
]) ])
def _extract_tool_response_chunk(self, tool_response: list[ToolInvokeMessage]) -> list:
"""
Extract tool response text
"""
all_chunks = []
node_data = cast(ToolNodeData, self.node_data)
icon = ToolManager.get_tool_icon(
tenant_id=self.tenant_id,
provider_type=node_data.provider_type,
provider_id=node_data.provider_id
)
for message in tool_response:
if message.type == ToolInvokeMessage.MessageType.CHUNK:
for chunk in message.message:
chunk.icon = icon
chunk.metadata = {
'_source': 'tool'
}
all_chunks.append(chunk)
return all_chunks
@classmethod @classmethod
def _extract_variable_selector_to_variable_mapping(cls, node_data: ToolNodeData) -> dict[str, list[str]]: def _extract_variable_selector_to_variable_mapping(cls, node_data: ToolNodeData) -> dict[str, list[str]]:
""" """