Skip to content

Commit

Permalink
feat: async chat completion request type
Browse files Browse the repository at this point in the history
  • Loading branch information
csgulati09 committed Dec 20, 2024
1 parent aec433b commit bb5989e
Showing 1 changed file with 65 additions and 2 deletions.
67 changes: 65 additions & 2 deletions portkey_ai/api_resources/apis/chat_complete.py
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,21 @@ def __init__(self, client: AsyncPortkey) -> None:
self.openai_client = client.openai_client

async def stream_create(
self, model, messages, stream, temperature, max_tokens, top_p, **kwargs
self,
model,
messages,
stream,
temperature,
max_tokens,
top_p,
audio,
max_completion_tokens,
metadata,
modalities,
prediction,
reasoning_effort,
store,
**kwargs,
) -> Union[ChatCompletions, AsyncIterator[ChatCompletionChunk]]:
async with self.openai_client.with_streaming_response.chat.completions.create(
model=model,
Expand All @@ -198,6 +212,13 @@ async def stream_create(
temperature=temperature,
max_tokens=max_tokens,
top_p=top_p,
audio=audio,
max_completion_tokens=max_completion_tokens,
metadata=metadata,
modalities=modalities,
prediction=prediction,
reasoning_effort=reasoning_effort,
store=store,
extra_body=kwargs,
) as response:
async for line in response.iter_lines():
Expand All @@ -215,7 +236,21 @@ async def stream_create(
pass

async def normal_create(
self, model, messages, stream, temperature, max_tokens, top_p, **kwargs
self,
model,
messages,
stream,
temperature,
max_tokens,
top_p,
audio,
max_completion_tokens,
metadata,
modalities,
prediction,
reasoning_effort,
store,
**kwargs,
) -> ChatCompletions:
response = await self.openai_client.with_raw_response.chat.completions.create(
model=model,
Expand All @@ -224,6 +259,13 @@ async def normal_create(
temperature=temperature,
max_tokens=max_tokens,
top_p=top_p,
audio=audio,
max_completion_tokens=max_completion_tokens,
metadata=metadata,
modalities=modalities,
prediction=prediction,
reasoning_effort=reasoning_effort,
store=store,
extra_body=kwargs,
)
data = ChatCompletions(**json.loads(response.text))
Expand All @@ -239,6 +281,13 @@ async def create(
temperature: Union[float, NotGiven] = NOT_GIVEN,
max_tokens: Union[int, NotGiven] = NOT_GIVEN,
top_p: Union[float, NotGiven] = NOT_GIVEN,
audio: Optional[Any] = NOT_GIVEN,
max_completion_tokens: Union[int, NotGiven] = NOT_GIVEN,
metadata: Union[Dict[str, str], NotGiven] = NOT_GIVEN,
modalities: Union[List[Any], NotGiven] = NOT_GIVEN,
prediction: Union[Any, NotGiven] = NOT_GIVEN,
reasoning_effort: Union[Any, NotGiven] = NOT_GIVEN,
store: Union[Optional[bool], NotGiven] = NOT_GIVEN,
**kwargs,
) -> Union[ChatCompletions, AsyncIterator[ChatCompletionChunk]]:
if stream is True:
Expand All @@ -249,6 +298,13 @@ async def create(
temperature=temperature,
max_tokens=max_tokens,
top_p=top_p,
audio=audio,
max_completion_tokens=max_completion_tokens,
metadata=metadata,
modalities=modalities,
prediction=prediction,
reasoning_effort=reasoning_effort,
store=store,
**kwargs,
)
else:
Expand All @@ -259,6 +315,13 @@ async def create(
temperature=temperature,
max_tokens=max_tokens,
top_p=top_p,
audio=audio,
max_completion_tokens=max_completion_tokens,
metadata=metadata,
modalities=modalities,
prediction=prediction,
reasoning_effort=reasoning_effort,
store=store,
**kwargs,
)

Expand Down

0 comments on commit bb5989e

Please sign in to comment.