Add Azure OpenAI backend support

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']"
         )

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,
+        )
+

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"]