add bing_search_chat.py and duckduckgo_search_chat.py

configs/model_config.py.example

@@ -107,3 +107,16 @@ PROMPT_TEMPLATE = """【指令】根据已知信息，简洁和专业的来回
 # API 是否开启跨域，默认为False，如果需要开启，请设置为True
 # is open cross domain
 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 = ""

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 starlette.responses import RedirectResponse
-from server.chat import chat, knowledge_base_chat, openai_chat
+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"],

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

server/chat/bing_chat.py

@@ -1,3 +0,0 @@
-# TODO: 完成 bing_chat agent 接口实现
-def bing_chat():
-    pass

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")

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")