1"""
2common.py — Shared Azure AI Foundry authentication and client setup.
3 
4Supports three connection methods in order of preference:
51. /v1/ project endpoint (simplest, preferred)
62. Foundry SDK with DefaultAzureCredential (no API key needed, cloud-native)
73. Azure OpenAI endpoint (classic)
8 
9AAD tokens are auto-refreshed via azure.identity for long-running scripts
10(monitor_training.py, generate_distillation_data.py, etc.).
11 
12Usage:
13    from common import get_clients, upload_file
14 
15    # Method 1: Project /v1/ endpoint (preferred)
16    clients = get_clients(base_url="https://<resource>.services.ai.azure.com/api/projects/<project>/openai/v1/",
17                          api_key="KEY")
18 
19    # Method 2: Foundry SDK (DefaultAzureCredential — no API key needed)
20    clients = get_clients(project_endpoint="https://<resource>.services.ai.azure.com/api/projects/<project>")
21 
22    # Method 3: Azure OpenAI endpoint
23    clients = get_clients(azure_endpoint="https://<resource>.openai.azure.com",
24                          api_key="KEY")
25"""
26import argparse
27import os
28import sys
29 
30 
31 
32try:
33    sys.stdout.reconfigure(encoding="utf-8")
34    sys.stderr.reconfigure(encoding="utf-8")
35except (AttributeError, OSError):
36    pass  # Stream not reconfigurable (older Python or non-tty); default encoding is fine
37_AZURE_COGSERVICES_SCOPE = "https://cognitiveservices.azure.com/.default"
38 
39 
40def _clamp_score(v, default=0):
41    """Clamp a judge score to [1, 10]. Returns `default` for missing/non-numeric values.
42 
43    LLM judges occasionally return out-of-range integers (e.g., 15) or non-numeric
44    strings ("high"). Without clamping, these distort aggregate scores or crash
45    `int()`. We use 0 as a sentinel for "missing/failed" so callers can filter via
46    `score > 0`.
47    """
48    if v is None:
49        return default
50    try:
51        return max(1, min(10, int(v)))
52    except (ValueError, TypeError):
53        return default
54 
55 
56class HelpOnErrorParser(argparse.ArgumentParser):
57    """ArgumentParser that prints full help when arguments are invalid.
58    
59    Standard ArgumentParser only prints a one-line usage summary on error,
60    which isn't helpful for first-time users. This prints the full --help.
61    """
62 
63    def error(self, message):
64        self.print_help(sys.stderr)
65        self.exit(2, f"\nerror: {message}\n")
66 
67 
68def _make_token_provider():
69    """Create an auto-refreshing AAD token provider for long-running scripts.
70    
71    Returns a callable that the OpenAI SDK calls before each request to get
72    a fresh token. Tokens are cached and refreshed ~5 min before expiry.
73    """
74    from azure.identity import DefaultAzureCredential
75    credential = DefaultAzureCredential()
76 
77    def get_token():
78        try:
79            token = credential.get_token(_AZURE_COGSERVICES_SCOPE)
80            return token.token
81        except Exception as e:
82            raise RuntimeError(
83                f"Azure AD authentication failed: {e}\n"
84                "Ensure you're logged in (az login) or have valid "
85                "AZURE_CLIENT_ID/AZURE_TENANT_ID/AZURE_CLIENT_SECRET set."
86            ) from e
87 
88    return get_token
89 
90 
91def get_clients(base_url=None, azure_endpoint=None, project_endpoint=None, api_key=None):
92    """Initialize and return OpenAI-compatible client.
93 
94    Tries in order:
95    1. Project /v1/ endpoint with openai.OpenAI() (simplest, preferred)
96    2. Foundry SDK with AIProjectClient.get_openai_client() (no API key needed)
97    3. Azure OpenAI endpoint with openai.AzureOpenAI() (classic)
98 
99    When using DefaultAzureCredential (no API key), tokens are auto-refreshed
100    so long-running scripts won't fail with 401 after ~60 min.
101 
102    Returns: (openai_client, method_name)
103    """
104    # Method 1: /v1/ project endpoint
105    base_url = base_url or os.environ.get("OPENAI_BASE_URL")
106    api_key = api_key or os.environ.get("AZURE_OPENAI_API_KEY")
107 
108    if base_url:
109        import openai
110        if not api_key:
111            try:
112                token_provider = _make_token_provider()
113                token_provider()  # verify it works
114                # Use a custom httpx auth class that refreshes the token on each request
115                import httpx
116 
117                class _AzureADAuth(httpx.Auth):
118                    def __init__(self, provider):
119                        self._provider = provider
120 
121                    def auth_flow(self, request):
122                        request.headers["Authorization"] = f"Bearer {self._provider()}"
123                        yield request
124 
125                client = openai.OpenAI(
126                    base_url=base_url,
127                    api_key="aad",  # required by SDK but overridden by auth
128                    http_client=httpx.Client(auth=_AzureADAuth(token_provider)),
129                )
130                print(f"✅ Connected via /v1/ project endpoint (DefaultAzureCredential, auto-refresh)")
131                return client, "project-v1-aad"
132            except Exception as e:
133                print(f"⚠️ No API key and DefaultAzureCredential failed: {e}")
134        else:
135            client = openai.OpenAI(base_url=base_url, api_key=api_key)
136            print(f"✅ Connected via /v1/ project endpoint")
137            return client, "project-v1"
138 
139    # Method 2: Foundry SDK
140    project_endpoint = project_endpoint or os.environ.get("AZURE_AI_PROJECT_ENDPOINT")
141    if project_endpoint:
142        try:
143            from azure.ai.projects import AIProjectClient
144            from azure.identity import DefaultAzureCredential
145 
146            credential = DefaultAzureCredential()
147            project_client = AIProjectClient(endpoint=project_endpoint, credential=credential)
148            openai_client = project_client.get_openai_client()
149            print(f"✅ Connected via Foundry SDK")
150            return openai_client, "foundry-sdk"
151        except Exception as e:
152            print(f"⚠️ Foundry SDK failed: {e}")
153 
154    # Method 3: Azure OpenAI endpoint
155    azure_endpoint = azure_endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
156    if azure_endpoint:
157        import openai
158        if api_key:
159            client = openai.AzureOpenAI(
160                azure_endpoint=azure_endpoint,
161                api_key=api_key,
162                api_version="2025-04-01-preview",
163            )
164            print(f"✅ Connected via Azure OpenAI endpoint")
165            return client, "azure-openai"
166        else:
167            # No API key — use DefaultAzureCredential with auto-refresh
168            try:
169                token_provider = _make_token_provider()
170                token_provider()  # verify it works
171                client = openai.AzureOpenAI(
172                    azure_endpoint=azure_endpoint,
173                    azure_ad_token_provider=token_provider,
174                    api_version="2025-04-01-preview",
175                )
176                print(f"✅ Connected via Azure OpenAI endpoint (DefaultAzureCredential, auto-refresh)")
177                return client, "azure-openai-aad"
178            except Exception as e:
179                print(f"⚠️ DefaultAzureCredential failed for Azure endpoint: {e}")
180 
181    print("❌ No valid connection method. Set one of:")
182    print("   OPENAI_BASE_URL (preferred)")
183    print("   AZURE_AI_PROJECT_ENDPOINT (Foundry SDK)")
184    print("   AZURE_OPENAI_ENDPOINT + AZURE_OPENAI_API_KEY")
185    raise SystemExit(1)
186 
187 
188def upload_file(openai_client, filepath: str, purpose: str = "fine-tune") -> str:
189    """Upload a file to Azure AI Foundry and wait for processing."""
190    print(f"📤 Uploading {filepath}...")
191    with open(filepath, "rb") as f:
192        file_obj = openai_client.files.create(file=f, purpose=purpose)
193    print(f"   File ID: {file_obj.id}")
194    print(f"   Waiting for processing...")
195    openai_client.files.wait_for_processing(file_obj.id)
196    print(f"   ✅ File ready")
197    return file_obj.id
198