Coverage for src/fastoai/routers/__init__.py: 65%
26 statements
« prev ^ index » next coverage.py v7.6.8, created at 2024-12-06 09:34 +0800
« prev ^ index » next coverage.py v7.6.8, created at 2024-12-06 09:34 +0800
1from typing import AsyncIterable, cast
3from fastapi import APIRouter, Depends
4from fastapi.responses import StreamingResponse
5from openai.types.chat.chat_completion import ChatCompletion
6from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
7from openai.types.chat.completion_create_params import CompletionCreateParams
8from pydantic import RootModel
10from ..dependencies import OpenAIDependency, get_user
11from .beta import router as beta_router
12from .files import router as files_router
13from .models import router as models_router
15router = APIRouter(dependencies=[Depends(get_user)])
17chat_router = APIRouter(tags=["Chat"])
20@chat_router.post("/chat/completions")
21async def create_chat_completions(
22 params: RootModel[CompletionCreateParams],
23 openai: OpenAIDependency,
24):
25 response = await openai.chat.completions.create(**params.model_dump())
26 if params.root.get("stream", False):
27 response = cast(AsyncIterable[ChatCompletionChunk], response)
29 async def _stream(chunks: AsyncIterable[ChatCompletionChunk]):
30 async for chunk in chunks:
31 yield f"data: {chunk.model_dump_json()}\n\n"
33 yield "data: [DONE]\n\n"
35 return StreamingResponse(_stream(response))
36 return cast(ChatCompletion, response)
39for subrouter in [
40 chat_router,
41 models_router,
42 files_router,
43 beta_router,
44]:
45 router.include_router(subrouter)