import asyncio
import uvicorn

from typing import AsyncIterable, Awaitable
from dotenv import load_dotenv
from fastapi import FastAPI
from fastapi.responses import FileResponse, StreamingResponse
from langchain.callbacks import AsyncIteratorCallbackHandler
from langchain.chat_models import ChatOpenAI
from langchain.schema import HumanMessage

load_dotenv()
async def wait_done(fn: Awaitable, event: asyncio.Event):
    try:
        await fn
    except Exception as e:
        print(e)
        event.set()
    finally:
        event.set()

async def call_openai(question: str) -> AsyncIterable[str]:
    callback = AsyncIteratorCallbackHandler()
    model = ChatOpenAI(streaming=True, verbose=True, callbacks=[callback])

    coroutine = wait_done(model.agenerate(messages=[[HumanMessage(content=question)]]), callback.done)
    task = asyncio.create_task(coroutine)

    async for token in callback.aiter():
        yield f"{token}"

    await task


app = FastAPI()

@app.post("/ask")
def ask(body: dict):
    return StreamingResponse(call_openai(body['question']), media_type="text/event-stream")

@app.get("/")
async def homepage():
    return FileResponse('statics/index.html')

if __name__ == "__main__":
    uvicorn.run(host="0.0.0.0", port=8888, app=app)
langchain.callbacks.AsyncIteratorCallbackHandler langchain.schema.HumanMessage langchain.chat_models.ChatOpenAI
[((345, 358), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (356, 358), False, 'from dotenv import load_dotenv\n'), ((959, 968), 'fastapi.FastAPI', 'FastAPI', ([], {}), '()\n', (966, 968), False, 'from fastapi import FastAPI\n'), ((616, 646), 'langchain.callbacks.AsyncIteratorCallbackHandler', 'AsyncIteratorCallbackHandler', ([], {}), '()\n', (644, 646), False, 'from langchain.callbacks import AsyncIteratorCallbackHandler\n'), ((659, 721), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'streaming': '(True)', 'verbose': '(True)', 'callbacks': '[callback]'}), '(streaming=True, verbose=True, callbacks=[callback])\n', (669, 721), False, 'from langchain.chat_models import ChatOpenAI\n'), ((837, 867), 'asyncio.create_task', 'asyncio.create_task', (['coroutine'], {}), '(coroutine)\n', (856, 867), False, 'import asyncio\n'), ((1149, 1183), 'fastapi.responses.FileResponse', 'FileResponse', (['"""statics/index.html"""'], {}), "('statics/index.html')\n", (1161, 1183), False, 'from fastapi.responses import FileResponse, StreamingResponse\n'), ((1216, 1263), 'uvicorn.run', 'uvicorn.run', ([], {'host': '"""0.0.0.0"""', 'port': '(8888)', 'app': 'app'}), "(host='0.0.0.0', port=8888, app=app)\n", (1227, 1263), False, 'import uvicorn\n'), ((776, 806), 'langchain.schema.HumanMessage', 'HumanMessage', ([], {'content': 'question'}), '(content=question)\n', (788, 806), False, 'from langchain.schema import HumanMessage\n')]