From 29a36855d955a24b8333f65b3089b7b4996952bb Mon Sep 17 00:00:00 2001 From: harshitkd Date: Mon, 5 Jan 2026 18:24:49 +0530 Subject: [PATCH] Add Azure OpenAI backend support - Create AzureOpenAIClient class in rlm/clients/azure_openai.py - Register azure_openai backend in rlm/clients/__init__.py - Add azure_openai to ClientBackend type in rlm/core/types.py This allows users to use Azure OpenAI as a backend by providing azure_endpoint, api_key, api_version, and optionally azure_deployment. --- rlm/clients/__init__.py | 6 +- rlm/clients/azure_openai.py | 142 ++++++++++++++++++++++++++++++++++++ rlm/core/types.py | 2 +- 3 files changed, 148 insertions(+), 2 deletions(-) create mode 100644 rlm/clients/azure_openai.py diff --git a/rlm/clients/__init__.py b/rlm/clients/__init__.py index 716062680..d07b05a5e 100644 --- a/rlm/clients/__init__.py +++ b/rlm/clients/__init__.py @@ -44,7 +44,11 @@ def get_client( from rlm.clients.anthropic import AnthropicClient return AnthropicClient(**backend_kwargs) + elif backend == "azure_openai": + from rlm.clients.azure_openai import AzureOpenAIClient + + return AzureOpenAIClient(**backend_kwargs) else: raise ValueError( - f"Unknown backend: {backend}. Supported backends: ['openai', 'vllm', 'portkey', 'openrouter', 'litellm', 'anthropic']" + f"Unknown backend: {backend}. Supported backends: ['openai', 'vllm', 'portkey', 'openrouter', 'litellm', 'anthropic', 'azure_openai']" ) diff --git a/rlm/clients/azure_openai.py b/rlm/clients/azure_openai.py new file mode 100644 index 000000000..c5207ba54 --- /dev/null +++ b/rlm/clients/azure_openai.py @@ -0,0 +1,142 @@ +import os +from collections import defaultdict +from typing import Any + +import openai +from dotenv import load_dotenv + +from rlm.clients.base_lm import BaseLM +from rlm.core.types import ModelUsageSummary, UsageSummary + +load_dotenv() + +# Load API key from environment variable +DEFAULT_AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY") + + +class AzureOpenAIClient(BaseLM): + """ + LM Client for running models with the Azure OpenAI API. + """ + + def __init__( + self, + api_key: str | None = None, + model_name: str | None = None, + azure_endpoint: str | None = None, + api_version: str | None = None, + azure_deployment: str | None = None, + **kwargs, + ): + super().__init__(model_name=model_name, **kwargs) + + if api_key is None: + api_key = DEFAULT_AZURE_OPENAI_API_KEY + + if azure_endpoint is None: + azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT") + + if api_version is None: + api_version = os.getenv("AZURE_OPENAI_API_VERSION", "2024-02-01") + + if azure_endpoint is None: + raise ValueError( + "azure_endpoint is required for Azure OpenAI client. " + "Set it via argument or AZURE_OPENAI_ENDPOINT environment variable." + ) + + self.client = openai.AzureOpenAI( + api_key=api_key, + azure_endpoint=azure_endpoint, + api_version=api_version, + azure_deployment=azure_deployment, + ) + self.async_client = openai.AsyncAzureOpenAI( + api_key=api_key, + azure_endpoint=azure_endpoint, + api_version=api_version, + azure_deployment=azure_deployment, + ) + self.model_name = model_name + self.azure_deployment = azure_deployment + + # Per-model usage tracking + self.model_call_counts: dict[str, int] = defaultdict(int) + self.model_input_tokens: dict[str, int] = defaultdict(int) + self.model_output_tokens: dict[str, int] = defaultdict(int) + self.model_total_tokens: dict[str, int] = defaultdict(int) + + def completion( + self, prompt: str | list[dict[str, Any]], model: str | None = None + ) -> str: + if isinstance(prompt, str): + messages = [{"role": "user", "content": prompt}] + elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt): + messages = prompt + else: + raise ValueError(f"Invalid prompt type: {type(prompt)}") + + model = model or self.model_name + if not model: + raise ValueError("Model name is required for Azure OpenAI client.") + + response = self.client.chat.completions.create( + model=model, + messages=messages, + ) + self._track_cost(response, model) + return response.choices[0].message.content + + async def acompletion( + self, prompt: str | list[dict[str, Any]], model: str | None = None + ) -> str: + if isinstance(prompt, str): + messages = [{"role": "user", "content": prompt}] + elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt): + messages = prompt + else: + raise ValueError(f"Invalid prompt type: {type(prompt)}") + + model = model or self.model_name + if not model: + raise ValueError("Model name is required for Azure OpenAI client.") + + response = await self.async_client.chat.completions.create( + model=model, + messages=messages, + ) + self._track_cost(response, model) + return response.choices[0].message.content + + def _track_cost(self, response: openai.ChatCompletion, model: str): + self.model_call_counts[model] += 1 + + usage = getattr(response, "usage", None) + if usage is None: + raise ValueError("No usage data received. Tracking tokens not possible.") + + self.model_input_tokens[model] += usage.prompt_tokens + self.model_output_tokens[model] += usage.completion_tokens + self.model_total_tokens[model] += usage.total_tokens + + # Track last call for handler to read + self.last_prompt_tokens = usage.prompt_tokens + self.last_completion_tokens = usage.completion_tokens + + def get_usage_summary(self) -> UsageSummary: + model_summaries = {} + for model in self.model_call_counts: + model_summaries[model] = ModelUsageSummary( + total_calls=self.model_call_counts[model], + total_input_tokens=self.model_input_tokens[model], + total_output_tokens=self.model_output_tokens[model], + ) + return UsageSummary(model_usage_summaries=model_summaries) + + def get_last_usage(self) -> ModelUsageSummary: + return ModelUsageSummary( + total_calls=1, + total_input_tokens=self.last_prompt_tokens, + total_output_tokens=self.last_completion_tokens, + ) + diff --git a/rlm/core/types.py b/rlm/core/types.py index ad1676284..3c485b08b 100644 --- a/rlm/core/types.py +++ b/rlm/core/types.py @@ -2,7 +2,7 @@ from types import ModuleType from typing import Any, Literal -ClientBackend = Literal["openai", "portkey", "openrouter", "vllm", "litellm", "anthropic"] +ClientBackend = Literal["openai", "portkey", "openrouter", "vllm", "litellm", "anthropic", "azure_openai"] EnvironmentType = Literal["local", "prime", "modal"]