"""ReflACT Model backend — Azure OpenAI wrapper with token tracking. Provides optimizer/target dual-deployment chat functions and a global TokenTracker for per-stage cost accounting. Previously llm/azure_openai.py. """ from __future__ import annotations import json import os import subprocess import threading import time from types import SimpleNamespace from typing import Any from openai import AzureOpenAI, OpenAI # Sentinel value used as the api_version when the "openai_compatible" # auth_mode is selected. Real Azure deployments never use this string, # so it doubles as a marker for downstream type narrowing. _OPENAI_COMPATIBLE_API_VERSION = "openai-compat" # ── Configuration ───────────────────────────────────────────────────────────── ENDPOINT = os.environ.get( "AZURE_OPENAI_ENDPOINT", "", # Set via env var or config: e.g. "https://your-resource.openai.azure.com/" ) API_VERSION = os.environ.get("AZURE_OPENAI_API_VERSION", "2024-12-01-preview") API_KEY = os.environ.get( "AZURE_OPENAI_API_KEY", "", ) AUTH_MODE = os.environ.get("AZURE_OPENAI_AUTH_MODE", "azure_cli").strip().lower() AD_SCOPE = os.environ.get( "AZURE_OPENAI_AD_SCOPE", "https://cognitiveservices.azure.com/.default", ) MANAGED_IDENTITY_CLIENT_ID = os.environ.get( "AZURE_OPENAI_MANAGED_IDENTITY_CLIENT_ID", "", ).strip() OPTIMIZER_ENDPOINT = ( os.environ.get("OPTIMIZER_AZURE_OPENAI_ENDPOINT") or os.environ.get("AZURE_OPENAI_OPTIMIZER_ENDPOINT") or ENDPOINT ) TARGET_ENDPOINT = ( os.environ.get("TARGET_AZURE_OPENAI_ENDPOINT") or os.environ.get("AZURE_OPENAI_TARGET_ENDPOINT") or ENDPOINT ) OPTIMIZER_API_VERSION = ( os.environ.get("OPTIMIZER_AZURE_OPENAI_API_VERSION") or os.environ.get("AZURE_OPENAI_OPTIMIZER_API_VERSION") or API_VERSION ) TARGET_API_VERSION = ( os.environ.get("TARGET_AZURE_OPENAI_API_VERSION") or os.environ.get("AZURE_OPENAI_TARGET_API_VERSION") or API_VERSION ) OPTIMIZER_API_KEY = ( os.environ.get("OPTIMIZER_AZURE_OPENAI_API_KEY") or os.environ.get("AZURE_OPENAI_OPTIMIZER_API_KEY") or API_KEY ) TARGET_API_KEY = ( os.environ.get("TARGET_AZURE_OPENAI_API_KEY") or os.environ.get("AZURE_OPENAI_TARGET_API_KEY") or API_KEY ) OPTIMIZER_AUTH_MODE = ( os.environ.get("OPTIMIZER_AZURE_OPENAI_AUTH_MODE") or os.environ.get("AZURE_OPENAI_OPTIMIZER_AUTH_MODE") or AUTH_MODE ).strip().lower() TARGET_AUTH_MODE = ( os.environ.get("TARGET_AZURE_OPENAI_AUTH_MODE") or os.environ.get("AZURE_OPENAI_TARGET_AUTH_MODE") or AUTH_MODE ).strip().lower() OPTIMIZER_AD_SCOPE = ( os.environ.get("OPTIMIZER_AZURE_OPENAI_AD_SCOPE") or os.environ.get("AZURE_OPENAI_OPTIMIZER_AD_SCOPE") or AD_SCOPE ) TARGET_AD_SCOPE = ( os.environ.get("TARGET_AZURE_OPENAI_AD_SCOPE") or os.environ.get("AZURE_OPENAI_TARGET_AD_SCOPE") or AD_SCOPE ) OPTIMIZER_MANAGED_IDENTITY_CLIENT_ID = ( os.environ.get("OPTIMIZER_AZURE_OPENAI_MANAGED_IDENTITY_CLIENT_ID") or os.environ.get("AZURE_OPENAI_OPTIMIZER_MANAGED_IDENTITY_CLIENT_ID") or MANAGED_IDENTITY_CLIENT_ID ).strip() TARGET_MANAGED_IDENTITY_CLIENT_ID = ( os.environ.get("TARGET_AZURE_OPENAI_MANAGED_IDENTITY_CLIENT_ID") or os.environ.get("AZURE_OPENAI_TARGET_MANAGED_IDENTITY_CLIENT_ID") or MANAGED_IDENTITY_CLIENT_ID ).strip() OPTIMIZER_DEPLOYMENT = os.environ.get("OPTIMIZER_DEPLOYMENT", "gpt-4o") TARGET_DEPLOYMENT = os.environ.get("TARGET_DEPLOYMENT", "gpt-4o") REASONING_EFFORT: str | None = None _AZ_CLI_TOKEN_CACHE: dict[str, dict[str, Any]] = {} # Deployments that require Responses API _RESPONSES_API_MODELS = { "gpt-5.3-codex", "gpt-5.1-codex", "gpt-5.2-codex", "gpt-5-codex", "codex-mini", "gpt-5.4-pro", } # ── Token Tracker ───────────────────────────────────────────────────────────── class TokenTracker: """Thread-safe per-stage token counter.""" def __init__(self) -> None: self._lock = threading.Lock() self._data: dict[str, dict] = {} def record( self, stage: str, prompt_tokens: int, completion_tokens: int, ) -> None: with self._lock: if stage not in self._data: self._data[stage] = { "calls": 0, "prompt_tokens": 0, "completion_tokens": 0, } d = self._data[stage] d["calls"] += 1 d["prompt_tokens"] += prompt_tokens d["completion_tokens"] += completion_tokens def summary(self) -> dict: with self._lock: out: dict = {} total_p = total_c = total_calls = 0 for stage, d in sorted(self._data.items()): out[stage] = { "calls": d["calls"], "prompt_tokens": d["prompt_tokens"], "completion_tokens": d["completion_tokens"], "total_tokens": d["prompt_tokens"] + d["completion_tokens"], } total_p += d["prompt_tokens"] total_c += d["completion_tokens"] total_calls += d["calls"] out["_total"] = { "calls": total_calls, "prompt_tokens": total_p, "completion_tokens": total_c, "total_tokens": total_p + total_c, } return out def reset(self) -> None: with self._lock: self._data.clear() def stage_snapshot(self, stage: str) -> dict: """Return a copy of one stage's counters (or zeros if not tracked yet).""" with self._lock: d = self._data.get(stage, {}) return { "calls": d.get("calls", 0), "prompt_tokens": d.get("prompt_tokens", 0), "completion_tokens": d.get("completion_tokens", 0), "total_tokens": d.get("prompt_tokens", 0) + d.get("completion_tokens", 0), } tracker = TokenTracker() # ── Client management ───────────────────────────────────────────────────────── _optimizer_client: AzureOpenAI | OpenAI | None = None _target_client: AzureOpenAI | OpenAI | None = None _optimizer_lock = threading.Lock() _target_lock = threading.Lock() def _role_config(role: str) -> dict[str, str]: if role == "optimizer": return { "endpoint": OPTIMIZER_ENDPOINT, "api_version": OPTIMIZER_API_VERSION, "api_key": OPTIMIZER_API_KEY, "auth_mode": OPTIMIZER_AUTH_MODE, "ad_scope": OPTIMIZER_AD_SCOPE, "managed_identity_client_id": OPTIMIZER_MANAGED_IDENTITY_CLIENT_ID, } if role == "target": return { "endpoint": TARGET_ENDPOINT, "api_version": TARGET_API_VERSION, "api_key": TARGET_API_KEY, "auth_mode": TARGET_AUTH_MODE, "ad_scope": TARGET_AD_SCOPE, "managed_identity_client_id": TARGET_MANAGED_IDENTITY_CLIENT_ID, } raise ValueError(f"Unknown Azure OpenAI client role: {role!r}") def _make_token_provider( auth_mode: str, ad_scope: str, managed_identity_client_id: str, ): try: from azure.identity import ( # type: ignore[import-not-found] AzureCliCredential, ManagedIdentityCredential, get_bearer_token_provider, ) except ImportError as e: if auth_mode == "azure_cli": return _make_azure_cli_token_provider(ad_scope) raise ImportError( "Azure AD auth requires azure-identity. Install it with `pip install azure-identity`." ) from e if auth_mode in {"managed_identity", "aad", "azure_ad"}: if managed_identity_client_id: credential = ManagedIdentityCredential(client_id=managed_identity_client_id) else: credential = ManagedIdentityCredential() elif auth_mode == "azure_cli": credential = AzureCliCredential() else: raise ValueError( "Unsupported Azure OpenAI auth mode " f"{auth_mode!r}; expected api_key, managed_identity, azure_ad, aad, or azure_cli." ) return get_bearer_token_provider(credential, ad_scope) def _make_azure_cli_token_provider(ad_scope: str): """Return an Azure CLI token provider compatible with AzureOpenAI. This fallback avoids requiring azure-identity in environments where `az` is already logged in. The SDK calls this provider whenever it needs a bearer token. """ resource = ad_scope.removesuffix("/.default") def _provider() -> str: now = int(time.time()) cache = _AZ_CLI_TOKEN_CACHE.setdefault(resource, {"token": "", "expires_on": 0}) cached = str(cache.get("token") or "") expires_on = int(cache.get("expires_on") or 0) if cached and expires_on - now > 300: return cached raw = subprocess.check_output( [ "az", "account", "get-access-token", "--resource", resource, "-o", "json", ], text=True, stderr=subprocess.STDOUT, ) payload = json.loads(raw) token = str(payload["accessToken"]) cache["token"] = token cache["expires_on"] = int(payload.get("expires_on") or now + 3000) return token return _provider def _make_client(role: str) -> AzureOpenAI | OpenAI: cfg = _role_config(role) if not cfg["endpoint"]: raise ValueError( f"Azure OpenAI endpoint is not configured for {role}. " "Pass --azure_openai_endpoint https://your-resource.openai.azure.com/ " "or set AZURE_OPENAI_ENDPOINT in your environment." ) auth_mode = cfg["auth_mode"] if auth_mode in {"openai_compatible", "compat", "openai"}: return OpenAI( base_url=cfg["endpoint"].rstrip("/"), api_key=cfg["api_key"] or "dummy", default_headers={"User-Agent": "SkillOpt"}, ) if auth_mode in {"api_key", "key"}: if not cfg["api_key"]: raise ValueError( f"Azure OpenAI API key is not configured for {role}. " "Set model.azure_openai_api_key in the config or export AZURE_OPENAI_API_KEY." ) return AzureOpenAI( api_version=cfg["api_version"], azure_endpoint=cfg["endpoint"], api_key=cfg["api_key"], ) return AzureOpenAI( api_version=cfg["api_version"], azure_endpoint=cfg["endpoint"], azure_ad_token_provider=_make_token_provider( auth_mode, cfg["ad_scope"], cfg["managed_identity_client_id"], ), ) def get_optimizer_client() -> AzureOpenAI | OpenAI: global _optimizer_client with _optimizer_lock: if _optimizer_client is None: _optimizer_client = _make_client("optimizer") return _optimizer_client def get_target_client() -> AzureOpenAI | OpenAI: global _target_client with _target_lock: if _target_client is None: # When using qwen_chat backend, return an OpenAI client pointing to vLLM from skillopt.model.backend_config import get_target_backend if get_target_backend() == "qwen_chat": from skillopt.model import qwen_backend as _qwen target_config = _qwen.TARGET_CONFIG _target_client = OpenAI( base_url=target_config.base_url, api_key=target_config.api_key or "dummy", ) else: _target_client = _make_client("target") return _target_client def _needs_responses_api(deployment: str) -> bool: dep = deployment.lower() return any(dep == m or dep.startswith(m + "-") for m in _RESPONSES_API_MODELS) # ── Core chat function ──────────────────────────────────────────────────────── def _chat_impl( client: AzureOpenAI | OpenAI, deployment: str, system: str, user: str, max_completion_tokens: int, retries: int, stage: str, reasoning_effort: str | None = None, timeout: int | None = None, ) -> tuple[str, dict]: """Call LLM, track tokens, return (text, usage_dict).""" last_err = None usage_info = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} for attempt in range(retries): try: if _needs_responses_api(deployment): kwargs: dict[str, Any] = { "model": deployment, "instructions": system, "input": [{"role": "user", "content": user}], "max_output_tokens": max_completion_tokens, } actual_effort = reasoning_effort or REASONING_EFFORT if actual_effort: kwargs["reasoning"] = {"effort": actual_effort} if timeout is not None: kwargs["timeout"] = timeout resp = client.responses.create(**kwargs) text = getattr(resp, "output_text", None) or "" if not text: for item in getattr(resp, "output", None) or []: for part in getattr(item, "content", []): if getattr(part, "type", "") == "output_text": text = part.text or "" if hasattr(resp, "usage") and resp.usage: usage_info = { "prompt_tokens": getattr(resp.usage, "input_tokens", 0) or 0, "completion_tokens": getattr(resp.usage, "output_tokens", 0) or 0, "total_tokens": ( (getattr(resp.usage, "input_tokens", 0) or 0) + (getattr(resp.usage, "output_tokens", 0) or 0) ), } else: kwargs: dict[str, Any] = dict( model=deployment, messages=[ {"role": "system", "content": system}, {"role": "user", "content": user}, ], max_completion_tokens=max_completion_tokens, ) actual_effort = reasoning_effort or REASONING_EFFORT if actual_effort is not None: kwargs["reasoning_effort"] = actual_effort if timeout is not None: kwargs["timeout"] = timeout resp = client.chat.completions.create(**kwargs) text = resp.choices[0].message.content or "" if resp.usage: usage_info = { "prompt_tokens": resp.usage.prompt_tokens or 0, "completion_tokens": resp.usage.completion_tokens or 0, "total_tokens": resp.usage.total_tokens or 0, } tracker.record( stage, usage_info["prompt_tokens"], usage_info["completion_tokens"], ) return text, usage_info except Exception as e: # noqa: BLE001 last_err = e sleep = min(2 ** attempt, 30) time.sleep(sleep) raise RuntimeError(f"LLM call failed after {retries} retries: {last_err}") def _chat_messages_impl( client: AzureOpenAI | OpenAI, deployment: str, messages: list[dict[str, Any]], max_completion_tokens: int, retries: int, stage: str, reasoning_effort: str | None = None, *, tools: list[dict[str, Any]] | None = None, tool_choice: str | dict[str, Any] | None = None, return_message: bool = False, timeout: int | None = None, ) -> tuple[Any, dict]: """Call the model with a pre-built message list.""" last_err = None usage_info = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} for attempt in range(retries): try: if _needs_responses_api(deployment): input_items, instructions = _messages_to_responses_input(messages) kwargs: dict[str, Any] = { "model": deployment, "input": input_items, "max_output_tokens": max_completion_tokens, } if instructions: kwargs["instructions"] = instructions actual_effort = reasoning_effort or REASONING_EFFORT if actual_effort: kwargs["reasoning"] = {"effort": actual_effort} if tools: kwargs["tools"] = [_chat_tool_to_responses_tool(tool) for tool in tools] if tool_choice is not None: kwargs["tool_choice"] = tool_choice if timeout is not None: kwargs["timeout"] = timeout resp = client.responses.create(**kwargs) message, text = _responses_to_chat_message(resp) if hasattr(resp, "usage") and resp.usage: usage_info = { "prompt_tokens": getattr(resp.usage, "input_tokens", 0) or 0, "completion_tokens": getattr(resp.usage, "output_tokens", 0) or 0, "total_tokens": ( (getattr(resp.usage, "input_tokens", 0) or 0) + (getattr(resp.usage, "output_tokens", 0) or 0) ), } else: kwargs = dict( model=deployment, messages=messages, max_completion_tokens=max_completion_tokens, ) actual_effort = reasoning_effort or REASONING_EFFORT if tools: kwargs["tools"] = tools if tool_choice is not None: kwargs["tool_choice"] = tool_choice # Some models (e.g. gpt-5.5) don't support reasoning_effort with function tools elif actual_effort is not None: kwargs["reasoning_effort"] = actual_effort if timeout is not None: kwargs["timeout"] = timeout resp = client.chat.completions.create(**kwargs) message = resp.choices[0].message text = message.content or "" if resp.usage: usage_info = { "prompt_tokens": resp.usage.prompt_tokens or 0, "completion_tokens": resp.usage.completion_tokens or 0, "total_tokens": resp.usage.total_tokens or 0, } tracker.record( stage, usage_info["prompt_tokens"], usage_info["completion_tokens"], ) return (message if return_message else text), usage_info except Exception as e: # noqa: BLE001 last_err = e sleep = min(2 ** attempt, 30) time.sleep(sleep) raise RuntimeError(f"LLM message call failed after {retries} retries: {last_err}") def _chat_tool_to_responses_tool(tool: dict[str, Any]) -> dict[str, Any]: """Convert a Chat Completions function tool to Responses API format.""" if tool.get("type") == "function" and isinstance(tool.get("function"), dict): fn = tool["function"] return { "type": "function", "name": fn.get("name", ""), "description": fn.get("description", ""), "parameters": fn.get("parameters", {"type": "object", "properties": {}}), } return tool def _messages_to_responses_input(messages: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], str]: """Convert chat-style messages, including tool results, to Responses input.""" instructions: list[str] = [] input_items: list[dict[str, Any]] = [] for message in messages: role = message.get("role") content = message.get("content") or "" if role == "system": if content: instructions.append(str(content)) continue if role == "tool": input_items.append({ "type": "function_call_output", "call_id": str(message.get("tool_call_id", "")), "output": str(content), }) continue if role == "assistant": if content: input_items.append({"role": "assistant", "content": str(content)}) for tool_call in message.get("tool_calls") or []: function = tool_call.get("function", {}) or {} input_items.append({ "type": "function_call", "call_id": str(tool_call.get("id", "")), "name": str(function.get("name", "")), "arguments": str(function.get("arguments", "{}") or "{}"), }) continue if role in {"user", "developer"}: input_items.append({"role": "user", "content": str(content)}) return input_items, "\n\n".join(instructions) def _responses_to_chat_message(resp: Any) -> tuple[Any, str]: """Convert Responses output into the subset of Chat message API we use.""" text = getattr(resp, "output_text", None) or "" tool_calls: list[dict[str, Any]] = [] for item in getattr(resp, "output", None) or []: item_type = getattr(item, "type", "") if item_type == "function_call": tool_calls.append({ "id": getattr(item, "call_id", "") or getattr(item, "id", ""), "type": "function", "function": { "name": getattr(item, "name", ""), "arguments": getattr(item, "arguments", "") or "{}", }, }) elif item_type == "message" and not text: content_parts = getattr(item, "content", []) or [] for part in content_parts: if getattr(part, "type", "") == "output_text": text += getattr(part, "text", "") or "" return SimpleNamespace(content=text, tool_calls=tool_calls), text # ── Public API ──────────────────────────────────────────────────────────────── def configure_azure_openai( *, endpoint: str | None = None, api_version: str | None = None, api_key: str | None = None, auth_mode: str | None = None, ad_scope: str | None = None, managed_identity_client_id: str | None = None, optimizer_endpoint: str | None = None, optimizer_api_version: str | None = None, optimizer_api_key: str | None = None, optimizer_auth_mode: str | None = None, optimizer_ad_scope: str | None = None, optimizer_managed_identity_client_id: str | None = None, target_endpoint: str | None = None, target_api_version: str | None = None, target_api_key: str | None = None, target_auth_mode: str | None = None, target_ad_scope: str | None = None, target_managed_identity_client_id: str | None = None, ) -> None: global ENDPOINT, API_VERSION, API_KEY, AUTH_MODE, AD_SCOPE, MANAGED_IDENTITY_CLIENT_ID global OPTIMIZER_ENDPOINT, OPTIMIZER_API_VERSION, OPTIMIZER_API_KEY, OPTIMIZER_AUTH_MODE global OPTIMIZER_AD_SCOPE, OPTIMIZER_MANAGED_IDENTITY_CLIENT_ID global TARGET_ENDPOINT, TARGET_API_VERSION, TARGET_API_KEY, TARGET_AUTH_MODE global TARGET_AD_SCOPE, TARGET_MANAGED_IDENTITY_CLIENT_ID global _optimizer_client, _target_client def _clean(value: str | None, *, lower: bool = False) -> str | None: if value is None: return None str_value = str(value).strip() if not str_value: return None if lower: str_value = str_value.lower() return str_value def _set(global_name: str, value: str | None, env_key: str) -> None: if value is None: return globals()[global_name] = value os.environ[env_key] = value shared_endpoint = _clean(endpoint) shared_api_version = _clean(api_version) shared_api_key = _clean(api_key) shared_auth_mode = _clean(auth_mode, lower=True) shared_ad_scope = _clean(ad_scope) shared_managed_identity_client_id = _clean(managed_identity_client_id) # Auto-configure for openai_compatible mode if shared_auth_mode in {"openai_compatible", "compat", "openai"}: if shared_api_version is None: shared_api_version = _OPENAI_COMPATIBLE_API_VERSION _set("ENDPOINT", shared_endpoint, "AZURE_OPENAI_ENDPOINT") _set("API_VERSION", shared_api_version, "AZURE_OPENAI_API_VERSION") _set("API_KEY", shared_api_key, "AZURE_OPENAI_API_KEY") _set("AUTH_MODE", shared_auth_mode, "AZURE_OPENAI_AUTH_MODE") _set("AD_SCOPE", shared_ad_scope, "AZURE_OPENAI_AD_SCOPE") _set( "MANAGED_IDENTITY_CLIENT_ID", shared_managed_identity_client_id, "AZURE_OPENAI_MANAGED_IDENTITY_CLIENT_ID", ) resolved_optimizer_endpoint = _clean(optimizer_endpoint) or shared_endpoint resolved_optimizer_api_version = _clean(optimizer_api_version) or shared_api_version resolved_optimizer_api_key = _clean(optimizer_api_key) or shared_api_key resolved_optimizer_auth_mode = _clean(optimizer_auth_mode, lower=True) or shared_auth_mode resolved_optimizer_ad_scope = _clean(optimizer_ad_scope) or shared_ad_scope resolved_optimizer_mi = ( _clean(optimizer_managed_identity_client_id) or shared_managed_identity_client_id ) # Auto-configure for openai_compatible mode if resolved_optimizer_auth_mode in {"openai_compatible", "compat", "openai"}: if resolved_optimizer_api_version is None: resolved_optimizer_api_version = _OPENAI_COMPATIBLE_API_VERSION resolved_target_endpoint = _clean(target_endpoint) or shared_endpoint resolved_target_api_version = _clean(target_api_version) or shared_api_version resolved_target_api_key = _clean(target_api_key) or shared_api_key resolved_target_auth_mode = _clean(target_auth_mode, lower=True) or shared_auth_mode resolved_target_ad_scope = _clean(target_ad_scope) or shared_ad_scope resolved_target_mi = ( _clean(target_managed_identity_client_id) or shared_managed_identity_client_id ) # Auto-configure for openai_compatible mode if resolved_target_auth_mode in {"openai_compatible", "compat", "openai"}: if resolved_target_api_version is None: resolved_target_api_version = _OPENAI_COMPATIBLE_API_VERSION _set("OPTIMIZER_ENDPOINT", resolved_optimizer_endpoint, "OPTIMIZER_AZURE_OPENAI_ENDPOINT") _set( "OPTIMIZER_API_VERSION", resolved_optimizer_api_version, "OPTIMIZER_AZURE_OPENAI_API_VERSION", ) _set("OPTIMIZER_API_KEY", resolved_optimizer_api_key, "OPTIMIZER_AZURE_OPENAI_API_KEY") _set("OPTIMIZER_AUTH_MODE", resolved_optimizer_auth_mode, "OPTIMIZER_AZURE_OPENAI_AUTH_MODE") _set("OPTIMIZER_AD_SCOPE", resolved_optimizer_ad_scope, "OPTIMIZER_AZURE_OPENAI_AD_SCOPE") _set( "OPTIMIZER_MANAGED_IDENTITY_CLIENT_ID", resolved_optimizer_mi, "OPTIMIZER_AZURE_OPENAI_MANAGED_IDENTITY_CLIENT_ID", ) _set("TARGET_ENDPOINT", resolved_target_endpoint, "TARGET_AZURE_OPENAI_ENDPOINT") _set( "TARGET_API_VERSION", resolved_target_api_version, "TARGET_AZURE_OPENAI_API_VERSION", ) _set("TARGET_API_KEY", resolved_target_api_key, "TARGET_AZURE_OPENAI_API_KEY") _set("TARGET_AUTH_MODE", resolved_target_auth_mode, "TARGET_AZURE_OPENAI_AUTH_MODE") _set("TARGET_AD_SCOPE", resolved_target_ad_scope, "TARGET_AZURE_OPENAI_AD_SCOPE") _set( "TARGET_MANAGED_IDENTITY_CLIENT_ID", resolved_target_mi, "TARGET_AZURE_OPENAI_MANAGED_IDENTITY_CLIENT_ID", ) with _optimizer_lock: _optimizer_client = None with _target_lock: _target_client = None def chat_optimizer( system: str, user: str, max_completion_tokens: int = 16384, retries: int = 5, stage: str = "optimizer", reasoning_effort: str | None = None, timeout: int | None = None, ) -> tuple[str, dict]: """Call the optimizer model. Returns (response_text, usage_dict).""" return _chat_impl( get_optimizer_client(), OPTIMIZER_DEPLOYMENT, system, user, max_completion_tokens, retries, stage, reasoning_effort, timeout, ) def chat_with_deployment( deployment: str, system: str, user: str, max_completion_tokens: int = 16384, retries: int = 5, stage: str = "custom", reasoning_effort: str | None = None, timeout: int | None = None, ) -> tuple[str, dict]: """Call an arbitrary deployment using the shared Azure client.""" return _chat_impl( get_optimizer_client(), deployment, system, user, max_completion_tokens, retries, stage, reasoning_effort, timeout, ) def chat_target( system: str, user: str, max_completion_tokens: int = 16384, retries: int = 5, stage: str = "target", reasoning_effort: str | None = None, timeout: int | None = None, ) -> tuple[str, dict]: """Call the target model. Returns (response_text, usage_dict).""" return _chat_impl( get_target_client(), TARGET_DEPLOYMENT, system, user, max_completion_tokens, retries, stage, reasoning_effort, timeout, ) def chat_optimizer_messages( messages: list[dict[str, Any]], max_completion_tokens: int = 16384, retries: int = 5, stage: str = "optimizer", reasoning_effort: str | None = None, *, tools: list[dict[str, Any]] | None = None, tool_choice: str | dict[str, Any] | None = None, return_message: bool = False, timeout: int | None = None, ) -> tuple[Any, dict]: """Call the optimizer model with a pre-built chat message list.""" return _chat_messages_impl( get_optimizer_client(), OPTIMIZER_DEPLOYMENT, messages, max_completion_tokens, retries, stage, reasoning_effort, tools=tools, tool_choice=tool_choice, return_message=return_message, timeout=timeout, ) def chat_messages_with_deployment( deployment: str, messages: list[dict[str, Any]], max_completion_tokens: int = 16384, retries: int = 5, stage: str = "custom", reasoning_effort: str | None = None, *, tools: list[dict[str, Any]] | None = None, tool_choice: str | dict[str, Any] | None = None, return_message: bool = False, timeout: int | None = None, ) -> tuple[Any, dict]: """Call an arbitrary deployment with a pre-built chat message list.""" return _chat_messages_impl( get_optimizer_client(), deployment, messages, max_completion_tokens, retries, stage, reasoning_effort, tools=tools, tool_choice=tool_choice, return_message=return_message, timeout=timeout, ) def chat_target_messages( messages: list[dict[str, Any]], max_completion_tokens: int = 16384, retries: int = 5, stage: str = "target", reasoning_effort: str | None = None, *, tools: list[dict[str, Any]] | None = None, tool_choice: str | dict[str, Any] | None = None, return_message: bool = False, timeout: int | None = None, ) -> tuple[Any, dict]: """Call the target model with a pre-built chat message list.""" return _chat_messages_impl( get_target_client(), TARGET_DEPLOYMENT, messages, max_completion_tokens, retries, stage, reasoning_effort, tools=tools, tool_choice=tool_choice, return_message=return_message, timeout=timeout, ) def get_token_summary() -> dict: """Return per-stage and total token usage.""" return tracker.summary() def reset_token_tracker() -> None: tracker.reset() def set_target_deployment(deployment: str) -> None: """Change target deployment at runtime.""" global _target_client, TARGET_DEPLOYMENT TARGET_DEPLOYMENT = deployment os.environ["TARGET_DEPLOYMENT"] = deployment os.environ["AZURE_OPENAI_DEPLOYMENT"] = deployment with _target_lock: _target_client = None try: import llm_client as _legacy _legacy.DEPLOYMENT = deployment _legacy._client = None except Exception: pass def set_reasoning_effort(effort: str | None) -> None: """Set reasoning effort for all LLM calls. None = off.""" global REASONING_EFFORT REASONING_EFFORT = effort if effort else None def get_reasoning_effort() -> str | None: """Return the process-wide reasoning effort for direct Azure client users.""" return REASONING_EFFORT def set_optimizer_deployment(deployment: str) -> None: """Change optimizer deployment at runtime.""" global _optimizer_client, OPTIMIZER_DEPLOYMENT OPTIMIZER_DEPLOYMENT = deployment os.environ["OPTIMIZER_DEPLOYMENT"] = deployment with _optimizer_lock: _optimizer_client = None