Skip to content

Commit

Permalink
add support for getting clients by index
Browse files Browse the repository at this point in the history
  • Loading branch information
trisongz committed Oct 6, 2023
1 parent 4ad50e0 commit f60033e
Show file tree
Hide file tree
Showing 2 changed files with 88 additions and 2 deletions.
76 changes: 75 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ result: CompletionResponse = asyncio.run(

```

### Initialize Clients Manually
### Initialize Clients Manually, and working with multiple clients

```python

Expand Down Expand Up @@ -166,6 +166,27 @@ result = oai.completions.create(
stream = True
)

# You can select the different clients by name or index
result = OpenAI['az'].completions.create(
prompt = 'say this is a test',
max_tokens = 4,
stream = True
)

# Use the default client (openai)
result = OpenAI['default'].completions.create(
prompt = 'say this is a test',
max_tokens = 4,
stream = True
)

# Will use the `default` client since it was initialized first
result = OpenAI[0].completions.create(
prompt = 'say this is a test',
max_tokens = 4,
stream = True
)

```

### Handling Errors, Retries, and Rotations
Expand Down Expand Up @@ -292,6 +313,59 @@ Result:

```

### Configure Azure Model Mapping

Your azure models may be named differently than the default mapping. By configuring the mapping, you can automatically map the models to the correct azure model (when using openai model names).

```python

from async_openai import OpenAI

"""
Default Azure Model Mapping
{
'gpt-3.5-turbo': 'gpt-35-turbo',
'gpt-3.5-turbo-16k': 'gpt-35-turbo-16k',
'gpt-3.5-turbo-instruct': 'gpt-35-turbo-instruct',
'gpt-3.5-turbo-0301': 'gpt-35-turbo-0301',
'gpt-3.5-turbo-0613': 'gpt-35-turbo-0613',
}
"""

AzureModelMapping = {
'gpt-3.5-turbo': 'azure-gpt-35-turbo',
'gpt-3.5-turbo-16k': 'azure-gpt-35-turbo-16k',
'gpt-3.5-turbo-instruct': 'azure-gpt-35-turbo-instruct',
'gpt-3.5-turbo-0301': 'azure-gpt-35-turbo-0301',
'gpt-3.5-turbo-0613': 'azure-gpt-35-turbo-0613',
}

OpenAI.configure(
api_key = "sk-XXXX",
organization = "org-XXXX",
debug_enabled = False,

# Azure Configuration
azure_api_base = 'https://....openai.azure.com/',
azure_api_version = '2023-07-01-preview',
azure_api_key = '....',
azure_model_mapping = AzureModelMapping,
)

# This will now use the azure endpoint as the default client
OpenAI.init_api_client('az', set_as_default = True, debug_enabled = True)

# This will automatically map "gpt-3.5-turbo-16k" -> "azure-gpt-35-turbo-16k"
result: ChatResponse = OpenAI.chat.create(
model = "gpt-3.5-turbo-16k",
messages = [
{"role": "user", "content": "Translate the following English text to French: “Multiple models, each with different capabilities and price points. Prices are per 1,000 tokens. You can think of tokens as pieces of words, where 1,000 tokens is about 750 words. This paragraph is 35 tokens”"}
],
auto_retry = True,
)


```


---
Expand Down
14 changes: 13 additions & 1 deletion async_openai/meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,14 @@ def get_api_client(self, client_name: Optional[str] = None, **kwargs) -> 'OpenAI
self.clients[client_name] = self.init_api_client(client_name = client_name, **kwargs)
return self.clients[client_name]

def __getitem__(self, key: Union[str, int]) -> 'OpenAIClient':
"""
Returns a client by name.
"""
if isinstance(key, int):
key = self.rotate_client_names[key] if self.rotate_client_names else self.client_names[key]
return self.clients[key]

# Model Mapping for Azure
DefaultModelMapping = {
'gpt-3.5-turbo': 'gpt-35-turbo',
Expand Down Expand Up @@ -717,10 +725,14 @@ async def __aenter__(cls):
async def __aexit__(cls, exc_type, exc_value, traceback):
await cls.async_close()

def __getitem__(cls, key: str) -> 'OpenAIClient':
def __getitem__(cls, key: Union[int, str]) -> 'OpenAIClient':
"""
Returns the OpenAI API Client.
"""
if cls.enable_rotating_clients:
return cls.apis[key]
if isinstance(key, int):
key = cls.client_names[key]
return cls._clients[key]

@property
Expand Down

0 comments on commit f60033e

Please sign in to comment.