from SimpleLLMFunc.interface import LLM_Interface, APIKeyPool
class MyProviderInterface(LLM_Interface):
def __init__(self, api_key_pool: APIKeyPool, model_name: str, base_url: str):
super().__init__(api_key_pool, model_name, base_url)
# Initialize your client here
async def chat(self, *, trace_id, stream=False, messages, timeout=None, **kwargs):
api_key = self.api_key_pool.get_key()
# Make your API call
response = await my_client.complete(
model=self.model_name,
messages=messages,
api_key=api_key,
)
# Return in ChatCompletion-compatible shape
return self._to_chat_completion(response)
async def chat_stream(self, *, trace_id, stream=True, messages, timeout=None, **kwargs):
api_key = self.api_key_pool.get_key()
async for chunk in my_client.stream(
model=self.model_name,
messages=messages,
api_key=api_key,
):
yield self._to_chunk(chunk)