forked from stacklok/codegate
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathbase.py
More file actions
61 lines (53 loc) · 1.77 KB
/
base.py
File metadata and controls
61 lines (53 loc) · 1.77 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
import inspect
from abc import ABC, abstractmethod
from collections.abc import Iterator
from typing import Any, AsyncIterator, Callable, Optional, Union
from fastapi.responses import JSONResponse, StreamingResponse
from codegate.clients.clients import ClientType
class BaseCompletionHandler(ABC):
"""
The completion handler is responsible for executing the completion request
and creating the streaming response.
"""
@abstractmethod
async def execute_completion(
self,
request: Any,
base_url: Optional[str],
api_key: Optional[str],
stream: bool = False, # TODO: remove this param?
is_fim_request: bool = False,
) -> Union[Any, AsyncIterator[Any]]:
"""Execute the completion request"""
pass
@abstractmethod
def _create_streaming_response(
self,
stream: AsyncIterator[Any],
client_type: ClientType = ClientType.GENERIC,
stream_generator: Callable | None = None,
) -> StreamingResponse:
pass
@abstractmethod
def _create_json_response(self, response: Any) -> JSONResponse:
pass
def create_response(
self,
response: Any,
client_type: ClientType,
stream_generator: Callable | None = None,
) -> Union[JSONResponse, StreamingResponse]:
"""
Create a FastAPI response from the completion response.
"""
if (
isinstance(response, Iterator)
or isinstance(response, AsyncIterator)
or inspect.isasyncgen(response)
):
return self._create_streaming_response(
response,
client_type,
stream_generator=stream_generator,
)
return self._create_json_response(response)