add dynamic clients for all APIs

This commit is contained in:
Ashwin Bharambe 2024-10-30 16:04:16 -07:00
parent f04b566c5c
commit 4067038f74
3 changed files with 198 additions and 4 deletions

View file

@ -6,7 +6,15 @@
from enum import Enum
from typing import List, Literal, Optional, Protocol, runtime_checkable, Union
from typing import (
AsyncIterator,
List,
Literal,
Optional,
Protocol,
runtime_checkable,
Union,
)
from llama_models.schema_utils import json_schema_type, webmethod
@ -224,7 +232,7 @@ class Inference(Protocol):
response_format: Optional[ResponseFormat] = None,
stream: Optional[bool] = False,
logprobs: Optional[LogProbConfig] = None,
) -> Union[CompletionResponse, CompletionResponseStreamChunk]: ...
) -> Union[CompletionResponse, AsyncIterator[CompletionResponseStreamChunk]]: ...
@webmethod(route="/inference/chat_completion")
async def chat_completion(
@ -239,7 +247,9 @@ class Inference(Protocol):
response_format: Optional[ResponseFormat] = None,
stream: Optional[bool] = False,
logprobs: Optional[LogProbConfig] = None,
) -> Union[ChatCompletionResponse, ChatCompletionResponseStreamChunk]: ...
) -> Union[
ChatCompletionResponse, AsyncIterator[ChatCompletionResponseStreamChunk]
]: ...
@webmethod(route="/inference/embeddings")
async def embeddings(