autogen icon indicating copy to clipboard operation
autogen copied to clipboard

[Feature Request]: Async ModelClient

Open mphilippnv opened this issue 1 year ago • 0 comments

Is your feature request related to a problem? Please describe.

I recently implemented a custom ModelClient class per the post here and it's working successfully. However, I would love to be able to implement an async version of the class. I think this would be very useful for lots of applications where "async all the way down" is pretty standard. Given how long inferencing operations can take, making things async to free up the application makes a lot of sense.

Of course, if I'm being obtuse or missing something obvious, do let me know!

Describe the solution you'd like

I think having a ModelClientAsync base class makes sense. This class could contain the same interface as the ModelClient class except with async definitions.

class ModelClientAsync(Protocol):
    RESPONSE_USAGE_KEYS = ["prompt_tokens", "completion_tokens", "total_tokens", "cost", "model"]

    class ModelClientResponseProtocol(Protocol):
        class Choice(Protocol):
            class Message(Protocol):
                content: Optional[str]

            message: Message

        choices: List[Choice]
        model: str

    async def create(self, params: Dict[str, Any]) -> ModelClientResponseProtocol: ...

    async def message_retrieval(
        self, response: ModelClientResponseProtocol
    ) -> Union[List[str], List[ModelClient.ModelClientResponseProtocol.Choice.Message]]: ...
       
    async def cost(self, response: ModelClientResponseProtocol) -> float: ...

    @staticmethod
    async def get_usage(response: ModelClientResponseProtocol) -> Dict: ... 

Alternatively, maybe just keep the single ModelClient class but add async methods for each: create_async, cost_async, get_usage_async, and message_retrieval_async.

In my own use case, I'd only need the create method to be async. But I think it would make sense to open up all methods to be async just in case users might get usage, cost, etc. from an API endpoint or something like that.

Additional context

Unsure if it's helpful for testing, but Azure has an async openAI library called AsyncAzureOpenAI. This can be used with the Azure LLM Gateway.

mphilippnv avatar Jul 30 '24 05:07 mphilippnv