1+ from fastapi import Body , Request
2+ from WebUI .configs import GetProviderByName
3+ from fastapi .responses import StreamingResponse
4+ from WebUI .Server .chat .utils import History
5+ from langchain .docstore .document import Document
6+ from langchain .chains import LLMChain
7+ import asyncio
8+ import json
9+ import os
10+ from WebUI .Server .utils import wrap_done , get_ChatOpenAI
11+ from fastapi .concurrency import run_in_threadpool
12+ from WebUI .Server .utils import get_prompt_template
13+ from langchain .callbacks import AsyncIteratorCallbackHandler
14+ from langchain .utilities .bing_search import BingSearchAPIWrapper
15+ from langchain .utilities .duckduckgo_search import DuckDuckGoSearchAPIWrapper
16+ from langchain .text_splitter import RecursiveCharacterTextSplitter
17+ from WebUI .configs .webuiconfig import InnerJsonConfigWebUIParse
18+ from langchain .prompts .chat import ChatPromptTemplate
19+ from typing import AsyncIterable , Dict , List , Optional
20+
21+ def bing_search (text , search_url , api_key , result_len , ** kwargs ):
22+ search = BingSearchAPIWrapper (bing_subscription_key = api_key ,
23+ bing_search_url = search_url )
24+ return search .results (text , result_len )
25+
26+
27+ def duckduckgo_search (text , search_url , api_key , result_len , ** kwargs ):
28+ search = DuckDuckGoSearchAPIWrapper ()
29+ return search .results (text , result_len )
30+
31+ def metaphor_search (
32+ text : str ,
33+ search_url : str ,
34+ api_key : str ,
35+ result_len : int ,
36+ chunk_size : int = 500 ,
37+ chunk_overlap : int = 50 ,
38+ ) -> List [Dict ]:
39+ from exa_py import Exa
40+ from markdownify import markdownify
41+ from strsimpy .normalized_levenshtein import NormalizedLevenshtein
42+
43+ highlights_options = {
44+ "num_sentences" : 7 , # how long our highlights should be
45+ "highlights_per_url" : 1 , # just get the best highlight for each URL
46+ }
47+
48+ info_for_llm = []
49+ exa = Exa (api_key = api_key )
50+ search_response = exa .search_and_contents (text , highlights = highlights_options , num_results = result_len , use_autoprompt = True )
51+ info = [sr for sr in search_response .results ]
52+ for x in info :
53+ x .highlights [0 ] = markdownify (x .highlights [0 ])
54+ info_for_llm = info
55+
56+ docs = [{"snippet" : x .highlights [0 ],
57+ "link" : x .url ,
58+ "title" : x .title }
59+ for x in info_for_llm ]
60+ return docs
61+
62+ SEARCH_ENGINES = {
63+ "bing" : bing_search ,
64+ "duckduckgo" : duckduckgo_search ,
65+ "metaphor" : metaphor_search ,
66+ }
67+
68+ def search_result2docs (search_results ):
69+ docs = []
70+ for result in search_results :
71+ doc = Document (page_content = result ["snippet" ] if "snippet" in result .keys () else "" ,
72+ metadata = {"source" : result ["link" ] if "link" in result .keys () else "" ,
73+ "filename" : result ["title" ] if "title" in result .keys () else "" })
74+ docs .append (doc )
75+ return docs
76+
77+
78+ async def lookup_search_engine (
79+ query : str ,
80+ search_engine_name : str ,
81+ top_k : int = 3 ,
82+ ):
83+ search_engine = SEARCH_ENGINES [search_engine_name ]
84+
85+ configinst = InnerJsonConfigWebUIParse ()
86+ webui_config = configinst .dump ()
87+ config = webui_config .get ("SearchEngine" ).get (search_engine_name )
88+ api_key = config .get ("api_key" , "" )
89+ search_url = config .get ("search_url" , "" )
90+ if search_engine_name == "bing" :
91+ if api_key == "" or api_key == "YOUR_API_KEY" :
92+ api_key = os .environ .get ("BING_SUBSCRIPTION_KEY" , "" )
93+ if search_url == "" :
94+ search_url = os .environ .get ("BING_SEARCH_URL" , "" )
95+ elif search_engine_name == "metaphor" :
96+ if api_key == "" or api_key == "YOUR_API_KEY" :
97+ api_key = os .environ .get ("METAPHOR_API_KEY" , "" )
98+ results = await run_in_threadpool (search_engine , query , search_url = search_url , api_key = api_key , result_len = top_k )
99+ docs = search_result2docs (results )
100+ return docs
101+
102+ async def search_engine_chat (
103+ query : str = Body (..., description = "User input: " , examples = ["chat" ]),
104+ search_engine_name : str = Body (..., description = "search engine name" , examples = ["duckduckgo" ]),
105+ history : List [dict ] = Body ([],
106+ description = "History chat" ,
107+ examples = [[
108+ {"role" : "user" , "content" : "Who are you?" },
109+ {"role" : "assistant" , "content" : "I am AI." }]]
110+ ),
111+ stream : bool = Body (False , description = "stream output" ),
112+ model_name : str = Body ("" , description = "model name" ),
113+ temperature : float = Body (0.7 , description = "LLM Temperature" , ge = 0.0 , le = 1.0 ),
114+ max_tokens : Optional [int ] = Body (None , description = "max tokens." ),
115+ prompt_name : str = Body ("default" , description = "" ),
116+ ) -> StreamingResponse :
117+ configinst = InnerJsonConfigWebUIParse ()
118+ webui_config = configinst .dump ()
119+ searchengine = webui_config .get ("SearchEngine" )
120+ top_k = searchengine .get ("top_k" , 3 )
121+
122+ history = [History .from_data (h ) for h in history ]
123+
124+ async def search_engine_chat_iterator (query : str ,
125+ search_engine_name : str ,
126+ top_k : int ,
127+ history : Optional [List [History ]],
128+ stream : bool ,
129+ model_name : str = "" ,
130+ temperature : float = 0.7 ,
131+ max_tokens : Optional [int ] = None ,
132+ prompt_name : str = prompt_name ,
133+ ) -> AsyncIterable [str ]:
134+ nonlocal webui_config
135+ callback = AsyncIteratorCallbackHandler ()
136+ if isinstance (max_tokens , int ) and max_tokens <= 0 :
137+ max_tokens = None
138+ provider = GetProviderByName (webui_config , model_name )
139+
140+ model = get_ChatOpenAI (
141+ provider = provider ,
142+ model_name = model_name ,
143+ temperature = temperature ,
144+ max_tokens = max_tokens ,
145+ callbacks = [callback ],
146+ )
147+
148+ docs = await lookup_search_engine (query , search_engine_name , top_k )
149+ context = "\n " .join ([doc .page_content for doc in docs ])
150+
151+ prompt_template = get_prompt_template ("search_engine_chat" , prompt_name )
152+ input_msg = History (role = "user" , content = prompt_template ).to_msg_template (False )
153+ chat_prompt = ChatPromptTemplate .from_messages (
154+ [i .to_msg_template () for i in history ] + [input_msg ])
155+
156+ chain = LLMChain (prompt = chat_prompt , llm = model )
157+
158+ task = asyncio .create_task (wrap_done (
159+ chain .acall ({"context" : context , "question" : query }),
160+ callback .done ),
161+ )
162+
163+ source_documents = [
164+ f"""from [{ inum + 1 } { doc .metadata ["source" ]} { doc .metadata ["source" ]} \n \n { doc .page_content } \n \n """
165+ for inum , doc in enumerate (docs )
166+ ]
167+
168+ if len (source_documents ) == 0 :
169+ source_documents .append (f"""<span style='color:red'>No relevant information were found. This response is generated based on the LLM Model '{ model_name } )
170+
171+ if stream :
172+ async for token in callback .aiter ():
173+ # Use server-sent-events to stream the response
174+ yield json .dumps ({"answer" : token }, ensure_ascii = False )
175+ yield json .dumps ({"docs" : source_documents }, ensure_ascii = False )
176+ else :
177+ answer = ""
178+ async for token in callback .aiter ():
179+ answer += token
180+ yield json .dumps ({"answer" : answer ,
181+ "docs" : source_documents },
182+ ensure_ascii = False )
183+ await task
184+
185+ return StreamingResponse (search_engine_chat_iterator (query = query ,
186+ search_engine_name = search_engine_name ,
187+ top_k = top_k ,
188+ history = history ,
189+ stream = stream ,
190+ model_name = model_name ,
191+ temperature = temperature ,
192+ max_tokens = max_tokens ,
193+ prompt_name = prompt_name ),
194+ )
0 commit comments