self-evolve-framework 1.6.0 → 1.6.1

package/template/skills/skillopt-sleep/scripts/framework/skillopt/model/azure_openai.py

@@ -1,915 +0,0 @@
-"""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