FastAPIでServer-Sent Eventsを実装する 

## 導入
ChatGPT APIはStream出力をサポートしており、これを使うことで本家ChatGPTのように結果を順次表示することが可能です。

pythonを使ってStreaming出力を利用する場合は以下のように実装できます

```python
response = openai.ChatCompletion.create(
    model="gpt-3.5-turbo-0613",
    messages=[
        {'role': 'user', 'content': 'Hello?'}
    ],
    stream=True
)
for chunk in response:
    chunk_message = chunk['choices'][0]['delta'].get('content', '')
    print(chunk_message)
```

langchainを使うとこう

```python
llm = OpenAI(
    model="gpt-3.5-turbo-instruct", 
)
for chunk in llm.stream("Hello?"):
    print(chunk)
```

このように非常に簡単に実装可能です。
しかし、これをサーバーサイドとして扱う場合には通常のAPIではなく、リアルタイムな通信処理を実装する必要があります。
この時に有効なのがServer-Sent Eventsです。

## Server-Sent Eventsとは 
[Server-Sent Events](https://developer.mozilla.org/ja/docs/Web/API/Server-sent_events/Using_server-sent_events)(以下SSE)とはリアルタイムな通信をHTTP経由で提供する技術です。
似た技術としてはWebSocketがありますが、WebSocketが双方向であるのに対し、SSEはサーバーからクライアントへの単方向通信となっています。
その分扱いが簡単なので、クライアント側からリアルタイムに情報を送る必要がない場合はこちらが適しています。

## FastAPIによる実装
さて、今回はサーバーサイドに[FastAPI](https://fastapi.tiangolo.com/ja/)を利用します。
FastAPIは非常に簡易なpython用サーバーサイドフレームワークであり、SwaggerUIを自動生成してくれるため、プロトタイプ制作に適しています。
FastAPIにおけるSSEは次のようなコードで実装できます。

```python
## ファイル名はmain.py

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import uvicorn
from pydantic import BaseModel
import asyncio

app = FastAPI()

class StreamRequest(BaseModel):
    query: str

async def send_token(query: str):
    for token in query:
        yield token
        await asyncio.sleep(0.1)

@app.post("/streaming")
async def ask_stream(request: StreamRequest) -> StreamingResponse:
    return StreamingResponse(
        send_token(request.query),
        media_type="text/event-stream",
    )

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
```

FastAPIの[StreamingResponse](https://fastapi.tiangolo.com/ja/advanced/custom-response/#streamingresponse)はGeneratorを受け取ってevent-streamを作成します。
今回はsend_token関数定義内で[yield](https://docs.python.org/ja/3/reference/simple_stmts.html#yield)を使用することでGenerator関数として扱っています。

適当なshellで`python main.py`を実行してサーバーを起動したのち、
次のコマンドを実行して動作を確認します。

```shell
curl -N -X POST  \
 -H "Content-Type: application/json" \
 -H "accept: text/event-stream" \
 -d '{"query":"Streaming responseのテスト"}' \
 http://127.0.0.1:8000/streaming
```

入力したテキストが良い感じに1文字ずつ出力されたら成功です。
ChatGPT APIに置き換えるとこんな感じになります。

```python
async def send_token(query: str):
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo-0613",
        messages=[
            {'role': 'user', 'content': query}
        ],
        stream=True
    )
    for chunk in response:
        chunk_message = chunk['choices'][0]['delta'].get('content', '')
        yield chunk_message

```
text/event-stream形式はテキストしか扱えませんが、json.dumpsメソッドなどを利用することで構造データを送ることも可能なため、使い方はいろいろあると思います。

## 複雑な処理が必要な場合の拡張

さて、基本的には上記の方法で良いのですが、時折コールバックのようなyieldを利用できない処理を使わなくてはならないことがあると思います。(というか、以前のlangchainがそうでした)
こういった場合、Iteratorを自作してやることで柔軟な実装が可能となります。

まずStreamingResponseに渡す[Iterator](https://docs.python.org/ja/3.8/library/stdtypes.html#typeiter)を用意します。
Iteratorはイテレーターオブジェクト自体を返す__iter__メソッドと、データを順次出力し、最後にStopIterationを出力する__next__メソッドを含むクラスとして以下のように実装できます。

```python
import queue, asyncio

class StreamIterator:
    def __init__(self):
        self.queue = queue.Queue()
    def __iter__(self):
        return self
    def __next__(self):
        item = self.queue.get()
        if item is StopIteration:
            raise item
        return item
    def send(self, data):
        self.queue.put(data)
    def close(self):
        self.queue.put(StopIteration)
```

この実装ではiteratorのsendメソッドを呼び出すことでテキストを送信できます。
また、処理中にエラーが起きた場合、iteratorが残っていると終了できなくなるためfinallyでcloseメソッドを実行します。

```python
async def main_process(query: str, iterator: StreamIterator):
    try:
        await send_token(query, iterator)
    except Exception as err:
        print(err)
    finally:
        iterator.close()

async def send_token(query: str, iterator: StreamIterator):
    for token in query:
        iterator.send(token)
        await asyncio.sleep(0.1)
```

これらのプロセスを[asyncio](https://docs.python.org/3.11/library/asyncio-eventloop.html)モジュールのイベントループ上で実行し、
StreamingResponseに渡したiteratorにmain_process内でデータを格納することでSSEを実装できます。

```python
@app.post("/streaming")
async def ask_stream(request: StreamRequest) -> StreamingResponse:
    iterator = StreamIterator()
    event_loop = asyncio.get_event_loop()
    event_loop.create_task(main_process(request.query, iterator))

    return StreamingResponse(
        iterator,
        media_type="text/event-stream",
    )
```

再び適当なshellで`python main.py`を実行してサーバーを起動したのち、
次のコマンドを実行して動作を確認します。

```shell
curl -N -X POST  \
 -H "Content-Type: application/json" \
 -H "accept: text/event-stream" \
 -d '{"query":"Streaming responseのテスト"}' \
 http://127.0.0.1:8000/streaming
```

入力したテキストが良い感じに1文字ずつ出力されたら成功です。

この方式ではiteratorを渡すことでコールバックなどyieldの扱いが難しい処理でもSSEでデータを送ることが可能となります。

## 結び
今回はFastAPIでServer-Sents Eventを実装してみました。
Iteratorの自作は多分しばらく使う機会はないと思いますが、仕組みの理解の上では役立つので暇があれば触ってみても良いかもしれません