Coverage for src/fastoai/routers/__init__.py: 65%

1 from typing import AsyncIterable, cast

2

3 from fastapi import APIRouter, Depends

4 from fastapi.responses import StreamingResponse

5 from openai.types.chat.chat_completion import ChatCompletion

6 from openai.types.chat.chat_completion_chunk import ChatCompletionChunk

7 from openai.types.chat.completion_create_params import CompletionCreateParams

8 from pydantic import RootModel

9

10 from ..dependencies import OpenAIDependency, get_user

11 from .beta import router as beta_router

12 from .files import router as files_router

13 from .models import router as models_router

14

15 router = APIRouter(dependencies=[Depends(get_user)])

16

17 chat_router = APIRouter(tags=["Chat"])

18

19

20@chat_router.post("/chat/completions")

21 async 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)

28

29 async def _stream(chunks: AsyncIterable[ChatCompletionChunk]):

30 async for chunk in chunks:

31 yield f"data: {chunk.model_dump_json()}\n\n"

32

33 yield "data: [DONE]\n\n"

34

35 return StreamingResponse(_stream(response))

36 return cast(ChatCompletion, response)

37

38

39 for subrouter in [

40 chat_router,

41 models_router,

42 files_router,

43 beta_router,

44]:

45 router.include_router(subrouter)