FastAPIを使ってLLMの出力をリアルタイム送信する方法

はじめに

本記事ではFastAPIのストリーミングレスポンス機能を利用して、LLMの出力をAPIでリアルタイムに送信する方法を解説します。

方法

FastAPIでストリーミングを行うには、StreamingResponseを使用します。 StreamingResponseはcontentを複数のチャンクに分割してストリームで返すためのレスポンスクラスです。 contentの型はIteratorかAsyncIteratorである必要があります。

下記ではOpenAIのCHatCompletionとLangChainを使った場合のサンプルコードを記載します。

from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from models.chat import ChatResponse, InitChatResponse, StreamRequest

router = APIRouter()

async def completion_stream(prompt: str) -> AsyncIterable[str]:
    stream = openai_client.chat.completions.create(
        stream=True,
        model=OPENAI_MODEL_NAME,
        temperature=OPENAI_TEMPERATURE,
        messages=[{"role": "user", "content": prompt}],
    )
    for item in stream:
        content = item.choices[0].delta.content
        if content:
            yield content


async def chain_stream(message: str) -> AsyncIterable[str | list[str | dict[Any, Any]]]:
    llm = ChatOpenAI(api_key=OPENAI_API_KEY)
    messages = [HumanMessagePromptTemplate.from_template(template="{message}")]
    prompt = ChatPromptTemplate.from_messages(messages)
    chain = prompt | llm
    stream = chain.astream({"message": message})
    async for msg in stream:
        yield msg.content

@router.post("/stream-completion", response_model=ChatResponse)
async def stream_completion(req: StreamRequest) -> StreamingResponse:
    try:
        return StreamingResponse(completion_stream(req.prompt), media_type="text/event-stream")
    except TimeoutError as e:
        print("Error:", e)
        raise HTTPException(status_code=408, detail="Timeout Error")
    except Exception as e:
        print("Error:", e)
        raise HTTPException(status_code=500, detail="Internal Service Error")


@router.post("/stream-chain", response_model=ChatResponse)
async def stream_chain(req: StreamRequest) -> StreamingResponse:
    try:
        return StreamingResponse(chain_stream(req.prompt), media_type="text/event-stream")
    except TimeoutError as e:
        print("Error:", e)
        raise HTTPException(status_code=408, detail="Timeout Error")
    except Exception as e:
        print("Error:", e)
        raise HTTPException(status_code=500, detail="Internal Service Error")

動作確認

サーバを起動し、リクエストを送ります。

curl -N -X POST -H "Content-Type: application/json" -d '{"prompt": "write poem aboutsparkling water"}' http://0.0.0.0:8080/api/stream-chain

レスポンスがストリーミングになっていることが確認できます。

まとめ

本記事ではFastAPIのストリーミングレスポンス機能を利用してChatGPTのようなストリーミングレスポンスをAPIで実現する方法を紹介しました。この方法はテキストのリアルタイムな送信以外にも動画などの大きなファイルの送信にも有効な方法な方法です。ぜひ使ってみてください。

参考

備考

Hakky ではエンジニアを募集中です！まずは話してみたいなどでも構いませんので、ぜひお気軽に採用ページからお問い合わせくださいませ。