chore: import upstream snapshot with attribution

python/packages/ag-ui/getting_started/server.py

@@ -0,0 +1,144 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AG-UI server example with server-side tools."""
+
+from __future__ import annotations
+
+import logging
+import os
+
+from agent_framework import Agent, tool
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework.openai import OpenAIChatCompletionClient
+from dotenv import load_dotenv
+from fastapi import Depends, FastAPI, HTTPException, Security
+from fastapi.security import APIKeyHeader
+
+load_dotenv()
+
+# Enable debug logging
+logging.basicConfig(
+    level=logging.DEBUG,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+# Read required configuration
+endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
+model = os.environ.get("AZURE_OPENAI_MODEL")
+
+if not endpoint:
+    raise ValueError("AZURE_OPENAI_ENDPOINT environment variable is required")
+if not model:
+    raise ValueError("AZURE_OPENAI_MODEL environment variable is required")
+
+
+# ============================================================================
+# AUTHENTICATION EXAMPLE
+# ============================================================================
+# This demonstrates how to secure the AG-UI endpoint with API key authentication.
+# In production, you should use a more robust authentication mechanism such as:
+# - OAuth 2.0 / OpenID Connect
+# - JWT tokens with proper validation
+# - Azure AD / Entra ID integration
+# - Your organization's identity provider
+#
+# The API key should be stored securely (e.g., Azure Key Vault, environment variables)
+# and rotated regularly.
+# ============================================================================
+
+# API key header configuration
+API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False)
+
+# Get the expected API key from environment variable
+# In production, use a secrets manager like Azure Key Vault
+EXPECTED_API_KEY = os.environ.get("AG_UI_API_KEY")
+
+
+async def verify_api_key(api_key: str | None = Security(API_KEY_HEADER)) -> None:
+    """Verify the API key provided in the request header.
+
+    Args:
+        api_key: The API key from the X-API-Key header
+
+    Raises:
+        HTTPException: If the API key is missing or invalid
+    """
+    if not EXPECTED_API_KEY:
+        # If no API key is configured, log a warning but allow the request
+        # This maintains backward compatibility but warns about the security risk
+        logger.warning(
+            "AG_UI_API_KEY environment variable not set. "
+            "The endpoint is accessible without authentication. "
+            "Set AG_UI_API_KEY to enable API key authentication."
+        )
+        return
+
+    if not api_key:
+        raise HTTPException(
+            status_code=401,
+            detail="Missing API key. Provide X-API-Key header.",
+        )
+
+    if api_key != EXPECTED_API_KEY:
+        raise HTTPException(
+            status_code=403,
+            detail="Invalid API key.",
+        )
+
+
+# Server-side tool (executes on server)
+@tool(description="Get the time zone for a location.")
+def get_time_zone(location: str) -> str:
+    """Get the time zone for a location.
+
+    Args:
+        location: The city or location name
+    """
+    print(f"[SERVER] get_time_zone tool called with location: {location}")
+    timezone_data = {
+        "seattle": "Pacific Time (UTC-8)",
+        "san francisco": "Pacific Time (UTC-8)",
+        "new york": "Eastern Time (UTC-5)",
+        "london": "Greenwich Mean Time (UTC+0)",
+    }
+    result = timezone_data.get(location.lower(), f"Time zone data not available for {location}")
+    print(f"[SERVER] get_time_zone returning: {result}")
+    return result
+
+
+# Create the AI agent with ONLY server-side tools
+# IMPORTANT: Do NOT include tools that the client provides!
+# In this example:
+# - get_time_zone: SERVER-ONLY tool (only server has this)
+# - get_weather: CLIENT-ONLY tool (client provides this, server should NOT include it)
+# The client will send get_weather tool metadata so the LLM knows about it,
+# and the function invocation mixin on AGUIChatClient will execute it client-side.
+# This matches the .NET AG-UI hybrid execution pattern.
+agent = Agent(
+    name="AGUIAssistant",
+    instructions="You are a helpful assistant. Use get_weather for weather and get_time_zone for time zones.",
+    client=OpenAIChatCompletionClient(
+        azure_endpoint=endpoint,
+        model=model,
+    ),
+    tools=[get_time_zone],  # ONLY server-side tools
+)
+
+# Create FastAPI app
+app = FastAPI(title="AG-UI Server")
+
+# Register the AG-UI endpoint with authentication
+# The dependencies parameter accepts FastAPI Depends() objects that run before the handler
+add_agent_framework_fastapi_endpoint(
+    app,
+    agent,
+    "/",
+    dependencies=[Depends(verify_api_key)],
+)
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="127.0.0.1", port=5100, log_level="debug", access_log=True)