agent: Azure OpenAI support for visit extractor + condenser

Mirrors inference/tool_visit.py: when AZURE_OPENAI_ENDPOINT is set we use
AzureOpenAI() (with AZURE_OPENAI_API_VERSION) and the effective model name
becomes AZURE_OPENAI_DEPLOYMENT, otherwise we fall back to the plain
OpenAI client + SUMMARY_MODEL_NAME / MEMORY_MODEL_NAME. Condenser trigger
gate accepts either MEMORY_MODEL_NAME or AZURE_OPENAI_DEPLOYMENT.

app.py

@@ -1567,16 +1567,52 @@ MEMORY_TOKEN_THRESHOLD = int(
     or "16000"
 )
 
-
-def _get_openai_client(api_key: str, base_url: Optional[str]):
-    """Lazy import so the Space still imports if `openai` isn't installed yet."""
+# Azure OpenAI support — mirrors inference/tool_visit.py logic. When
+# AZURE_OPENAI_ENDPOINT is set, we use AzureOpenAI() instead of OpenAI()
+# and AZURE_OPENAI_DEPLOYMENT overrides the per-purpose model name.
+AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT", "").strip()
+AZURE_OPENAI_API_VERSION = os.getenv("AZURE_OPENAI_API_VERSION", "").strip() or "2024-06-01"
+AZURE_OPENAI_DEPLOYMENT = os.getenv("AZURE_OPENAI_DEPLOYMENT", "").strip()
+
+
+def _get_chat_client_and_model(
+    api_key: str, base_url: Optional[str], fallback_model_name: str
+) -> Tuple[Optional[Any], str]:
+    """Construct an OpenAI-compatible chat client. Auto-switches to
+    AzureOpenAI when AZURE_OPENAI_ENDPOINT is configured; in that case the
+    effective model name becomes AZURE_OPENAI_DEPLOYMENT (Azure uses
+    deployment names, not raw model ids). Returns (client, model_name)."""
+    if AZURE_OPENAI_ENDPOINT:
+        try:
+            from openai import AzureOpenAI
+        except Exception:
+            return None, fallback_model_name
+        if not api_key:
+            return None, fallback_model_name
+        client = AzureOpenAI(
+            api_key=api_key,
+            azure_endpoint=AZURE_OPENAI_ENDPOINT,
+            api_version=AZURE_OPENAI_API_VERSION,
+        )
+        return client, (AZURE_OPENAI_DEPLOYMENT or fallback_model_name)
     try:
         from openai import OpenAI
     except Exception:
-        return None
+        return None, fallback_model_name
     if not api_key:
-        return None
-    return OpenAI(api_key=api_key, base_url=base_url) if base_url else OpenAI(api_key=api_key)
+        return None, fallback_model_name
+    client = (
+        OpenAI(api_key=api_key, base_url=base_url) if base_url else OpenAI(api_key=api_key)
+    )
+    return client, fallback_model_name
+
+
+# Backwards-compat shim: older callers asked for just a client without
+# Azure awareness. Keep returning the non-Azure client so we don't break
+# anything if a future patch imports it.
+def _get_openai_client(api_key: str, base_url: Optional[str]):
+    client, _ = _get_chat_client_and_model(api_key, base_url, fallback_model_name="")
+    return client
 
 
 def _approx_token_count(text: str) -> int:
@@ -1616,22 +1652,28 @@ _LAST_EXTRACT_ERR: Optional[str] = None
 
 def _llm_extract(webpage_content: str, goal: str) -> Optional[str]:
     """Run the SUMMARY model as the visit extractor. Mirrors
-    inference/prompt.py:build_visit_extractor_messages + tool_visit's call."""
+    inference/prompt.py:build_visit_extractor_messages + tool_visit's call.
+    Picks AzureOpenAI when AZURE_OPENAI_ENDPOINT is set."""
     global _LAST_EXTRACT_ERR
     _LAST_EXTRACT_ERR = None
-    if not SUMMARY_MODEL_NAME:
-        _LAST_EXTRACT_ERR = "SUMMARY_MODEL_NAME env var not set"
-        return None
     if not SUMMARY_API_KEY:
-        _LAST_EXTRACT_ERR = "API_KEY / SUMMARY_OPENAI_API_KEY env var not set"
+        _LAST_EXTRACT_ERR = "API_KEY env var not set"
         return None
-    client = _get_openai_client(SUMMARY_API_KEY, SUMMARY_API_BASE)
+    client, model_name = _get_chat_client_and_model(
+        SUMMARY_API_KEY, SUMMARY_API_BASE, SUMMARY_MODEL_NAME
+    )
     if client is None:
         _LAST_EXTRACT_ERR = "openai client could not be constructed (package missing?)"
         return None
+    if not model_name:
+        _LAST_EXTRACT_ERR = (
+            "no model name (set SUMMARY_MODEL_NAME or, on Azure, "
+            "AZURE_OPENAI_DEPLOYMENT)"
+        )
+        return None
     try:
         resp = client.chat.completions.create(
-            model=SUMMARY_MODEL_NAME,
+            model=model_name,
             messages=[
                 {
                     "role": "user",
@@ -1650,9 +1692,12 @@ def _llm_extract(webpage_content: str, goal: str) -> Optional[str]:
 
 def _llm_condense(events_text: str, prev_state: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
     """Run the MEMORY model as the State Summarizer. Returns a parsed JSON
-    state dict, or None if condensation failed."""
-    client = _get_openai_client(MEMORY_API_KEY, MEMORY_API_BASE)
-    if client is None or not MEMORY_MODEL_NAME:
+    state dict, or None if condensation failed. Picks AzureOpenAI when
+    AZURE_OPENAI_ENDPOINT is set."""
+    client, model_name = _get_chat_client_and_model(
+        MEMORY_API_KEY, MEMORY_API_BASE, MEMORY_MODEL_NAME
+    )
+    if client is None or not model_name:
         return None
     user_payload = json.dumps(
         {
@@ -1663,7 +1708,7 @@ def _llm_condense(events_text: str, prev_state: Optional[Dict[str, Any]]) -> Opt
     )
     try:
         resp = client.chat.completions.create(
-            model=MEMORY_MODEL_NAME,
+            model=model_name,
             messages=[
                 {"role": "system", "content": MEMORY_SYSTEM_PROMPT},
                 {"role": "user", "content": user_payload},
@@ -1980,7 +2025,7 @@ def build_research_agent(
         # context as [system, original_question, RESEARCH_STATE_SUMMARY].
         if (
             strategy == "condenser"
-            and MEMORY_MODEL_NAME
+            and (MEMORY_MODEL_NAME or AZURE_OPENAI_DEPLOYMENT)
             and MEMORY_API_KEY
             and turn > 1
             and _messages_token_count(messages) > MEMORY_TOKEN_THRESHOLD
@@ -2016,7 +2061,7 @@ def build_research_agent(
                 yield _emit()
         elif (
             strategy == "condenser"
-            and (not MEMORY_MODEL_NAME or not MEMORY_API_KEY)
+            and not ((MEMORY_MODEL_NAME or AZURE_OPENAI_DEPLOYMENT) and MEMORY_API_KEY)
             and state.trusted_notes
             and turn > 1
             and turn % 3 == 0