Skip to content

Commit

Permalink
feat: refactor OpenAI wrapper to use API client factory
Browse files Browse the repository at this point in the history
- Use APIClientFactory to get API client instance
- Set API key on global settings object
- Remove direct Client instantiation
  • Loading branch information
joehewett committed Jan 31, 2025
1 parent 2398927 commit 1bd5c80
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 6 deletions.
7 changes: 5 additions & 2 deletions src/asteroid_sdk/registration/initialise_project.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,11 @@ def asteroid_init(
api_key: Optional API key to override the default from environment variables
"""
if api_key:
# If the user provided an API key, override the client in settings
# 1) Set the key on the global settings object
logging.info("Overriding API key env variable with provided API key")
settings.api_key = api_key

# 2) Overwrite the API client directly
APIClientFactory._instance = Client(
base_url=settings.api_url,
headers={"X-Asteroid-Api-Key": api_key}
Expand All @@ -53,7 +56,7 @@ def asteroid_init(
supervision_config.set_execution_settings(execution_settings)

register_tools_and_supervisors_from_registry(run_id=run_id,
message_supervisors=message_supervisors)
message_supervisors=message_supervisors)

return run_id

Expand Down
9 changes: 5 additions & 4 deletions src/asteroid_sdk/wrappers/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@
from asteroid_sdk.supervision.helpers.openai_helper import OpenAiSupervisionHelper
from asteroid_sdk.api.generated.asteroid_api_client.api.run.get_run import sync as get_run_sync

from asteroid_sdk.registration.helper import APIClientFactory

# Conditionally import Langfuse if enabled
if settings.langfuse_enabled:
try:
Expand Down Expand Up @@ -288,10 +290,9 @@ def asteroid_openai_client(
raise ValueError("Invalid OpenAI client: missing chat attribute")

try:
client = Client(
base_url=settings.api_url,
headers={"X-Asteroid-Api-Key": f"{settings.api_key}"},
)
# Get the client from the factory
client = APIClientFactory.get_client()

supervision_manager = _create_supervision_manager(client)
openai_client.chat.completions = CompletionsWrapper(
openai_client.chat.completions, supervision_manager, run_id, execution_mode
Expand Down

0 comments on commit 1bd5c80

Please sign in to comment.