Skip to content

Commit

Permalink
Merge pull request #9 from mattalford/azure-openai-api
Browse files Browse the repository at this point in the history
Azure OpenAI API
  • Loading branch information
Holmeswww authored Apr 23, 2024
2 parents df43b1e + 4d0ab19 commit f7f3523
Show file tree
Hide file tree
Showing 3 changed files with 76 additions and 30 deletions.
4 changes: 3 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,8 @@ Currently, the built-in API supports OpenAI and Anthropic, see https://pypi.org/

To use the OpenAI models, set environment variables `OPENAI_KEY` and `OPENAI_ORG`. Alternatively, you can put the openai 'key' and 'organization' in the first 2 lines of `~/.openai/openai.key`.

To use the Azure OpenAI models, set environment variables `AZURE_OPENAI_API_KEY`, `AZURE_OPENAI_API_VERSION`, `AZURE_OPENAI_ENDPOINT`, and `AZURE_DEPLOYMENT_NAME`. Alternatively, you can store the Azure OpenAI API key, API version, Azure endpoint, and deployment name in the first 4 lines of `~/.openai/azure_openai.key`.

To use the Anthropic models, set environment variable `ANTHROPIC_KEY`. Alternatively, you can put the anthropic 'key' in 3rd line of `~/.openai/openai.key`.

# Using AgentKit without Programming Experience
Expand Down Expand Up @@ -156,4 +158,4 @@ To support advanced capabilities such as branching, AgentKit offers API to dynam
<source media="(prefers-color-scheme: light)" srcset="https://api.star-history.com/svg?repos=holmeswww/agentkit&type=Date" />
<img alt="Star History Chart" src="https://api.star-history.com/svg?repos=holmeswww/agentkit&type=Date" />
</picture>
</a>
</a>
4 changes: 3 additions & 1 deletion docs/source/getting started.rst
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,8 @@ Currently, the built-in API supports OpenAI and Anthropic, see https://pypi.org/

To use the OpenAI models, set environment variables ``OPENAI_KEY`` and ``OPENAI_ORG``. Alternatively, you can put the openai 'key' and 'organization' in the first 2 lines of ``~/.openai/openai.key``.

To use the Azure OpenAI models, set environment variables ``AZURE_OPENAI_API_KEY``, ``AZURE_OPENAI_API_VERSION``, ``AZURE_OPENAI_ENDPOINT``, and ``AZURE_DEPLOYMENT_NAME``. Alternatively, you can store the Azure OpenAI API key, API version, Azure endpoint, and deployment name in the first 4 lines of ``~/.openai/azure_openai.key``.

To use the Anthropic models, set environment variable ``ANTHROPIC_KEY``. Alternatively, you can put the anthropic 'key' in 3rd line of ``~/.openai/openai.key``.

``LLM_API_FUNCTION`` can be any LLM querying function that takes ``msg:list`` and ``shrink_idx:int``, and outputs ``llm_result:str`` and ``usage:dict``. Where ``msg`` is a prompt (`OpenAI format`_ by default), and ``shrink_idx:int`` is an index at which the LLM should reduce the length of the prompt in case of overflow.
Expand All @@ -95,4 +97,4 @@ AgentKit tracks token usage of each node through the ``LLM_API_FUNCTION`` with:
.. _OpenAI format: https://platform.openai.com/docs/guides/text-generation/chat-completions-api
.. _OpenAI format: https://platform.openai.com/docs/guides/text-generation/chat-completions-api
98 changes: 70 additions & 28 deletions src/agentkit/llm_api/GPT.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,29 +2,65 @@
import openai
except ImportError:
raise ImportError("Please install openai to use built-in LLM API.")
from openai import OpenAI
from openai import OpenAI, AzureOpenAI
import time
import os
from .utils import match_model
from .base import BaseModel

if os.environ.get("OPENAI_KEY") is None:
print("Environment variable for OpenAI key not found, using OpenAI API key from ~/.openai/openai.key.")
if not os.path.exists(os.path.join(os.path.expanduser('~'), ".openai/openai.key")):
raise FileNotFoundError("Please create a file at ~/.openai/openai.key with your OpenAI API key and organization ID. The first line should be your API key and the second line should be your organization ID.")
else:
with open(os.path.join(os.path.expanduser('~'), ".openai/openai.key"), 'r') as f:
org_key = f.readlines()
OpenAI_KEY = org_key[0].strip()
OpenAI_ORG = org_key[1].strip()
else:
print("Using OpenAI API key from environment variable.")
OpenAI_KEY = os.environ.get("OPENAI_KEY")
OpenAI_ORG = os.environ.get("OPENAI_ORG")
client = OpenAI(
api_key=OpenAI_KEY,
organization=OpenAI_ORG,
)
def initialize_client():
if os.environ.get("OPENAI_KEY"):
return OpenAI(
api_key=os.environ["OPENAI_KEY"],
organization=os.environ["OPENAI_ORG"],
)
elif os.environ.get("AZURE_OPENAI_API_KEY"):
return AzureOpenAI(
api_key=os.environ["AZURE_OPENAI_API_KEY"],
api_version=os.environ["AZURE_OPENAI_API_VERSION"],
azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
), os.environ.get("AZURE_DEPLOYMENT_NAME")

openai_key_path = os.path.join(os.path.expanduser('~'), ".openai", "openai.key")
azure_key_path = os.path.join(os.path.expanduser('~'), ".openai", "azure_openai.key")

if os.path.exists(openai_key_path):
with open(openai_key_path, 'r') as f:
lines = f.read().splitlines()
if len(lines) >= 2:
return OpenAI(
api_key=lines[0].strip(),
organization=lines[1].strip(),
)
raise FileNotFoundError(
"""
Please create a file at ~/.openai/openai.key with your OpenAI API key and organization ID.
The first line should be your API key and the second line should be your organization ID.
"""
)

if os.path.exists(azure_key_path):
with open(azure_key_path, 'r') as f:
lines = f.read().splitlines()
if len(lines) >= 3:
return AzureOpenAI(
api_key=lines[0].strip(),
api_version=lines[1].strip(),
azure_endpoint=lines[2].strip(),
), lines[3].strip()

raise FileNotFoundError(
"""
Please create a file at ~/.openai/azure_openai.key with your Azure OpenAI API key,
API version, Azure endpoint, and deployment name.
The first line should be your API key, the second line should be your API version,
the third line should be your Azure endpoint, and the fourth line should be your deployment name.
"""
)

client_and_deployment = initialize_client()
client = client_and_deployment[0] if isinstance(client_and_deployment, tuple) else client_and_deployment
deployment_name = client_and_deployment[1] if isinstance(client_and_deployment, tuple) else None

class GPT_chat(BaseModel):

Expand All @@ -45,16 +81,22 @@ def query_chat(self, messages, shrink_idx, max_gen=512, temp=0.):
messages = self.shrink_msg(messages, shrink_idx, model_max-max_gen)
while(True):
try:
completion = client.chat.completions.create(
model=self.name,
messages=messages,
temperature=temp,
max_tokens=max_gen,
)
if isinstance(client, OpenAI):
completion = client.chat.completions.create(
model=self.name,
messages=messages,
temperature=temp,
max_tokens=max_gen,
)
elif isinstance(client, AzureOpenAI):
completion = client.completions.create(
model=deployment_name,
messages=messages,
temperature=temp,
max_tokens=max_gen,
)
return completion.choices[0].message.content, {"prompt":completion.usage.prompt_tokens, "completion":completion.usage.completion_tokens, "total":completion.usage.total_tokens}
except openai.AuthenticationError as e:
raise e
except (openai.RateLimitError, openai.APIStatusError, openai.APITimeoutError, openai.APIConnectionError, openai.InternalServerError) as e:
except (openai.RateLimitError, openai.APIStatusError, openai.APITimeoutError, openai.APIConnectionError, openai.InternalServerError):
time.sleep(30)
except Exception as e:
e = str(e)
Expand All @@ -68,4 +110,4 @@ def query_chat(self, messages, shrink_idx, max_gen=512, temp=0.):
messages = self.shrink_msg_by(messages, shrink_idx, val-model_max)
else:
time.sleep(5)
print(e)
print(e)

0 comments on commit f7f3523

Please sign in to comment.