mirror of
https://github.com/wahyd4/SAI2.git
synced 2026-08-09 05:06:05 +10:00
修改主程序文件，跟lepton解耦
This commit is contained in:
fatwang2
@@ -1,52 +0,0 @@
|
||||
# Lepton Search
|
||||
Build your own conversational search engine using less than 500 lines of code.
|
||||
|
||||
See a live demo site https://search.lepton.run/
|
||||
|
||||
The source code of this project lives [here](https://github.com/leptonai/search_with_lepton/). This README will detail how to set up and deploy this project on Lepton's platform.
|
||||
|
||||
## Setup Search Engine API
|
||||
|
||||
You have a few options for setting up your search engine API. You can use Bing or Google, or if you just want to very quickly try the demo out, use the lepton demo API directly.
|
||||
|
||||
### Bing
|
||||
|
||||
If you are using Bing, you can subscribe to the bing search api [here](https://www.microsoft.com/en-us/bing/apis/bing-web-search-api). After that, write down the Bing search api subscription key. We follow the convention and name it `BING_SEARCH_V7_SUBSCRIPTION_KEY`. We recommend you store the key as a secret in Lepton.
|
||||
|
||||
### Google
|
||||
|
||||
If you choose to use Google, you can follow the instructions [here](https://developers.google.com/custom-search/v1/overview) to get your Google search api key. We follow the convention and name it `GOOGLE_SEARCH_API_KEY`. We recommend you store the key as a secret in Lepton. You will also get a search engine CX id, which you will need as well.
|
||||
|
||||
### SearchApi
|
||||
|
||||
If you want to use SearchApi, a 3rd party Google Search API, you can retrieve the API key by registering [here](https://www.searchapi.io/). We follow the convention and name it `SEARCHAPI_API_KEY`. We recommend you store the key as a secret in Lepton.
|
||||
|
||||
### Lepton Demo API
|
||||
|
||||
If you choose to use the lepton demo api, you don't need to do anything - your workspace credential will give you access to the demo api. Note that this does incur an API call cost.
|
||||
|
||||
|
||||
## Deployment Configurations
|
||||
|
||||
Here are the configurations you can set for your deployment:
|
||||
* Name: The name of your deployment, like "my-search"
|
||||
* Resource Shape: most of heavy lifting will be done by the LLM server and the search engine API, so you can choose a small resource shape. `cpu.small` is usually good enough.
|
||||
|
||||
Then, set the following environmental variables.
|
||||
|
||||
* `BACKEND`: the search backend to use. If you don't have bing or google set up, simply use `LEPTON` to try the demo. Otherwise, do `BING`, `GOOGLE` or `SEARCHAPI`.
|
||||
* `LLM_MODEL`: the LLM model to run. We recommend using `mixtral-8x7b`, but if you want to experiment other models, you can try the ones hosted on LeptonAI, for example, `llama2-70b`, `llama2-13b`, `llama2-7b`. Note that small models won't work that well.
|
||||
* `KV_NAME`: the Lepton KV to use to store the search results. You can use the default `search-with-lepton`.
|
||||
* `RELATED_QUESTIONS`: whether to generate related questions. If you set this to `true`, the search engine will generate related questions for you. Otherwise, it will not.
|
||||
* `GOOGLE_SEARCH_CX`: if you are using google, specify the search cx. Otherwise, leave it empty.
|
||||
* `LEPTON_ENABLE_AUTH_BY_COOKIE`: this is to allow web UI access to the deployment. Set it to `true`.
|
||||
|
||||
In addition, you will need to set the following secrets:
|
||||
* `LEPTON_WORKSPACE_TOKEN`: this is required to call Lepton's LLM and KV apis. You can find your workspace token at [Settings](https://dashboard.lepton.ai/workspace-redirect/settings).
|
||||
* `BING_SEARCH_V7_SUBSCRIPTION_KEY`: if you are using Bing, you need to specify the subscription key. Otherwise it is not needed.
|
||||
* `GOOGLE_SEARCH_API_KEY`: if you are using Google, you need to specify the search api key. Note that you should also specify the cx in the env. If you are not using Google, it is not needed.
|
||||
* `SEARCHAPI_API_KEY`: if you are using SearchApi, a 3rd party Google Search API, you need to specify the api key.
|
||||
|
||||
Once these fields are set, click `Deploy` button at the bottom of the page to create the deployment. You can see the deployment has now been created under [Deployments](https://dashboard.lepton.ai/workspace-redirect/deployments). Click on the deployment name to check the details. You’ll be able to see the deployment URL and status on this page.
|
||||
|
||||
Once the status is turned into `Ready`, click the URL on the deployment card to access it. Enjoy!
|
||||
+618
@@ -0,0 +1,618 @@
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import requests
|
||||
import traceback
|
||||
import httpx
|
||||
from typing import Annotated, List, Generator
|
||||
from openai import AsyncOpenAI
|
||||
from loguru import logger
|
||||
|
||||
import sanic
|
||||
from sanic import Sanic
|
||||
import sanic.exceptions
|
||||
from sanic.exceptions import HTTPException, InvalidUsage
|
||||
from sqlitedict import SqliteDict
|
||||
|
||||
app = Sanic("search")
|
||||
|
||||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
|
||||
################################################################################
|
||||
# Constant values for the RAG model.
|
||||
################################################################################
|
||||
|
||||
# Search engine related. You don't really need to change this.
|
||||
BING_SEARCH_V7_ENDPOINT = "https://api.bing.microsoft.com/v7.0/search"
|
||||
BING_MKT = "en-US"
|
||||
GOOGLE_SEARCH_ENDPOINT = "https://customsearch.googleapis.com/customsearch/v1"
|
||||
SERPER_SEARCH_ENDPOINT = "https://google.serper.dev/search"
|
||||
SEARCHAPI_SEARCH_ENDPOINT = "https://www.searchapi.io/api/v1/search"
|
||||
|
||||
|
||||
# Specify the number of references from the search engine you want to use.
|
||||
# 8 is usually a good number.
|
||||
REFERENCE_COUNT = 8
|
||||
|
||||
# Specify the default timeout for the search engine. If the search engine
|
||||
# does not respond within this time, we will return an error.
|
||||
DEFAULT_SEARCH_ENGINE_TIMEOUT = 5
|
||||
|
||||
|
||||
# If the user did not provide a query, we will use this default query.
|
||||
_default_query = "Who said 'live long and prosper'?"
|
||||
|
||||
# This is really the most important part of the rag model. It gives instructions
|
||||
# to the model on how to generate the answer. Of course, different models may
|
||||
# behave differently, and we haven't tuned the prompt to make it optimal - this
|
||||
# is left to you, application creators, as an open problem.
|
||||
_rag_query_text = """
|
||||
You are a large language AI assistant built by Lepton AI. You are given a user question, and please write clean, concise and accurate answer to the question. You will be given a set of related contexts to the question, each starting with a reference number like [[citation:x]], where x is a number. Please use the context and cite the context at the end of each sentence if applicable.
|
||||
|
||||
Your answer must be correct, accurate and written by an expert using an unbiased and professional tone. Please limit to 1024 tokens. Do not give any information that is not related to the question, and do not repeat. Say "information is missing on" followed by the related topic, if the given context do not provide sufficient information.
|
||||
|
||||
Please cite the contexts with the reference numbers, in the format [citation:x]. If a sentence comes from multiple contexts, please list all applicable citations, like [citation:3][citation:5]. Other than code and specific names and citations, your answer must be written in the same language as the question.
|
||||
|
||||
Here are the set of contexts:
|
||||
|
||||
{context}
|
||||
|
||||
Remember, don't blindly repeat the contexts verbatim. And here is the user question:
|
||||
"""
|
||||
|
||||
# A set of stop words to use - this is not a complete set, and you may want to
|
||||
# add more given your observation.
|
||||
stop_words = [
|
||||
"<|im_end|>",
|
||||
"[End]",
|
||||
"[end]",
|
||||
"\nReferences:\n",
|
||||
"\nSources:\n",
|
||||
"End.",
|
||||
]
|
||||
|
||||
# This is the prompt that asks the model to generate related questions to the
|
||||
# original question and the contexts.
|
||||
# Ideally, one want to include both the original question and the answer from the
|
||||
# model, but we are not doing that here: if we need to wait for the answer, then
|
||||
# the generation of the related questions will usually have to start only after
|
||||
# the whole answer is generated. This creates a noticeable delay in the response
|
||||
# time. As a result, and as you will see in the code, we will be sending out two
|
||||
# consecutive requests to the model: one for the answer, and one for the related
|
||||
# questions. This is not ideal, but it is a good tradeoff between response time
|
||||
# and quality.
|
||||
_more_questions_prompt = """
|
||||
You are a helpful assistant that helps the user to ask related questions, based on user's original question and the related contexts. Please identify worthwhile topics that can be follow-ups, and write questions no longer than 20 words each. Please make sure that specifics, like events, names, locations, are included in follow up questions so they can be asked standalone. For example, if the original question asks about "the Manhattan project", in the follow up question, do not just say "the project", but use the full name "the Manhattan project". Your related questions must be in the same language as the original question.
|
||||
|
||||
Here are the contexts of the question:
|
||||
|
||||
{context}
|
||||
|
||||
Remember, based on the original question and related contexts, suggest three such further questions. Do NOT repeat the original question. Each related question should be no longer than 20 words. Here is the original question:
|
||||
"""
|
||||
|
||||
|
||||
class KVWrapper(object):
|
||||
def __init__(self, kv_name):
|
||||
self._db = SqliteDict(filename=kv_name)
|
||||
|
||||
def get(self, key: str):
|
||||
v = self._db[key]
|
||||
if v is None:
|
||||
raise KeyError(key)
|
||||
return v
|
||||
|
||||
def put(self, key: str, value: str):
|
||||
self._db[key] = value
|
||||
self._db.commit()
|
||||
|
||||
|
||||
def search_with_bing(query: str, subscription_key: str):
|
||||
"""
|
||||
Search with bing and return the contexts.
|
||||
"""
|
||||
params = {"q": query, "mkt": BING_MKT}
|
||||
response = requests.get(
|
||||
BING_SEARCH_V7_ENDPOINT,
|
||||
headers={"Ocp-Apim-Subscription-Key": subscription_key},
|
||||
params=params,
|
||||
timeout=DEFAULT_SEARCH_ENGINE_TIMEOUT,
|
||||
)
|
||||
if not response.ok:
|
||||
logger.error(f"{response.status_code} {response.text}")
|
||||
raise HTTPException("Search engine error.")
|
||||
json_content = response.json()
|
||||
try:
|
||||
contexts = json_content["webPages"]["value"][:REFERENCE_COUNT]
|
||||
except KeyError:
|
||||
logger.error(f"Error encountered: {json_content}")
|
||||
return []
|
||||
return contexts
|
||||
|
||||
|
||||
def search_with_google(query: str, subscription_key: str, cx: str):
|
||||
"""
|
||||
Search with google and return the contexts.
|
||||
"""
|
||||
params = {
|
||||
"key": subscription_key,
|
||||
"cx": cx,
|
||||
"q": query,
|
||||
"num": REFERENCE_COUNT,
|
||||
}
|
||||
response = requests.get(
|
||||
GOOGLE_SEARCH_ENDPOINT, params=params, timeout=DEFAULT_SEARCH_ENGINE_TIMEOUT
|
||||
)
|
||||
if not response.ok:
|
||||
logger.error(f"{response.status_code} {response.text}")
|
||||
raise HTTPException("Search engine error.")
|
||||
json_content = response.json()
|
||||
try:
|
||||
contexts = json_content["items"][:REFERENCE_COUNT]
|
||||
for item in contexts:
|
||||
item["name"] = item["title"]
|
||||
item["url"] = item["link"]
|
||||
except KeyError:
|
||||
logger.error(f"Error encountered: {json_content}")
|
||||
return []
|
||||
return contexts
|
||||
|
||||
|
||||
def search_with_serper(query: str, subscription_key: str):
|
||||
"""
|
||||
Search with serper and return the contexts.
|
||||
"""
|
||||
payload = json.dumps(
|
||||
{
|
||||
"q": query,
|
||||
"num": (
|
||||
REFERENCE_COUNT
|
||||
if REFERENCE_COUNT % 10 == 0
|
||||
else (REFERENCE_COUNT // 10 + 1) * 10
|
||||
),
|
||||
}
|
||||
)
|
||||
headers = {"X-API-KEY": subscription_key, "Content-Type": "application/json"}
|
||||
logger.info(
|
||||
f"{payload} {headers} {subscription_key} {query} {SERPER_SEARCH_ENDPOINT}"
|
||||
)
|
||||
response = requests.post(
|
||||
SERPER_SEARCH_ENDPOINT,
|
||||
headers=headers,
|
||||
data=payload,
|
||||
timeout=DEFAULT_SEARCH_ENGINE_TIMEOUT,
|
||||
)
|
||||
if not response.ok:
|
||||
logger.error(f"{response.status_code} {response.text}")
|
||||
raise HTTPException("Search engine error.")
|
||||
json_content = response.json()
|
||||
try:
|
||||
# convert to the same format as bing/google
|
||||
contexts = []
|
||||
if json_content.get("knowledgeGraph"):
|
||||
url = json_content["knowledgeGraph"].get("descriptionUrl") or json_content[
|
||||
"knowledgeGraph"
|
||||
].get("website")
|
||||
snippet = json_content["knowledgeGraph"].get("description")
|
||||
if url and snippet:
|
||||
contexts.append(
|
||||
{
|
||||
"name": json_content["knowledgeGraph"].get("title", ""),
|
||||
"url": url,
|
||||
"snippet": snippet,
|
||||
}
|
||||
)
|
||||
if json_content.get("answerBox"):
|
||||
url = json_content["answerBox"].get("url")
|
||||
snippet = json_content["answerBox"].get("snippet") or json_content[
|
||||
"answerBox"
|
||||
].get("answer")
|
||||
if url and snippet:
|
||||
contexts.append(
|
||||
{
|
||||
"name": json_content["answerBox"].get("title", ""),
|
||||
"url": url,
|
||||
"snippet": snippet,
|
||||
}
|
||||
)
|
||||
contexts += [
|
||||
{"name": c["title"], "url": c["link"], "snippet": c.get("snippet", "")}
|
||||
for c in json_content["organic"]
|
||||
]
|
||||
return contexts[:REFERENCE_COUNT]
|
||||
except KeyError:
|
||||
logger.error(f"Error encountered: {json_content}")
|
||||
return []
|
||||
|
||||
|
||||
def search_with_searchapi(query: str, subscription_key: str):
|
||||
"""
|
||||
Search with SearchApi.io and return the contexts.
|
||||
"""
|
||||
payload = {
|
||||
"q": query,
|
||||
"engine": "google",
|
||||
"num": (
|
||||
REFERENCE_COUNT
|
||||
if REFERENCE_COUNT % 10 == 0
|
||||
else (REFERENCE_COUNT // 10 + 1) * 10
|
||||
),
|
||||
}
|
||||
headers = {
|
||||
"Authorization": f"Bearer {subscription_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
logger.info(
|
||||
f"{payload} {headers} {subscription_key} {query} {SEARCHAPI_SEARCH_ENDPOINT}"
|
||||
)
|
||||
response = requests.get(
|
||||
SEARCHAPI_SEARCH_ENDPOINT,
|
||||
headers=headers,
|
||||
params=payload,
|
||||
timeout=30,
|
||||
)
|
||||
if not response.ok:
|
||||
logger.error(f"{response.status_code} {response.text}")
|
||||
raise HTTPException("Search engine error.")
|
||||
json_content = response.json()
|
||||
try:
|
||||
# convert to the same format as bing/google
|
||||
contexts = []
|
||||
|
||||
if json_content.get("answer_box"):
|
||||
if json_content["answer_box"].get("organic_result"):
|
||||
title = (
|
||||
json_content["answer_box"].get("organic_result").get("title", "")
|
||||
)
|
||||
url = json_content["answer_box"].get("organic_result").get("link", "")
|
||||
if json_content["answer_box"].get("type") == "population_graph":
|
||||
title = json_content["answer_box"].get("place", "")
|
||||
url = json_content["answer_box"].get("explore_more_link", "")
|
||||
|
||||
title = json_content["answer_box"].get("title", "")
|
||||
url = json_content["answer_box"].get("link")
|
||||
snippet = json_content["answer_box"].get("answer") or json_content[
|
||||
"answer_box"
|
||||
].get("snippet")
|
||||
|
||||
if url and snippet:
|
||||
contexts.append({"name": title, "url": url, "snippet": snippet})
|
||||
|
||||
if json_content.get("knowledge_graph"):
|
||||
if json_content["knowledge_graph"].get("source"):
|
||||
url = json_content["knowledge_graph"].get("source").get("link", "")
|
||||
|
||||
url = json_content["knowledge_graph"].get("website", "")
|
||||
snippet = json_content["knowledge_graph"].get("description")
|
||||
|
||||
if url and snippet:
|
||||
contexts.append(
|
||||
{
|
||||
"name": json_content["knowledge_graph"].get("title", ""),
|
||||
"url": url,
|
||||
"snippet": snippet,
|
||||
}
|
||||
)
|
||||
|
||||
contexts += [
|
||||
{"name": c["title"], "url": c["link"], "snippet": c.get("snippet", "")}
|
||||
for c in json_content["organic_results"]
|
||||
]
|
||||
|
||||
if json_content.get("related_questions"):
|
||||
for question in json_content["related_questions"]:
|
||||
if question.get("source"):
|
||||
url = question.get("source").get("link", "")
|
||||
else:
|
||||
url = ""
|
||||
|
||||
snippet = question.get("answer", "")
|
||||
|
||||
if url and snippet:
|
||||
contexts.append(
|
||||
{
|
||||
"name": question.get("question", ""),
|
||||
"url": url,
|
||||
"snippet": snippet,
|
||||
}
|
||||
)
|
||||
|
||||
return contexts[:REFERENCE_COUNT]
|
||||
except KeyError:
|
||||
logger.error(f"Error encountered: {json_content}")
|
||||
return []
|
||||
|
||||
|
||||
def new_async_client(_app):
|
||||
return AsyncOpenAI(
|
||||
api_key=os.environ["OPENAI_API_KEY"],
|
||||
base_url=os.environ["OPENAI_BASE_URL"],
|
||||
http_client=_app.ctx.http_session,
|
||||
)
|
||||
|
||||
|
||||
@app.before_server_start
|
||||
async def server_init(_app, loop):
|
||||
"""
|
||||
Initializes global configs.
|
||||
"""
|
||||
_app.ctx.backend = os.environ["BACKEND"].upper()
|
||||
# if _app.ctx.backend == "LEPTON":
|
||||
# from leptonai import Client
|
||||
|
||||
# _app.ctx.leptonsearch_client = Client(
|
||||
# "https://search-api.lepton.run/",
|
||||
# token=os.environ.get("LEPTON_WORKSPACE_TOKEN"),
|
||||
# stream=True,
|
||||
# timeout=httpx.Timeout(connect=10, read=120, write=120, pool=10),
|
||||
# )
|
||||
if _app.ctx.backend == "BING":
|
||||
_app.ctx.search_api_key = os.environ["BING_SEARCH_V7_SUBSCRIPTION_KEY"]
|
||||
_app.ctx.search_function = lambda query: search_with_bing(
|
||||
query,
|
||||
_app.ctx.search_api_key,
|
||||
)
|
||||
elif _app.ctx.backend == "GOOGLE":
|
||||
_app.ctx.search_api_key = os.environ["GOOGLE_SEARCH_API_KEY"]
|
||||
_app.ctx.search_function = lambda query: search_with_google(
|
||||
query,
|
||||
_app.ctx.search_api_key,
|
||||
os.environ["GOOGLE_SEARCH_CX"],
|
||||
)
|
||||
elif _app.ctx.backend == "SERPER":
|
||||
_app.ctx.search_api_key = os.environ["SERPER_SEARCH_API_KEY"]
|
||||
_app.ctx.search_function = lambda query: search_with_serper(
|
||||
query,
|
||||
_app.ctx.search_api_key,
|
||||
)
|
||||
elif _app.ctx.backend == "SEARCHAPI":
|
||||
_app.ctx.search_api_key = os.environ["SEARCHAPI_API_KEY"]
|
||||
_app.ctx.search_function = lambda query: search_with_searchapi(
|
||||
query,
|
||||
_app.ctx.search_api_key,
|
||||
)
|
||||
else:
|
||||
raise RuntimeError("Backend must be BING, GOOGLE, SERPER or SEARCHAPI.")
|
||||
_app.ctx.model = os.environ["LLM_MODEL"]
|
||||
_app.ctx.handler_max_concurrency = 16
|
||||
# An executor to carry out async tasks, such as uploading to KV.
|
||||
_app.ctx.executor = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=_app.ctx.handler_max_concurrency * 2
|
||||
)
|
||||
# Create the KV to store the search results.
|
||||
logger.info("Creating KV. May take a while for the first time.")
|
||||
_app.ctx.kv = KVWrapper(os.getenv("KV_NAME") or "search.db")
|
||||
# whether we should generate related questions.
|
||||
_app.ctx.should_do_related_questions = bool(
|
||||
os.environ["RELATED_QUESTIONS"] in ("1", "yes", "true")
|
||||
)
|
||||
# Create httpx Session
|
||||
_app.ctx.http_session = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(connect=10, read=120, write=120, pool=10),
|
||||
)
|
||||
|
||||
|
||||
async def get_related_questions(_app, query, contexts):
|
||||
"""
|
||||
Gets related questions based on the query and context.
|
||||
"""
|
||||
|
||||
def ask_related_questions(
|
||||
questions: Annotated[
|
||||
List[str],
|
||||
[
|
||||
(
|
||||
"question",
|
||||
Annotated[
|
||||
str,
|
||||
"related question to the original question and context.",
|
||||
],
|
||||
)
|
||||
],
|
||||
],
|
||||
):
|
||||
"""
|
||||
ask further questions that are related to the input and output.
|
||||
"""
|
||||
pass
|
||||
|
||||
try:
|
||||
openai_client = new_async_client(_app)
|
||||
response = await openai_client.chat.completions.create(
|
||||
model=_app.ctx.model,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": _more_questions_prompt.format(
|
||||
context="\n\n".join([c["snippet"] for c in contexts])
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": query,
|
||||
},
|
||||
],
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
# "function": tool.get_tools_spec(ask_related_questions),
|
||||
}
|
||||
],
|
||||
max_tokens=512,
|
||||
)
|
||||
related = response.choices[0].message.tool_calls[0].function.arguments
|
||||
if isinstance(related, str):
|
||||
related = json.loads(related)
|
||||
logger.trace(f"Related questions: {related}")
|
||||
return related["questions"][:5]
|
||||
except Exception as e:
|
||||
# For any exceptions, we will just return an empty list.
|
||||
logger.error(
|
||||
"encountered error while generating related questions:"
|
||||
f" {e}\n{traceback.format_exc()}"
|
||||
)
|
||||
return []
|
||||
|
||||
|
||||
async def _raw_stream_response(
|
||||
_app, contexts, llm_response, related_questions_future
|
||||
) -> Generator[str, None, None]:
|
||||
"""
|
||||
A generator that yields the raw stream response. You do not need to call
|
||||
this directly. Instead, use the stream_and_upload_to_kv which will also
|
||||
upload the response to KV.
|
||||
"""
|
||||
# First, yield the contexts.
|
||||
yield json.dumps(contexts)
|
||||
yield "\n\n__LLM_RESPONSE__\n\n"
|
||||
# Second, yield the llm response.
|
||||
if not contexts:
|
||||
# Prepend a warning to the user
|
||||
yield (
|
||||
"(The search engine returned nothing for this query. Please take the"
|
||||
" answer with a grain of salt.)\n\n"
|
||||
)
|
||||
async for chunk in llm_response:
|
||||
if chunk.choices:
|
||||
yield chunk.choices[0].delta.content or ""
|
||||
# Third, yield the related questions. If any error happens, we will just
|
||||
# return an empty list.
|
||||
if related_questions_future is not None:
|
||||
related_questions = await related_questions_future
|
||||
try:
|
||||
result = json.dumps(related_questions)
|
||||
except Exception as e:
|
||||
logger.error(f"encountered error: {e}\n{traceback.format_exc()}")
|
||||
result = "[]"
|
||||
yield "\n\n__RELATED_QUESTIONS__\n\n"
|
||||
yield result
|
||||
|
||||
|
||||
def get_query_object(request):
|
||||
params = {k: v[0] for k, v in request.args.items()}
|
||||
if request.method == "POST":
|
||||
if "form" in request.content_type:
|
||||
params.update({k: v[0] for k, v in request.form.items()})
|
||||
else:
|
||||
try:
|
||||
if request.json:
|
||||
params.update(request.json)
|
||||
except InvalidUsage:
|
||||
pass
|
||||
return params
|
||||
|
||||
|
||||
@app.route("/query", methods=["POST"])
|
||||
async def query_function(request: sanic.Request):
|
||||
"""
|
||||
Query the search engine and returns the response.
|
||||
|
||||
The query can have the following fields:
|
||||
- query: the user query.
|
||||
- search_uuid: a uuid that is used to store or retrieve the search result. If
|
||||
the uuid does not exist, generate and write to the kv. If the kv
|
||||
fails, we generate regardless, in favor of availability. If the uuid
|
||||
exists, return the stored result.
|
||||
- generate_related_questions: if set to false, will not generate related
|
||||
questions. Otherwise, will depend on the environment variable
|
||||
RELATED_QUESTIONS. Default: true.
|
||||
"""
|
||||
_app = request.app
|
||||
params = get_query_object(request)
|
||||
query = params.get("query", None)
|
||||
search_uuid = params.get("search_uuid", None)
|
||||
generate_related_questions = params.get("generate_related_questions", True)
|
||||
if not query:
|
||||
raise HTTPException("query must be provided.")
|
||||
# Note that, if uuid exists, we don't check if the stored query is the same
|
||||
# as the current query, and simply return the stored result. This is to enable
|
||||
# the user to share a searched link to others and have others see the same result.
|
||||
if search_uuid:
|
||||
try:
|
||||
result = await _app.loop.run_in_executor(
|
||||
_app.ctx.executor, lambda sid: _app.ctx.kv.get(sid), search_uuid
|
||||
)
|
||||
return sanic.text(result)
|
||||
except KeyError:
|
||||
logger.info(f"Key {search_uuid} not found, will generate again.")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"KV error: {e}\n{traceback.format_exc()}, will generate again."
|
||||
)
|
||||
else:
|
||||
raise HTTPException("search_uuid must be provided.")
|
||||
|
||||
# if _app.ctx.backend == "LEPTON":
|
||||
# # delegate to the lepton search api.
|
||||
# result = _app.ctx.leptonsearch_client.query(
|
||||
# query=query,
|
||||
# search_uuid=search_uuid,
|
||||
# generate_related_questions=generate_related_questions,
|
||||
# )
|
||||
# return StreamingResponse(content=result, media_type="text/html")
|
||||
|
||||
# First, do a search query.
|
||||
# query = query or _default_query
|
||||
# Basic attack protection: remove "[INST]" or "[/INST]" from the query
|
||||
query = re.sub(r"\[/?INST\]", "", query)
|
||||
contexts = await _app.loop.run_in_executor(
|
||||
_app.ctx.executor, _app.ctx.search_function, query
|
||||
)
|
||||
|
||||
system_prompt = _rag_query_text.format(
|
||||
context="\n\n".join(
|
||||
[f"[[citation:{i+1}]] {c['snippet']}" for i, c in enumerate(contexts)]
|
||||
)
|
||||
)
|
||||
try:
|
||||
openai_client = new_async_client(_app)
|
||||
llm_response = await openai_client.chat.completions.create(
|
||||
model=_app.ctx.model,
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": query},
|
||||
],
|
||||
max_tokens=1024,
|
||||
stop=stop_words,
|
||||
stream=True,
|
||||
temperature=0.9,
|
||||
)
|
||||
if _app.ctx.should_do_related_questions and generate_related_questions:
|
||||
# While the answer is being generated, we can start generating
|
||||
# related questions as a future.
|
||||
related_questions_future = get_related_questions(_app, query, contexts)
|
||||
|
||||
else:
|
||||
related_questions_future = None
|
||||
except Exception as e:
|
||||
logger.error(f"encountered error: {e}\n{traceback.format_exc()}")
|
||||
return sanic.json({"message": "Internal server error."}, 503)
|
||||
|
||||
response = await request.respond(content_type="text/html")
|
||||
# First, stream and yield the results.
|
||||
all_yielded_results = []
|
||||
|
||||
async for result in _raw_stream_response(
|
||||
_app, contexts, llm_response, related_questions_future
|
||||
):
|
||||
all_yielded_results.append(result)
|
||||
await response.send(result)
|
||||
|
||||
# Second, upload to KV. Note that if uploading to KV fails, we will silently
|
||||
# ignore it, because we don't want to affect the user experience.
|
||||
await response.eof()
|
||||
_ = _app.ctx.executor.submit(
|
||||
_app.ctx.kv.put, search_uuid, "".join(all_yielded_results)
|
||||
)
|
||||
|
||||
|
||||
app.static("/ui", os.path.join(BASE_DIR, "ui/"), name="/")
|
||||
app.static("/", os.path.join(BASE_DIR, "ui/index.html"), name="ui")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
port = int(os.getenv("PORT") or 8800)
|
||||
workers = int(os.getenv("WORKERS") or 1)
|
||||
app.run(host="0.0.0.0", port=port, workers=workers, debug=False)
|
||||
@@ -1,644 +0,0 @@
|
||||
import concurrent.futures
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import requests
|
||||
import traceback
|
||||
from typing import Annotated, List, Generator, Optional
|
||||
|
||||
from fastapi import HTTPException
|
||||
from fastapi.responses import HTMLResponse, StreamingResponse, RedirectResponse
|
||||
import httpx
|
||||
from loguru import logger
|
||||
|
||||
import leptonai
|
||||
from leptonai import Client
|
||||
from leptonai.kv import KV
|
||||
from leptonai.photon import Photon, StaticFiles
|
||||
from leptonai.photon.types import to_bool
|
||||
from leptonai.api.workspace import WorkspaceInfoLocalRecord
|
||||
from leptonai.util import tool
|
||||
|
||||
################################################################################
|
||||
# Constant values for the RAG model.
|
||||
################################################################################
|
||||
|
||||
# Search engine related. You don't really need to change this.
|
||||
BING_SEARCH_V7_ENDPOINT = "https://api.bing.microsoft.com/v7.0/search"
|
||||
BING_MKT = "en-US"
|
||||
GOOGLE_SEARCH_ENDPOINT = "https://customsearch.googleapis.com/customsearch/v1"
|
||||
SERPER_SEARCH_ENDPOINT = "https://google.serper.dev/search"
|
||||
SEARCHAPI_SEARCH_ENDPOINT = "https://www.searchapi.io/api/v1/search"
|
||||
|
||||
# Specify the number of references from the search engine you want to use.
|
||||
# 8 is usually a good number.
|
||||
REFERENCE_COUNT = 8
|
||||
|
||||
# Specify the default timeout for the search engine. If the search engine
|
||||
# does not respond within this time, we will return an error.
|
||||
DEFAULT_SEARCH_ENGINE_TIMEOUT = 5
|
||||
|
||||
|
||||
# If the user did not provide a query, we will use this default query.
|
||||
_default_query = "Who said 'live long and prosper'?"
|
||||
|
||||
# This is really the most important part of the rag model. It gives instructions
|
||||
# to the model on how to generate the answer. Of course, different models may
|
||||
# behave differently, and we haven't tuned the prompt to make it optimal - this
|
||||
# is left to you, application creators, as an open problem.
|
||||
_rag_query_text = """
|
||||
You are a large language AI assistant built by Lepton AI. You are given a user question, and please write clean, concise and accurate answer to the question. You will be given a set of related contexts to the question, each starting with a reference number like [[citation:x]], where x is a number. Please use the context and cite the context at the end of each sentence if applicable.
|
||||
|
||||
Your answer must be correct, accurate and written by an expert using an unbiased and professional tone. Please limit to 1024 tokens. Do not give any information that is not related to the question, and do not repeat. Say "information is missing on" followed by the related topic, if the given context do not provide sufficient information.
|
||||
|
||||
Please cite the contexts with the reference numbers, in the format [citation:x]. If a sentence comes from multiple contexts, please list all applicable citations, like [citation:3][citation:5]. Other than code and specific names and citations, your answer must be written in the same language as the question.
|
||||
|
||||
Here are the set of contexts:
|
||||
|
||||
{context}
|
||||
|
||||
Remember, don't blindly repeat the contexts verbatim. And here is the user question:
|
||||
"""
|
||||
|
||||
# A set of stop words to use - this is not a complete set, and you may want to
|
||||
# add more given your observation.
|
||||
stop_words = [
|
||||
"<|im_end|>",
|
||||
"[End]",
|
||||
"[end]",
|
||||
"\nReferences:\n",
|
||||
"\nSources:\n",
|
||||
"End.",
|
||||
]
|
||||
|
||||
# This is the prompt that asks the model to generate related questions to the
|
||||
# original question and the contexts.
|
||||
# Ideally, one want to include both the original question and the answer from the
|
||||
# model, but we are not doing that here: if we need to wait for the answer, then
|
||||
# the generation of the related questions will usually have to start only after
|
||||
# the whole answer is generated. This creates a noticeable delay in the response
|
||||
# time. As a result, and as you will see in the code, we will be sending out two
|
||||
# consecutive requests to the model: one for the answer, and one for the related
|
||||
# questions. This is not ideal, but it is a good tradeoff between response time
|
||||
# and quality.
|
||||
_more_questions_prompt = """
|
||||
You are a helpful assistant that helps the user to ask related questions, based on user's original question and the related contexts. Please identify worthwhile topics that can be follow-ups, and write questions no longer than 20 words each. Please make sure that specifics, like events, names, locations, are included in follow up questions so they can be asked standalone. For example, if the original question asks about "the Manhattan project", in the follow up question, do not just say "the project", but use the full name "the Manhattan project". Your related questions must be in the same language as the original question.
|
||||
|
||||
Here are the contexts of the question:
|
||||
|
||||
{context}
|
||||
|
||||
Remember, based on the original question and related contexts, suggest three such further questions. Do NOT repeat the original question. Each related question should be no longer than 20 words. Here is the original question:
|
||||
"""
|
||||
|
||||
|
||||
def search_with_bing(query: str, subscription_key: str):
|
||||
"""
|
||||
Search with bing and return the contexts.
|
||||
"""
|
||||
params = {"q": query, "mkt": BING_MKT}
|
||||
response = requests.get(
|
||||
BING_SEARCH_V7_ENDPOINT,
|
||||
headers={"Ocp-Apim-Subscription-Key": subscription_key},
|
||||
params=params,
|
||||
timeout=DEFAULT_SEARCH_ENGINE_TIMEOUT,
|
||||
)
|
||||
if not response.ok:
|
||||
logger.error(f"{response.status_code} {response.text}")
|
||||
raise HTTPException(response.status_code, "Search engine error.")
|
||||
json_content = response.json()
|
||||
try:
|
||||
contexts = json_content["webPages"]["value"][:REFERENCE_COUNT]
|
||||
except KeyError:
|
||||
logger.error(f"Error encountered: {json_content}")
|
||||
return []
|
||||
return contexts
|
||||
|
||||
|
||||
def search_with_google(query: str, subscription_key: str, cx: str):
|
||||
"""
|
||||
Search with google and return the contexts.
|
||||
"""
|
||||
params = {
|
||||
"key": subscription_key,
|
||||
"cx": cx,
|
||||
"q": query,
|
||||
"num": REFERENCE_COUNT,
|
||||
}
|
||||
response = requests.get(
|
||||
GOOGLE_SEARCH_ENDPOINT, params=params, timeout=DEFAULT_SEARCH_ENGINE_TIMEOUT
|
||||
)
|
||||
if not response.ok:
|
||||
logger.error(f"{response.status_code} {response.text}")
|
||||
raise HTTPException(response.status_code, "Search engine error.")
|
||||
json_content = response.json()
|
||||
try:
|
||||
contexts = json_content["items"][:REFERENCE_COUNT]
|
||||
except KeyError:
|
||||
logger.error(f"Error encountered: {json_content}")
|
||||
return []
|
||||
return contexts
|
||||
|
||||
|
||||
def search_with_serper(query: str, subscription_key: str):
|
||||
"""
|
||||
Search with serper and return the contexts.
|
||||
"""
|
||||
payload = json.dumps({
|
||||
"q": query,
|
||||
"num": (
|
||||
REFERENCE_COUNT
|
||||
if REFERENCE_COUNT % 10 == 0
|
||||
else (REFERENCE_COUNT // 10 + 1) * 10
|
||||
),
|
||||
})
|
||||
headers = {"X-API-KEY": subscription_key, "Content-Type": "application/json"}
|
||||
logger.info(
|
||||
f"{payload} {headers} {subscription_key} {query} {SERPER_SEARCH_ENDPOINT}"
|
||||
)
|
||||
response = requests.post(
|
||||
SERPER_SEARCH_ENDPOINT,
|
||||
headers=headers,
|
||||
data=payload,
|
||||
timeout=DEFAULT_SEARCH_ENGINE_TIMEOUT,
|
||||
)
|
||||
if not response.ok:
|
||||
logger.error(f"{response.status_code} {response.text}")
|
||||
raise HTTPException(response.status_code, "Search engine error.")
|
||||
json_content = response.json()
|
||||
try:
|
||||
# convert to the same format as bing/google
|
||||
contexts = []
|
||||
if json_content.get("knowledgeGraph"):
|
||||
url = json_content["knowledgeGraph"].get("descriptionUrl") or json_content["knowledgeGraph"].get("website")
|
||||
snippet = json_content["knowledgeGraph"].get("description")
|
||||
if url and snippet:
|
||||
contexts.append({
|
||||
"name": json_content["knowledgeGraph"].get("title",""),
|
||||
"url": url,
|
||||
"snippet": snippet
|
||||
})
|
||||
if json_content.get("answerBox"):
|
||||
url = json_content["answerBox"].get("url")
|
||||
snippet = json_content["answerBox"].get("snippet") or json_content["answerBox"].get("answer")
|
||||
if url and snippet:
|
||||
contexts.append({
|
||||
"name": json_content["answerBox"].get("title",""),
|
||||
"url": url,
|
||||
"snippet": snippet
|
||||
})
|
||||
contexts += [
|
||||
{"name": c["title"], "url": c["link"], "snippet": c.get("snippet","")}
|
||||
for c in json_content["organic"]
|
||||
]
|
||||
return contexts[:REFERENCE_COUNT]
|
||||
except KeyError:
|
||||
logger.error(f"Error encountered: {json_content}")
|
||||
return []
|
||||
|
||||
def search_with_searchapi(query: str, subscription_key: str):
|
||||
"""
|
||||
Search with SearchApi.io and return the contexts.
|
||||
"""
|
||||
payload = {
|
||||
"q": query,
|
||||
"engine": "google",
|
||||
"num": (
|
||||
REFERENCE_COUNT
|
||||
if REFERENCE_COUNT % 10 == 0
|
||||
else (REFERENCE_COUNT // 10 + 1) * 10
|
||||
),
|
||||
}
|
||||
headers = {"Authorization": f"Bearer {subscription_key}", "Content-Type": "application/json"}
|
||||
logger.info(
|
||||
f"{payload} {headers} {subscription_key} {query} {SEARCHAPI_SEARCH_ENDPOINT}"
|
||||
)
|
||||
response = requests.get(
|
||||
SEARCHAPI_SEARCH_ENDPOINT,
|
||||
headers=headers,
|
||||
params=payload,
|
||||
timeout=30,
|
||||
)
|
||||
if not response.ok:
|
||||
logger.error(f"{response.status_code} {response.text}")
|
||||
raise HTTPException(response.status_code, "Search engine error.")
|
||||
json_content = response.json()
|
||||
try:
|
||||
# convert to the same format as bing/google
|
||||
contexts = []
|
||||
|
||||
if json_content.get("answer_box"):
|
||||
if json_content["answer_box"].get("organic_result"):
|
||||
title = json_content["answer_box"].get("organic_result").get("title", "")
|
||||
url = json_content["answer_box"].get("organic_result").get("link", "")
|
||||
if json_content["answer_box"].get("type") == "population_graph":
|
||||
title = json_content["answer_box"].get("place", "")
|
||||
url = json_content["answer_box"].get("explore_more_link", "")
|
||||
|
||||
title = json_content["answer_box"].get("title", "")
|
||||
url = json_content["answer_box"].get("link")
|
||||
snippet = json_content["answer_box"].get("answer") or json_content["answer_box"].get("snippet")
|
||||
|
||||
if url and snippet:
|
||||
contexts.append({
|
||||
"name": title,
|
||||
"url": url,
|
||||
"snippet": snippet
|
||||
})
|
||||
|
||||
if json_content.get("knowledge_graph"):
|
||||
if json_content["knowledge_graph"].get("source"):
|
||||
url = json_content["knowledge_graph"].get("source").get("link", "")
|
||||
|
||||
url = json_content["knowledge_graph"].get("website", "")
|
||||
snippet = json_content["knowledge_graph"].get("description")
|
||||
|
||||
if url and snippet:
|
||||
contexts.append({
|
||||
"name": json_content["knowledge_graph"].get("title", ""),
|
||||
"url": url,
|
||||
"snippet": snippet
|
||||
})
|
||||
|
||||
contexts += [
|
||||
{"name": c["title"], "url": c["link"], "snippet": c.get("snippet", "")}
|
||||
for c in json_content["organic_results"]
|
||||
]
|
||||
|
||||
if json_content.get("related_questions"):
|
||||
for question in json_content["related_questions"]:
|
||||
if question.get("source"):
|
||||
url = question.get("source").get("link", "")
|
||||
else:
|
||||
url = ""
|
||||
|
||||
snippet = question.get("answer", "")
|
||||
|
||||
if url and snippet:
|
||||
contexts.append({
|
||||
"name": question.get("question", ""),
|
||||
"url": url,
|
||||
"snippet": snippet
|
||||
})
|
||||
|
||||
return contexts[:REFERENCE_COUNT]
|
||||
except KeyError:
|
||||
logger.error(f"Error encountered: {json_content}")
|
||||
return []
|
||||
|
||||
class RAG(Photon):
|
||||
"""
|
||||
Retrieval-Augmented Generation Demo from Lepton AI.
|
||||
|
||||
This is a minimal example to show how to build a RAG engine with Lepton AI.
|
||||
It uses search engine to obtain results based on user queries, and then uses
|
||||
LLM models to generate the answer as well as related questions. The results
|
||||
are then stored in a KV so that it can be retrieved later.
|
||||
"""
|
||||
|
||||
requirement_dependency = [
|
||||
"openai", # for openai client usage.
|
||||
]
|
||||
|
||||
extra_files = glob.glob("ui/**/*", recursive=True)
|
||||
|
||||
deployment_template = {
|
||||
# All actual computations are carried out via remote apis, so
|
||||
# we will use a cpu.small instance which is already enough for most of
|
||||
# the work.
|
||||
"resource_shape": "cpu.small",
|
||||
# You most likely don't need to change this.
|
||||
"env": {
|
||||
# Choose the backend. Currently, we support BING and GOOGLE. For
|
||||
# simplicity, in this demo, if you specify the backend as LEPTON,
|
||||
# we will use the hosted serverless version of lepton search api
|
||||
# at https://search-api.lepton.run/ to do the search and RAG, which
|
||||
# runs the same code (slightly modified and might contain improvements)
|
||||
# as this demo.
|
||||
"BACKEND": "LEPTON",
|
||||
# If you are using google, specify the search cx.
|
||||
"GOOGLE_SEARCH_CX": "",
|
||||
# Specify the LLM model you are going to use.
|
||||
"LLM_MODEL": "mixtral-8x7b",
|
||||
# For all the search queries and results, we will use the Lepton KV to
|
||||
# store them so that we can retrieve them later. Specify the name of the
|
||||
# KV here.
|
||||
"KV_NAME": "search-with-lepton",
|
||||
# If set to true, will generate related questions. Otherwise, will not.
|
||||
"RELATED_QUESTIONS": "true",
|
||||
# On the lepton platform, allow web access when you are logged in.
|
||||
"LEPTON_ENABLE_AUTH_BY_COOKIE": "true",
|
||||
},
|
||||
# Secrets you need to have: search api subscription key, and lepton
|
||||
# workspace token to query lepton's llama models.
|
||||
"secret": [
|
||||
# If you use BING, you need to specify the subscription key. Otherwise
|
||||
# it is not needed.
|
||||
"BING_SEARCH_V7_SUBSCRIPTION_KEY",
|
||||
# If you use GOOGLE, you need to specify the search api key. Note that
|
||||
# you should also specify the cx in the env.
|
||||
"GOOGLE_SEARCH_API_KEY",
|
||||
# If you use Serper, you need to specify the search api key.
|
||||
"SERPER_SEARCH_API_KEY",
|
||||
# If you use SearchApi, you need to specify the search api key.
|
||||
"SEARCHAPI_API_KEY",
|
||||
# You need to specify the workspace token to query lepton's LLM models.
|
||||
"LEPTON_WORKSPACE_TOKEN",
|
||||
],
|
||||
}
|
||||
|
||||
# It's just a bunch of api calls, so our own deployment can be made massively
|
||||
# concurrent.
|
||||
handler_max_concurrency = 16
|
||||
|
||||
def local_client(self):
|
||||
"""
|
||||
Gets a thread-local client, so in case openai clients are not thread safe,
|
||||
each thread will have its own client.
|
||||
"""
|
||||
import openai
|
||||
|
||||
thread_local = threading.local()
|
||||
try:
|
||||
return thread_local.client
|
||||
except AttributeError:
|
||||
thread_local.client = openai.OpenAI(
|
||||
base_url=f"https://{self.model}.lepton.run/api/v1/",
|
||||
api_key=os.environ.get("LEPTON_WORKSPACE_TOKEN")
|
||||
or WorkspaceInfoLocalRecord.get_current_workspace_token(),
|
||||
# We will set the connect timeout to be 10 seconds, and read/write
|
||||
# timeout to be 120 seconds, in case the inference server is
|
||||
# overloaded.
|
||||
timeout=httpx.Timeout(connect=10, read=120, write=120, pool=10),
|
||||
)
|
||||
return thread_local.client
|
||||
|
||||
def init(self):
|
||||
"""
|
||||
Initializes photon configs.
|
||||
"""
|
||||
# First, log in to the workspace.
|
||||
leptonai.api.workspace.login()
|
||||
self.backend = os.environ["BACKEND"].upper()
|
||||
if self.backend == "LEPTON":
|
||||
self.leptonsearch_client = Client(
|
||||
"https://search-api.lepton.run/",
|
||||
token=os.environ.get("LEPTON_WORKSPACE_TOKEN")
|
||||
or WorkspaceInfoLocalRecord.get_current_workspace_token(),
|
||||
stream=True,
|
||||
timeout=httpx.Timeout(connect=10, read=120, write=120, pool=10),
|
||||
)
|
||||
elif self.backend == "BING":
|
||||
self.search_api_key = os.environ["BING_SEARCH_V7_SUBSCRIPTION_KEY"]
|
||||
self.search_function = lambda query: search_with_bing(
|
||||
query,
|
||||
self.search_api_key,
|
||||
)
|
||||
elif self.backend == "GOOGLE":
|
||||
self.search_api_key = os.environ["GOOGLE_SEARCH_API_KEY"]
|
||||
self.search_function = lambda query: search_with_google(
|
||||
query,
|
||||
self.search_api_key,
|
||||
os.environ["GOOGLE_SEARCH_CX"],
|
||||
)
|
||||
elif self.backend == "SERPER":
|
||||
self.search_api_key = os.environ["SERPER_SEARCH_API_KEY"]
|
||||
self.search_function = lambda query: search_with_serper(
|
||||
query,
|
||||
self.search_api_key,
|
||||
)
|
||||
elif self.backend == "SEARCHAPI":
|
||||
self.search_api_key = os.environ["SEARCHAPI_API_KEY"]
|
||||
self.search_function = lambda query: search_with_searchapi(
|
||||
query,
|
||||
self.search_api_key,
|
||||
)
|
||||
else:
|
||||
raise RuntimeError("Backend must be LEPTON, BING, GOOGLE, SERPER or SEARCHAPI.")
|
||||
self.model = os.environ["LLM_MODEL"]
|
||||
# An executor to carry out async tasks, such as uploading to KV.
|
||||
self.executor = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=self.handler_max_concurrency * 2
|
||||
)
|
||||
# Create the KV to store the search results.
|
||||
logger.info("Creating KV. May take a while for the first time.")
|
||||
self.kv = KV(
|
||||
os.environ["KV_NAME"], create_if_not_exists=True, error_if_exists=False
|
||||
)
|
||||
# whether we should generate related questions.
|
||||
self.should_do_related_questions = to_bool(os.environ["RELATED_QUESTIONS"])
|
||||
|
||||
def get_related_questions(self, query, contexts):
|
||||
"""
|
||||
Gets related questions based on the query and context.
|
||||
"""
|
||||
|
||||
def ask_related_questions(
|
||||
questions: Annotated[
|
||||
List[str],
|
||||
[(
|
||||
"question",
|
||||
Annotated[
|
||||
str, "related question to the original question and context."
|
||||
],
|
||||
)],
|
||||
]
|
||||
):
|
||||
"""
|
||||
ask further questions that are related to the input and output.
|
||||
"""
|
||||
pass
|
||||
|
||||
try:
|
||||
response = self.local_client().chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": _more_questions_prompt.format(
|
||||
context="\n\n".join([c["snippet"] for c in contexts])
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": query,
|
||||
},
|
||||
],
|
||||
tools=[{
|
||||
"type": "function",
|
||||
"function": tool.get_tools_spec(ask_related_questions),
|
||||
}],
|
||||
max_tokens=512,
|
||||
)
|
||||
related = response.choices[0].message.tool_calls[0].function.arguments
|
||||
if isinstance(related, str):
|
||||
related = json.loads(related)
|
||||
logger.trace(f"Related questions: {related}")
|
||||
return related["questions"][:5]
|
||||
except Exception as e:
|
||||
# For any exceptions, we will just return an empty list.
|
||||
logger.error(
|
||||
"encountered error while generating related questions:"
|
||||
f" {e}\n{traceback.format_exc()}"
|
||||
)
|
||||
return []
|
||||
|
||||
def _raw_stream_response(
|
||||
self, contexts, llm_response, related_questions_future
|
||||
) -> Generator[str, None, None]:
|
||||
"""
|
||||
A generator that yields the raw stream response. You do not need to call
|
||||
this directly. Instead, use the stream_and_upload_to_kv which will also
|
||||
upload the response to KV.
|
||||
"""
|
||||
# First, yield the contexts.
|
||||
yield json.dumps(contexts)
|
||||
yield "\n\n__LLM_RESPONSE__\n\n"
|
||||
# Second, yield the llm response.
|
||||
if not contexts:
|
||||
# Prepend a warning to the user
|
||||
yield (
|
||||
"(The search engine returned nothing for this query. Please take the"
|
||||
" answer with a grain of salt.)\n\n"
|
||||
)
|
||||
for chunk in llm_response:
|
||||
if chunk.choices:
|
||||
yield chunk.choices[0].delta.content or ""
|
||||
# Third, yield the related questions. If any error happens, we will just
|
||||
# return an empty list.
|
||||
if related_questions_future is not None:
|
||||
related_questions = related_questions_future.result()
|
||||
try:
|
||||
result = json.dumps(related_questions)
|
||||
except Exception as e:
|
||||
logger.error(f"encountered error: {e}\n{traceback.format_exc()}")
|
||||
result = "[]"
|
||||
yield "\n\n__RELATED_QUESTIONS__\n\n"
|
||||
yield result
|
||||
|
||||
def stream_and_upload_to_kv(
|
||||
self, contexts, llm_response, related_questions_future, search_uuid
|
||||
) -> Generator[str, None, None]:
|
||||
"""
|
||||
Streams the result and uploads to KV.
|
||||
"""
|
||||
# First, stream and yield the results.
|
||||
all_yielded_results = []
|
||||
for result in self._raw_stream_response(
|
||||
contexts, llm_response, related_questions_future
|
||||
):
|
||||
all_yielded_results.append(result)
|
||||
yield result
|
||||
# Second, upload to KV. Note that if uploading to KV fails, we will silently
|
||||
# ignore it, because we don't want to affect the user experience.
|
||||
_ = self.executor.submit(self.kv.put, search_uuid, "".join(all_yielded_results))
|
||||
|
||||
@Photon.handler(method="POST", path="/query")
|
||||
def query_function(
|
||||
self,
|
||||
query: str,
|
||||
search_uuid: str,
|
||||
generate_related_questions: Optional[bool] = True,
|
||||
) -> StreamingResponse:
|
||||
"""
|
||||
Query the search engine and returns the response.
|
||||
|
||||
The query can have the following fields:
|
||||
- query: the user query.
|
||||
- search_uuid: a uuid that is used to store or retrieve the search result. If
|
||||
the uuid does not exist, generate and write to the kv. If the kv
|
||||
fails, we generate regardless, in favor of availability. If the uuid
|
||||
exists, return the stored result.
|
||||
- generate_related_questions: if set to false, will not generate related
|
||||
questions. Otherwise, will depend on the environment variable
|
||||
RELATED_QUESTIONS. Default: true.
|
||||
"""
|
||||
# Note that, if uuid exists, we don't check if the stored query is the same
|
||||
# as the current query, and simply return the stored result. This is to enable
|
||||
# the user to share a searched link to others and have others see the same result.
|
||||
if search_uuid:
|
||||
try:
|
||||
result = self.kv.get(search_uuid)
|
||||
|
||||
def str_to_generator(result: str) -> Generator[str, None, None]:
|
||||
yield result
|
||||
|
||||
return StreamingResponse(str_to_generator(result))
|
||||
except KeyError:
|
||||
logger.info(f"Key {search_uuid} not found, will generate again.")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"KV error: {e}\n{traceback.format_exc()}, will generate again."
|
||||
)
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail="search_uuid must be provided.")
|
||||
|
||||
if self.backend == "LEPTON":
|
||||
# delegate to the lepton search api.
|
||||
result = self.leptonsearch_client.query(
|
||||
query=query,
|
||||
search_uuid=search_uuid,
|
||||
generate_related_questions=generate_related_questions,
|
||||
)
|
||||
return StreamingResponse(content=result, media_type="text/html")
|
||||
|
||||
# First, do a search query.
|
||||
query = query or _default_query
|
||||
# Basic attack protection: remove "[INST]" or "[/INST]" from the query
|
||||
query = re.sub(r"\[/?INST\]", "", query)
|
||||
contexts = self.search_function(query)
|
||||
|
||||
system_prompt = _rag_query_text.format(
|
||||
context="\n\n".join(
|
||||
[f"[[citation:{i+1}]] {c['snippet']}" for i, c in enumerate(contexts)]
|
||||
)
|
||||
)
|
||||
try:
|
||||
client = self.local_client()
|
||||
llm_response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": query},
|
||||
],
|
||||
max_tokens=1024,
|
||||
stop=stop_words,
|
||||
stream=True,
|
||||
temperature=0.9,
|
||||
)
|
||||
if self.should_do_related_questions and generate_related_questions:
|
||||
# While the answer is being generated, we can start generating
|
||||
# related questions as a future.
|
||||
related_questions_future = self.executor.submit(
|
||||
self.get_related_questions, query, contexts
|
||||
)
|
||||
else:
|
||||
related_questions_future = None
|
||||
except Exception as e:
|
||||
logger.error(f"encountered error: {e}\n{traceback.format_exc()}")
|
||||
return HTMLResponse("Internal server error.", 503)
|
||||
|
||||
return StreamingResponse(
|
||||
self.stream_and_upload_to_kv(
|
||||
contexts, llm_response, related_questions_future, search_uuid
|
||||
),
|
||||
media_type="text/html",
|
||||
)
|
||||
|
||||
@Photon.handler(mount=True)
|
||||
def ui(self):
|
||||
return StaticFiles(directory="ui")
|
||||
|
||||
@Photon.handler(method="GET", path="/")
|
||||
def index(self) -> RedirectResponse:
|
||||
"""
|
||||
Redirects "/" to the ui page.
|
||||
"""
|
||||
return RedirectResponse(url="/ui/index.html")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
rag = RAG()
|
||||
rag.launch()
|
||||
Reference in New Issue
Block a user