Minimal Working Example of FastAPI using LangChain Stream Response

main.py

# Answer taken from https://gist.github.com/jvelezmagic/f3653cc2ddab1c91e86751c8b423a1b6


from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from pydantic import BaseModel

from typing import AsyncGenerator


class ChatRequest(BaseModel):
    session_id: str
    message: str


app = FastAPI(
    title="QA Chatbot Streaming using FastAPI, LangChain Expression Language , OpenAI, and Chroma",
    version="0.1.0",
)


async def generate_response(
    context: str,
    message: str,
) -> AsyncGenerator[str, None]:
    prompt = PromptTemplate.from_template(
        """Answer the question based only on the following context:
{context}
Question: {question}"""
    )

    llm = ChatOpenAI(
        temperature=0,
        openai_api_key=YOUR_OPENAI_KEY,
    )

    chain = prompt | llm  # type: ignore

    response = ""
    async for token in chain.astream({"context": context, "question": message}):  # type: ignore
        yield token.content
        response += token.content


@app.post("/chat")
async def chat(
    request: ChatRequest,
) -> StreamingResponse:
    return StreamingResponse(
        generate_response(
            context="",
            message=request.message,
        ),
        media_type="text/plain",
    )