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

1from typing import AsyncIterable, cast 

2 

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 

9 

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 

14 

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

16 

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

18 

19 

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) 

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 

39for subrouter in [ 

40 chat_router, 

41 models_router, 

42 files_router, 

43 beta_router, 

44]: 

45 router.include_router(subrouter)