diff --git a/main.py b/main.py index 442c8fd..cb94d9e 100644 --- a/main.py +++ b/main.py @@ -1,4 +1,4 @@ -from src.infuzu import (create_chat_completion, ChatCompletionsHandlerRequestMessage) +from src.infuzu import (create_chat_completion, ChatCompletionsHandlerRequestMessage, ChatCompletionsObject) from dotenv import load_dotenv @@ -12,7 +12,7 @@ try: - response: dict[str, any] = create_chat_completion(messages=messages) + response: ChatCompletionsObject = create_chat_completion(messages=messages) print(response) except Exception as e: print(f"Error: {e}") diff --git a/src/infuzu/__init__.py b/src/infuzu/__init__.py index a8e0648..4d59bc5 100644 --- a/src/infuzu/__init__.py +++ b/src/infuzu/__init__.py @@ -1,12 +1,41 @@ from .api_client import ( - create_chat_completion, ChatCompletionsHandlerRequestMessage, ChatCompletionsRequestContentPart + create_chat_completion, + ChatCompletionsHandlerRequestMessage, + ChatCompletionsRequestContentPart, + ChatCompletionsChoiceMessageAudioObject, + ChatCompletionsChoiceMessageFunctionCallObject, + ChatCompletionsChoiceMessageToolCallFunctionObject, + chatCompletionsChoiceMessageToolCallObject, + ChatCompletionsChoiceMessageObject, + ChatCompletionsChoiceLogprobsItemTopLogprobObject, + ChatCompletionsLogprobsItemObject, + ChatCompletionsChoiceLogprobsObject, + ChatCompletionsChoiceModelObject, + ChatCompletionsChoiceErrorObject, + ChatCompletionsChoiceLatencyObject, + ChatCompletionsChoiceObject, + ChatCompletionsObject, ) from .errors import (InfuzuAPIError, APIWarning, APIError) + __all__: list[str] = [ "create_chat_completion", "ChatCompletionsRequestContentPart", "ChatCompletionsHandlerRequestMessage", + "ChatCompletionsChoiceMessageAudioObject", + "ChatCompletionsChoiceMessageFunctionCallObject", + "ChatCompletionsChoiceMessageToolCallFunctionObject", + "chatCompletionsChoiceMessageToolCallObject", + "ChatCompletionsChoiceMessageObject", + "ChatCompletionsChoiceLogprobsItemTopLogprobObject", + "ChatCompletionsLogprobsItemObject", + "ChatCompletionsChoiceLogprobsObject", + "ChatCompletionsChoiceModelObject", + "ChatCompletionsChoiceErrorObject", + "ChatCompletionsChoiceLatencyObject", + "ChatCompletionsChoiceObject", + "ChatCompletionsObject", "InfuzuAPIError", "APIWarning", diff --git a/src/infuzu/api_client.py b/src/infuzu/api_client.py index 8b23e2d..b8c7b7c 100644 --- a/src/infuzu/api_client.py +++ b/src/infuzu/api_client.py @@ -1,7 +1,9 @@ +import time +import uuid import httpx import os from typing import (Optional, Dict, Union, List) -from pydantic import (BaseModel, validator) +from pydantic import (BaseModel, validator, Field) from .errors import InfuzuAPIError @@ -62,6 +64,131 @@ def role_must_be_valid(cls, v): return v +class ChatCompletionsChoiceMessageAudioObject(BaseModel): + id: Optional[str] = None + expired_at: Optional[int] = None + data: Optional[str] = None + transcript: Optional[str] = None + + class Config: + extra: str = "allow" + + +class ChatCompletionsChoiceMessageFunctionCallObject(BaseModel): + name: Optional[str] = None + arguments: Optional[str] = None + + class Config: + extra: str = "allow" + + +class ChatCompletionsChoiceMessageToolCallFunctionObject(BaseModel): + name: Optional[str] = None + arguments: Optional[str] = None + + class Config: + extra: str = "allow" + + +class chatCompletionsChoiceMessageToolCallObject(BaseModel): + id: Optional[str] = None + type: Optional[str] = None + function: Optional[ChatCompletionsChoiceMessageToolCallFunctionObject] = None + + class Config: + extra: str = "allow" + + +class ChatCompletionsChoiceMessageObject(BaseModel): + content: Optional[str] = None + refusal: Optional[str] = None + tool_calls: Optional[List[chatCompletionsChoiceMessageToolCallObject]] = None + role: Optional[str] = None + function_call: Optional[ChatCompletionsChoiceMessageFunctionCallObject] = None + audio: Optional[ChatCompletionsChoiceMessageAudioObject] = None + + class Config: + extra: str = "allow" + + +class ChatCompletionsChoiceLogprobsItemTopLogprobObject(BaseModel): + token: Optional[str] = None + logprob: Optional[int] = None + bytes: Optional[List[int]] = None + + class Config: + extra: str = "allow" + + +class ChatCompletionsLogprobsItemObject(BaseModel): + token: Optional[str] = None + logprob: Optional[int] = None + bytes: Optional[List[int]] = None + content: Optional[List[ChatCompletionsChoiceLogprobsItemTopLogprobObject]] = None + + class Config: + extra: str = "allow" + + +class ChatCompletionsChoiceLogprobsObject(BaseModel): + content: Optional[List[ChatCompletionsLogprobsItemObject]] = None + refusal: Optional[List[ChatCompletionsLogprobsItemObject]] = None + + class Config: + extra: str = "allow" + + +class ChatCompletionsChoiceModelObject(BaseModel): + ref: Optional[str] = None + rank: Optional[int] = None + + class Config: + extra: str = "allow" + + +class ChatCompletionsChoiceErrorObject(BaseModel): + message: Optional[str] = None + code: Optional[str] = None + + class Config: + extra: str = "allow" + + +class ChatCompletionsChoiceLatencyObject(BaseModel): + start: Optional[int] = Field(None, alias='start_latency') + end: Optional[int] = Field(None, alias='end_latency') + + class Config: + extra: str = "allow" + + +class ChatCompletionsChoiceObject(BaseModel): + finish_reason: Optional[str] = None + index: Optional[int] = None + message: Optional[ChatCompletionsChoiceMessageObject] = None + logprobs: Optional[ChatCompletionsChoiceLogprobsObject] = None + model: Optional[ChatCompletionsChoiceModelObject] = None + error: Optional[ChatCompletionsChoiceErrorObject] = None + latency: Optional[ChatCompletionsChoiceLatencyObject] = None + + class Config: + extra: str = "allow" + + +class ChatCompletionsObject(BaseModel): + id: Optional[str] = None + choices: Optional[List[ChatCompletionsChoiceObject]] = None + created: Optional[int] = None + model: Optional[str] = None + service_tier: Optional[str] = None + system_fingerprint: Optional[str] = None + object: Optional[str] = None + usage: Optional[Dict[str, int]] = None + + class Config: + extra: str = "allow" + + API_BASE_URL = "https://chat.infuzu.com/api" @@ -69,7 +196,7 @@ def create_chat_completion( messages: List[ChatCompletionsHandlerRequestMessage], api_key: Optional[str] = None, model: Optional[Union[str, InfuzuModelParams]] = None, -) -> Dict: +) -> ChatCompletionsObject: """ Creates a chat completion using the Infuzu API. @@ -121,6 +248,13 @@ def create_chat_completion( ) response.raise_for_status() - return response.json() + json_response: dict[str, any] = response.json() + + json_response.setdefault('id', f"chatcmpl-{uuid.uuid4()}") + json_response.setdefault('created', int(time.time())) + json_response.setdefault('model', 'infuzu-ims') + json_response.setdefault('object', 'chat.completion') + + return ChatCompletionsObject(**json_response) except httpx.HTTPStatusError as e: raise InfuzuAPIError(e)