diff --git a/README.md b/README.md index 0aed382..d5007bf 100644 --- a/README.md +++ b/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://.azure-api.net/v1 +GENAI_GATEWAY_API_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. diff --git a/agent.py b/agent.py index 04a0ee8..86bc3eb 100644 --- a/agent.py +++ b/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 diff --git a/config.py b/config.py index c591a60..16b7188 100644 --- a/config.py +++ b/config.py @@ -3,7 +3,7 @@ """ 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 @@ -11,7 +11,7 @@ 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,6 +26,13 @@ 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://.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.""" @@ -33,13 +40,23 @@ class Config: 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 diff --git a/env.example b/env.example index 9200ca9..fb2b30b 100644 --- a/env.example +++ b/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://.azure-api.net/v1 +# GENAI_GATEWAY_API_KEY= +# GENAI_GATEWAY_API_VERSION=2024-06-01-preview