mirror of
https://github.com/microsoft/SkillOpt.git
synced 2026-07-07 00:15:39 +08:00
916 lines
33 KiB
Python
916 lines
33 KiB
Python
"""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
|