Commit fdd75ff3 authored by John Wang's avatar John Wang

Merge branch 'feat/universal-chat' into deploy/dev

parents 7b3806a7 0b8fd9df
...@@ -5,6 +5,7 @@ from langchain import WikipediaAPIWrapper ...@@ -5,6 +5,7 @@ from langchain import WikipediaAPIWrapper
from langchain.callbacks.manager import Callbacks from langchain.callbacks.manager import Callbacks
from langchain.memory.chat_memory import BaseChatMemory from langchain.memory.chat_memory import BaseChatMemory
from langchain.tools import BaseTool, Tool, WikipediaQueryRun from langchain.tools import BaseTool, Tool, WikipediaQueryRun
from pydantic import BaseModel, Field
from core.agent.agent_executor import AgentExecutor, PlanningStrategy, AgentConfiguration from core.agent.agent_executor import AgentExecutor, PlanningStrategy, AgentConfiguration
from core.callback_handler.agent_loop_gather_callback_handler import AgentLoopGatherCallbackHandler from core.callback_handler.agent_loop_gather_callback_handler import AgentLoopGatherCallbackHandler
...@@ -17,7 +18,7 @@ from core.llm.llm_builder import LLMBuilder ...@@ -17,7 +18,7 @@ from core.llm.llm_builder import LLMBuilder
from core.llm.streamable_chat_open_ai import StreamableChatOpenAI from core.llm.streamable_chat_open_ai import StreamableChatOpenAI
from core.tool.dataset_retriever_tool import DatasetRetrieverTool from core.tool.dataset_retriever_tool import DatasetRetrieverTool
from core.tool.provider.serpapi_provider import SerpAPIToolProvider from core.tool.provider.serpapi_provider import SerpAPIToolProvider
from core.tool.serpapi_wrapper import OptimizedSerpAPIWrapper from core.tool.serpapi_wrapper import OptimizedSerpAPIWrapper, OptimizedSerpAPIInput
from core.tool.web_reader_tool import WebReaderTool from core.tool.web_reader_tool import WebReaderTool
from extensions.ext_database import db from extensions.ext_database import db
from models.dataset import Dataset, DatasetProcessRule from models.dataset import Dataset, DatasetProcessRule
...@@ -224,15 +225,20 @@ class OrchestratorRuleParser: ...@@ -224,15 +225,20 @@ class OrchestratorRuleParser:
"is not up to date." "is not up to date."
"Input should be a search query.", "Input should be a search query.",
func=OptimizedSerpAPIWrapper(**func_kwargs).run, func=OptimizedSerpAPIWrapper(**func_kwargs).run,
args_schema=OptimizedSerpAPIInput,
callbacks=[DifyStdOutCallbackHandler()] callbacks=[DifyStdOutCallbackHandler()]
) )
return tool return tool
def to_wikipedia_tool(self) -> Optional[BaseTool]: def to_wikipedia_tool(self) -> Optional[BaseTool]:
class WikipediaInput(BaseModel):
query: str = Field(..., description="search query.")
return WikipediaQueryRun( return WikipediaQueryRun(
name="wikipedia", name="wikipedia",
api_wrapper=WikipediaAPIWrapper(doc_content_chars_max=4000), api_wrapper=WikipediaAPIWrapper(doc_content_chars_max=4000),
args_schema=WikipediaInput,
callbacks=[DifyStdOutCallbackHandler()] callbacks=[DifyStdOutCallbackHandler()]
) )
......
from langchain import SerpAPIWrapper from langchain import SerpAPIWrapper
from pydantic import Field, BaseModel
class OptimizedSerpAPIInput(BaseModel):
query: str = Field(..., description="search query.")
class OptimizedSerpAPIWrapper(SerpAPIWrapper): class OptimizedSerpAPIWrapper(SerpAPIWrapper):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment