From 7c01a2a25313a8f2540ca753d887384f6dbf6ddb Mon Sep 17 00:00:00 2001 From: imClumsyPanda Date: Tue, 1 Aug 2023 16:39:17 +0800 Subject: [PATCH] add bing_search_chat.py and duckduckgo_search_chat.py --- configs/model_config.py.example | 15 +++++- server/api.py | 13 ++++-- server/chat/__init__.py | 2 + server/chat/bing_chat.py | 3 -- server/chat/bing_search_chat.py | 67 +++++++++++++++++++++++++++ server/chat/duckduckgo_search_chat.py | 62 +++++++++++++++++++++++++ 6 files changed, 155 insertions(+), 7 deletions(-) delete mode 100644 server/chat/bing_chat.py create mode 100644 server/chat/bing_search_chat.py create mode 100644 server/chat/duckduckgo_search_chat.py diff --git a/configs/model_config.py.example b/configs/model_config.py.example index eb4c4a9..7751bb4 100644 --- a/configs/model_config.py.example +++ b/configs/model_config.py.example @@ -106,4 +106,17 @@ PROMPT_TEMPLATE = """【指令】根据已知信息，简洁和专业的来回 # API 是否开启跨域，默认为False，如果需要开启，请设置为True # is open cross domain -OPEN_CROSS_DOMAIN = False \ No newline at end of file +OPEN_CROSS_DOMAIN = False + +# Bing 搜索必备变量 +# 使用 Bing 搜索需要使用 Bing Subscription Key,需要在azure port中申请试用bing search +# 具体申请方式请见 +# https://learn.microsoft.com/en-us/bing/search-apis/bing-web-search/create-bing-search-service-resource +# 使用python创建bing api 搜索实例详见: +# https://learn.microsoft.com/en-us/bing/search-apis/bing-web-search/quickstarts/rest/python +BING_SEARCH_URL = "https://api.bing.microsoft.com/v7.0/search" +# 注意不是bing Webmaster Tools的api key， + +# 此外，如果是在服务器上，报Failed to establish a new connection: [Errno 110] Connection timed out +# 是因为服务器加了防火墙，需要联系管理员加白名单，如果公司的服务器的话，就别想了GG +BING_SUBSCRIPTION_KEY = "" \ No newline at end of file diff --git a/server/api.py b/server/api.py index d58336c..9bbbbce 100644 --- a/server/api.py +++ b/server/api.py @@ -7,8 +7,9 @@ import argparse import uvicorn from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware -from starlette.responses import RedirectResponse, StreamingResponse -from server.chat import chat, knowledge_base_chat, openai_chat +from starlette.responses import RedirectResponse +from server.chat import (chat, knowledge_base_chat, openai_chat, + bing_search_chat, duckduckgo_search_chat) from server.knowledge_base import (list_kbs, create_kb, delete_kb, list_docs, upload_doc, delete_doc, update_doc) from server.utils import BaseResponse, ListResponse @@ -50,7 +51,13 @@ def create_app(): tags=["Chat"], summary="与知识库对话")(knowledge_base_chat) - # app.post("/chat/bing_search_chat", tags=["Chat"], summary="与Bing搜索对话")(bing_search_chat) + app.post("/chat/bing_search_chat", + tags=["Chat"], + summary="与Bing搜索对话")(bing_search_chat) + + app.post("/chat/duckduckgo_search_chat", + tags=["Chat"], + summary="与DuckDuckGo搜索对话")(duckduckgo_search_chat) app.get("/knowledge_base/list_knowledge_bases", tags=["Knowledge Base Management"], diff --git a/server/chat/__init__.py b/server/chat/__init__.py index 9fef958..728db91 100644 --- a/server/chat/__init__.py +++ b/server/chat/__init__.py @@ -1,3 +1,5 @@ from .chat import chat from .knowledge_base_chat import knowledge_base_chat from .openai_chat import openai_chat +from .duckduckgo_search_chat import duckduckgo_search_chat +from .bing_search_chat import bing_search_chat diff --git a/server/chat/bing_chat.py b/server/chat/bing_chat.py deleted file mode 100644 index 71453d2..0000000 --- a/server/chat/bing_chat.py +++ /dev/null @@ -1,3 +0,0 @@ -# TODO: 完成 bing_chat agent 接口实现 -def bing_chat(): - pass \ No newline at end of file diff --git a/server/chat/bing_search_chat.py b/server/chat/bing_search_chat.py new file mode 100644 index 0000000..bbb9089 --- /dev/null +++ b/server/chat/bing_search_chat.py @@ -0,0 +1,67 @@ +from langchain.utilities import BingSearchAPIWrapper +from configs.model_config import BING_SEARCH_URL, BING_SUBSCRIPTION_KEY +from fastapi import Body +from fastapi.responses import StreamingResponse +from configs.model_config import (llm_model_dict, LLM_MODEL, PROMPT_TEMPLATE) +from server.chat.utils import wrap_done +from langchain.chat_models import ChatOpenAI +from langchain import LLMChain +from langchain.callbacks import AsyncIteratorCallbackHandler +from typing import AsyncIterable +import asyncio +from langchain.prompts import PromptTemplate +from langchain.docstore.document import Document + +def bing_search(text, result_len=3): + if not (BING_SEARCH_URL and BING_SUBSCRIPTION_KEY): + return [{"snippet": "please set BING_SUBSCRIPTION_KEY and BING_SEARCH_URL in os ENV", + "title": "env info is not found", + "link": "https://python.langchain.com/en/latest/modules/agents/tools/examples/bing_search.html"}] + search = BingSearchAPIWrapper(bing_subscription_key=BING_SUBSCRIPTION_KEY, + bing_search_url=BING_SEARCH_URL) + return search.results(text, result_len) + + +def search_result2docs(search_results): + docs = [] + for result in search_results: + doc = Document(page_content=result["snippet"] if "snippet" in result.keys() else "", + metadata={"source": result["link"] if "link" in result.keys() else "", + "filename": result["title"] if "title" in result.keys() else ""}) + docs.append(doc) + return docs + + +def bing_search_chat(query: str = Body(..., description="用户输入", example="你好"), + ): + async def bing_search_chat_iterator(query: str, + ) -> AsyncIterable[str]: + callback = AsyncIteratorCallbackHandler() + model = ChatOpenAI( + streaming=True, + verbose=True, + callbacks=[callback], + openai_api_key=llm_model_dict[LLM_MODEL]["api_key"], + openai_api_base=llm_model_dict[LLM_MODEL]["api_base_url"], + model_name=LLM_MODEL + ) + + results = bing_search(query, result_len=3) + docs = search_result2docs(results) + context = "\n".join([doc.page_content for doc in docs]) + prompt = PromptTemplate(template=PROMPT_TEMPLATE, input_variables=["context", "question"]) + + chain = LLMChain(prompt=prompt, llm=model) + + # Begin a task that runs in the background. + task = asyncio.create_task(wrap_done( + chain.acall({"context": context, "question": query}), + callback.done), + ) + + async for token in callback.aiter(): + # Use server-sent-events to stream the response + yield token + await task + + return StreamingResponse(bing_search_chat_iterator(query), media_type="text/event-stream") diff --git a/server/chat/duckduckgo_search_chat.py b/server/chat/duckduckgo_search_chat.py new file mode 100644 index 0000000..c10e065 --- /dev/null +++ b/server/chat/duckduckgo_search_chat.py @@ -0,0 +1,62 @@ +from langchain.utilities import DuckDuckGoSearchAPIWrapper +from fastapi import Body +from fastapi.responses import StreamingResponse +from configs.model_config import (llm_model_dict, LLM_MODEL, PROMPT_TEMPLATE) +from server.chat.utils import wrap_done +from langchain.chat_models import ChatOpenAI +from langchain import LLMChain +from langchain.callbacks import AsyncIteratorCallbackHandler +from typing import AsyncIterable +import asyncio +from langchain.prompts import PromptTemplate +from langchain.docstore.document import Document + + +def duckduckgo_search(text, result_len=3): + search = DuckDuckGoSearchAPIWrapper() + return search.results(text, result_len) + + +def search_result2docs(search_results): + docs = [] + for result in search_results: + doc = Document(page_content=result["snippet"] if "snippet" in result.keys() else "", + metadata={"source": result["link"] if "link" in result.keys() else "", + "filename": result["title"] if "title" in result.keys() else ""}) + docs.append(doc) + return docs + + +def duckduckgo_search_chat(query: str = Body(..., description="用户输入", example="你好"), + ): + async def duckduckgo_search_chat_iterator(query: str, + ) -> AsyncIterable[str]: + callback = AsyncIteratorCallbackHandler() + model = ChatOpenAI( + streaming=True, + verbose=True, + callbacks=[callback], + openai_api_key=llm_model_dict[LLM_MODEL]["api_key"], + openai_api_base=llm_model_dict[LLM_MODEL]["api_base_url"], + model_name=LLM_MODEL + ) + + results = duckduckgo_search(query, result_len=3) + docs = search_result2docs(results) + context = "\n".join([doc.page_content for doc in docs]) + prompt = PromptTemplate(template=PROMPT_TEMPLATE, input_variables=["context", "question"]) + + chain = LLMChain(prompt=prompt, llm=model) + + # Begin a task that runs in the background. + task = asyncio.create_task(wrap_done( + chain.acall({"context": context, "question": query}), + callback.done), + ) + + async for token in callback.aiter(): + # Use server-sent-events to stream the response + yield token + await task + + return StreamingResponse(duckduckgo_search_chat_iterator(query), media_type="text/event-stream")