mirror of
https://github.com/StanGirard/quivr.git
synced 2024-09-11 13:15:41 +03:00
feat: stream router (#353)
* wip: stream router * feat: chatai streaming * chore: add comments * feat: streaming for chains * chore: comments
This commit is contained in:
parent
90bd49527b
commit
3e753f2d56
@ -10,6 +10,7 @@ from routes.chat_routes import chat_router
|
||||
from routes.crawl_routes import crawl_router
|
||||
from routes.explore_routes import explore_router
|
||||
from routes.misc_routes import misc_router
|
||||
from routes.stream_routes import stream_router
|
||||
from routes.upload_routes import upload_router
|
||||
from routes.user_routes import user_router
|
||||
|
||||
@ -19,12 +20,14 @@ app = FastAPI()
|
||||
|
||||
add_cors_middleware(app)
|
||||
max_brain_size = os.getenv("MAX_BRAIN_SIZE")
|
||||
max_brain_size_with_own_key = os.getenv("MAX_BRAIN_SIZE_WITH_KEY",209715200)
|
||||
max_brain_size_with_own_key = os.getenv("MAX_BRAIN_SIZE_WITH_KEY", 209715200)
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
async def startup_event():
|
||||
pypandoc.download_pandoc()
|
||||
|
||||
|
||||
app.include_router(brain_router)
|
||||
app.include_router(chat_router)
|
||||
app.include_router(crawl_router)
|
||||
@ -33,3 +36,4 @@ app.include_router(misc_router)
|
||||
app.include_router(upload_router)
|
||||
app.include_router(user_router)
|
||||
app.include_router(api_key_router)
|
||||
app.include_router(stream_router)
|
||||
|
123
backend/routes/stream_routes.py
Normal file
123
backend/routes/stream_routes.py
Normal file
@ -0,0 +1,123 @@
|
||||
import asyncio
|
||||
import os
|
||||
from typing import AsyncIterable, Awaitable
|
||||
|
||||
from auth.auth_bearer import AuthBearer, get_current_user
|
||||
from fastapi import APIRouter, Depends
|
||||
from fastapi.responses import StreamingResponse
|
||||
from langchain.callbacks import AsyncIteratorCallbackHandler
|
||||
from langchain.chains import ConversationalRetrievalChain
|
||||
from langchain.chains.llm import LLMChain
|
||||
from langchain.chains.question_answering import load_qa_chain
|
||||
from langchain.chat_models import ChatOpenAI
|
||||
from langchain.embeddings import OpenAIEmbeddings
|
||||
from llm.prompt.CONDENSE_PROMPT import CONDENSE_QUESTION_PROMPT
|
||||
from logger import get_logger
|
||||
from models.chats import ChatMessage
|
||||
from models.settings import CommonsDep, common_dependencies
|
||||
from models.users import User
|
||||
from supabase import create_client
|
||||
from utils.users import fetch_user_id_from_credentials
|
||||
from vectorstore.supabase import CustomSupabaseVectorStore
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
stream_router = APIRouter()
|
||||
|
||||
openai_api_key = os.getenv("OPENAI_API_KEY")
|
||||
supabase_url = os.getenv("SUPABASE_URL")
|
||||
supabase_service_key = os.getenv("SUPABASE_SERVICE_KEY")
|
||||
|
||||
|
||||
async def send_message(
|
||||
chat_message: ChatMessage, chain, callback
|
||||
) -> AsyncIterable[str]:
|
||||
async def wrap_done(fn: Awaitable, event: asyncio.Event):
|
||||
"""Wrap an awaitable with a event to signal when it's done or an exception is raised."""
|
||||
try:
|
||||
resp = await fn
|
||||
logger.debug("Done: %s", resp)
|
||||
except Exception as e:
|
||||
logger.error(f"Caught exception: {e}")
|
||||
finally:
|
||||
# Signal the aiter to stop.
|
||||
event.set()
|
||||
|
||||
# Use the agenerate method for models.
|
||||
# Use the acall method for chains.
|
||||
task = asyncio.create_task(
|
||||
wrap_done(
|
||||
chain.acall(
|
||||
{
|
||||
"question": chat_message.question,
|
||||
"chat_history": chat_message.history,
|
||||
}
|
||||
),
|
||||
callback.done,
|
||||
)
|
||||
)
|
||||
|
||||
# Use the aiter method of the callback to stream the response with server-sent-events
|
||||
async for token in callback.aiter():
|
||||
logger.info("Token: %s", token)
|
||||
yield f"data: {token}\n\n"
|
||||
|
||||
await task
|
||||
|
||||
|
||||
def create_chain(commons: CommonsDep, current_user: User):
|
||||
user_id = fetch_user_id_from_credentials(commons, {"email": current_user.email})
|
||||
|
||||
embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)
|
||||
|
||||
supabase_client = create_client(supabase_url, supabase_service_key)
|
||||
|
||||
vector_store = CustomSupabaseVectorStore(
|
||||
supabase_client, embeddings, table_name="vectors", user_id=user_id
|
||||
)
|
||||
|
||||
generator_llm = ChatOpenAI(
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
# Callback provides the on_llm_new_token method
|
||||
callback = AsyncIteratorCallbackHandler()
|
||||
|
||||
streaming_llm = ChatOpenAI(
|
||||
temperature=0,
|
||||
streaming=True,
|
||||
callbacks=[callback],
|
||||
)
|
||||
question_generator = LLMChain(
|
||||
llm=generator_llm,
|
||||
prompt=CONDENSE_QUESTION_PROMPT,
|
||||
)
|
||||
doc_chain = load_qa_chain(
|
||||
llm=streaming_llm,
|
||||
chain_type="stuff",
|
||||
)
|
||||
|
||||
return (
|
||||
ConversationalRetrievalChain(
|
||||
combine_docs_chain=doc_chain,
|
||||
question_generator=question_generator,
|
||||
retriever=vector_store.as_retriever(),
|
||||
verbose=True,
|
||||
),
|
||||
callback,
|
||||
)
|
||||
|
||||
|
||||
@stream_router.post("/stream", dependencies=[Depends(AuthBearer())], tags=["Stream"])
|
||||
async def stream(
|
||||
chat_message: ChatMessage,
|
||||
current_user: User = Depends(get_current_user),
|
||||
) -> StreamingResponse:
|
||||
commons = common_dependencies()
|
||||
|
||||
qa_chain, callback = create_chain(commons, current_user)
|
||||
|
||||
return StreamingResponse(
|
||||
send_message(chat_message, qa_chain, callback),
|
||||
media_type="text/event-stream",
|
||||
)
|
Loading…
Reference in New Issue
Block a user