Add Azure APIM Gen-AI Gateway support and update configuration options

README.md

@@ -134,3 +134,26 @@ UserCtx:
 ---
 
 **🚀 Ready to see AI memory in action? Start chatting and watch your travel preferences get smarter!**
+
+---
+
+## 🔐 (Optional) Azure APIM Gen-AI Gateway
+
+You can route ALL model calls through an Azure API Management (APIM) Gen-AI Gateway that exposes an OpenAI-compatible endpoint. This lets you centralize policies like authentication, rate-limits, prompt safety, and usage analytics without changing application code paths.
+
+How this app integrates:
+- When `GENAI_GATEWAY_ENABLED=true`, the app sets `OPENAI_BASE_URL` to `GENAI_GATEWAY_BASE_URL` and uses `GENAI_GATEWAY_API_KEY` as the `OPENAI_API_KEY`.
+- The AutoGen OpenAI client (`OpenAIChatCompletionClient`) reads these env vars and sends all requests via APIM.
+
+Env variables in `.env`:
+```
+GENAI_GATEWAY_ENABLED=true
+GENAI_GATEWAY_BASE_URL=https://<apim-name>.azure-api.net/v1
+GENAI_GATEWAY_API_KEY=<your-apim-subscription-key>
+# Optional if your gateway requires it as a query string
+GENAI_GATEWAY_API_VERSION=2024-06-01-preview
+```
+
+Notes:
+- Ensure your APIM endpoint is OpenAI-compatible and typically includes `/v1`.
+- If you disable the gateway (default), the app will use `OPENAI_API_KEY` directly against OpenAI.

agent.py

@@ -74,7 +74,16 @@ def __init__(self, config: Optional[AppConfig] = None):
         self.config = config
 
         # Set environment variables for SDK clients
-        os.environ["OPENAI_API_KEY"] = config.openai_api_key
+        # Route OpenAI traffic through Azure APIM Gen-AI Gateway if enabled
+        if getattr(config, "genai_gateway_enabled", False):
+            if config.genai_gateway_api_key:
+                os.environ["OPENAI_API_KEY"] = config.genai_gateway_api_key
+            if config.genai_gateway_base_url:
+                # Expect an OpenAI-compatible base URL (often ends with '/v1')
+                os.environ["OPENAI_BASE_URL"] = config.genai_gateway_base_url
+        else:
+            os.environ["OPENAI_API_KEY"] = config.openai_api_key
+            # Do not force OPENAI_BASE_URL when not using gateway; leave default
         os.environ["TAVILY_API_KEY"] = config.tavily_api_key
 
         # Initialize shared clients

config.py

@@ -3,15 +3,15 @@
 """
 import os
 from typing import Optional
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, Field, field_validator, model_validator
 from pydantic_settings import BaseSettings
 
 
 class AppConfig(BaseSettings):
     """Application configuration with validation."""
 
     # API Keys
-    openai_api_key: str = Field(..., env="OPENAI_API_KEY", description="OpenAI API key")
+    openai_api_key: str | None = Field(default=None, env="OPENAI_API_KEY", description="OpenAI API key")
     tavily_api_key: str = Field(..., env="TAVILY_API_KEY", description="Tavily API key")
 
     # Model Configuration
@@ -26,20 +26,37 @@ class AppConfig(BaseSettings):
     server_name: str = Field(default="0.0.0.0", env="SERVER_NAME", description="Server host")
     server_port: int = Field(default=7860, env="SERVER_PORT", description="Server port")
     share: bool = Field(default=False, env="SHARE", description="Enable public sharing")
+
+    # Azure APIM Gen-AI Gateway (optional)
+    # When enabled, all OpenAI-compatible calls will be routed through the gateway.
+    genai_gateway_enabled: bool = Field(default=False, env="GENAI_GATEWAY_ENABLED", description="Enable Azure APIM Gen-AI Gateway routing")
+    genai_gateway_base_url: str | None = Field(default=None, env="GENAI_GATEWAY_BASE_URL", description="Base URL for the OpenAI-compatible gateway endpoint (e.g., https://<apim-name>.azure-api.net/v1)")
+    genai_gateway_api_key: str | None = Field(default=None, env="GENAI_GATEWAY_API_KEY", description="Subscription/API key for the gateway (used instead of OPENAI_API_KEY when enabled)")
+    genai_gateway_api_version: str | None = Field(default=None, env="GENAI_GATEWAY_API_VERSION", description="Optional API version query string required by the gateway, if any")
 
     class Config:
         """Pydantic config."""
         env_file = ".env"
         env_file_encoding = "utf-8"
         case_sensitive = False
 
-    @field_validator("openai_api_key")
-    @classmethod
-    def validate_openai_key(cls, v):
-        """Validate OpenAI API key format."""
-        if not v.startswith("sk-"):
-            raise ValueError("OpenAI API key must start with 'sk-'")
-        return v
+    # Per-field validation for OPENAI_API_KEY is skipped so that gateway-only configs can omit it.
+
+    @model_validator(mode="after")
+    def validate_gateway_vs_openai(self) -> "AppConfig":
+        """Cross-field validation for gateway configuration vs direct OpenAI usage."""
+        if self.genai_gateway_enabled:
+            if not self.genai_gateway_base_url:
+                raise ValueError("GENAI_GATEWAY_ENABLED is true but GENAI_GATEWAY_BASE_URL is not set")
+            if not self.genai_gateway_api_key:
+                raise ValueError("GENAI_GATEWAY_ENABLED is true but GENAI_GATEWAY_API_KEY is not set")
+        else:
+            # When not using the gateway, enforce typical OpenAI key format.
+            if not self.openai_api_key or not isinstance(self.openai_api_key, str):
+                raise ValueError("OPENAI_API_KEY is required when GENAI_GATEWAY_ENABLED=false")
+            if not self.openai_api_key.startswith("sk-"):
+                raise ValueError("OpenAI API key must start with 'sk-' when GENAI_GATEWAY_ENABLED=false")
+        return self
 
 
 
@@ -64,9 +81,17 @@ def validate_dependencies() -> bool:
 
     # Test OpenAI API
     try:
-        client = OpenAI(api_key=config.openai_api_key)
+        # Route through gateway if enabled (OpenAI-compatible endpoint)
+        if config.genai_gateway_enabled:
+            base_url = config.genai_gateway_base_url
+            # Ensure no trailing slash inconsistencies
+            if base_url and not base_url.rstrip().endswith("/v1") and "/v1" not in base_url:
+                print("ℹ️ Note: Expected an OpenAI-compatible base URL that includes '/v1'. Current:", base_url)
+            client = OpenAI(api_key=config.genai_gateway_api_key, base_url=base_url)
+        else:
+            client = OpenAI(api_key=config.openai_api_key)
         # Just test the client creation, not making an actual API call
-        print("✅ OpenAI API key configured")
+        print("✅ OpenAI client configured")
     except Exception as e:
         print(f"❌ OpenAI API error: {e}")
         return False

env.example

@@ -17,3 +17,10 @@ REDIS_URL=redis://localhost:6379
 SERVER_NAME=0.0.0.0
 SERVER_PORT=7860
 SHARE=false
+
+# Optional Azure APIM Gen-AI Gateway (OpenAI-compatible) routing
+# Set these to route ALL OpenAI traffic through your APIM gateway
+# GENAI_GATEWAY_ENABLED=true
+# GENAI_GATEWAY_BASE_URL=https://<apim-name>.azure-api.net/v1
+# GENAI_GATEWAY_API_KEY=<your-apim-subscription-key>
+# GENAI_GATEWAY_API_VERSION=2024-06-01-preview