Files
agent-framework/python/packages/ag-ui
T
Eduard van Valkenburg 6138487888 Python: Phase 2: Embedding clients for Ollama, Bedrock, and Azure AI Inference (#4207)
* Phase 2: Embedding clients for Ollama, Bedrock, and Azure AI Inference

Add embedding client implementations to existing provider packages:

- OllamaEmbeddingClient: Text embeddings via Ollama's embed API
- BedrockEmbeddingClient: Text embeddings via Amazon Titan on Bedrock
- AzureAIInferenceEmbeddingClient: Text and image embeddings via Azure AI
  Inference, supporting Content | str input with separate model IDs for
  text (AZURE_AI_INFERENCE_EMBEDDING_MODEL_ID) and image
  (AZURE_AI_INFERENCE_IMAGE_EMBEDDING_MODEL_ID) endpoints

Additional changes:
- Rename EmbeddingCoT -> EmbeddingT, EmbeddingOptionsCoT -> EmbeddingOptionsT
- Add otel_provider_name passthrough to all embedding clients
- Register integration pytest marker in all packages
- Add lazy-loading namespace exports for Ollama and Bedrock embeddings
- Add image embedding sample using Cohere-embed-v3-english
- Add azure-ai-inference dependency to azure-ai package

Part of #1188

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix mypy duplicate name and ruff lint issues

- Rename second 'vector' variable to 'img_vector' in image embedding loop
- Combine nested with statements in tests
- Remove unused result assignments in tests

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* updates from feedback

* Fix CI failures in embedding usage handling

- Fix Azure AI embedding mypy issues by normalizing vectors to list[float],
  safely accumulating optional usage token fields, and filtering None entries
  before constructing GeneratedEmbeddings
- Avoid Bandit false positive by initializing usage details as an empty dict
- Update OpenAI embedding tests to assert canonical usage keys
  (input_token_count/total_token_count)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
6138487888 ยท 2026-02-25 17:45:08 +00:00
History
..
2025-11-05 05:25:24 +00:00

Agent Framework AG-UI Integration

AG-UI protocol integration for Agent Framework, enabling seamless integration with AG-UI's web interface and streaming protocol.

Installation

pip install agent-framework-ag-ui

Quick Start

Server (Host an AI Agent)

from fastapi import FastAPI
from agent_framework import Agent
from agent_framework.azure import AzureOpenAIChatClient
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint

# Create your agent
agent = Agent(
    name="my_agent",
    instructions="You are a helpful assistant.",
    client=AzureOpenAIChatClient(
        endpoint="https://your-resource.openai.azure.com/",
        deployment_name="gpt-4o-mini",
        api_key="your-api-key",
    ),
)

# Create FastAPI app and add AG-UI endpoint
app = FastAPI()
add_agent_framework_fastapi_endpoint(app, agent, "/")

# Run with: uvicorn main:app --reload

Server (Host a Workflow)

from fastapi import FastAPI
from agent_framework import WorkflowBuilder, WorkflowContext, executor
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint

@executor(id="start")
async def start(message: str, ctx: WorkflowContext) -> None:
    await ctx.yield_output(f"Workflow received: {message}")

workflow = WorkflowBuilder(start_executor=start).build()

app = FastAPI()
add_agent_framework_fastapi_endpoint(app, workflow, "/")

Server (Thread-Scoped WorkflowBuilder)

Use workflow_factory when your workflow keeps runtime state (for example pending request_info interrupts) and must be isolated per AG-UI thread:

from fastapi import FastAPI
from agent_framework import Workflow, WorkflowBuilder
from agent_framework.ag_ui import AgentFrameworkWorkflow, add_agent_framework_fastapi_endpoint

def build_workflow_for_thread(thread_id: str) -> Workflow:
    # Build a fresh workflow instance for each thread id.
    return WorkflowBuilder(start_executor=...).build()

app = FastAPI()
thread_scoped_workflow = AgentFrameworkWorkflow(
    workflow_factory=build_workflow_for_thread,
    name="my_workflow",
)
add_agent_framework_fastapi_endpoint(app, thread_scoped_workflow, "/")

Client (Connect to an AG-UI Server)

import asyncio
from agent_framework.ag_ui import AGUIChatClient

async def main():
    async with AGUIChatClient(endpoint="http://localhost:8000/") as client:
        # Stream responses
        async for update in client.get_response("Hello!", stream=True):
            for content in update.contents:
                if content.type == "text" and content.text:
                    print(content.text, end="", flush=True)
        print()

asyncio.run(main())

The AGUIChatClient supports:

  • Streaming and non-streaming responses
  • Hybrid tool execution (client-side + server-side tools)
  • Automatic thread management for conversation continuity
  • Integration with Agent for client-side history management
  • Interrupt metadata passthrough (availableInterrupts and resume)

Documentation

  • Getting Started Tutorial - Step-by-step guide to building AG-UI servers and clients
    • Server setup with FastAPI
    • Client examples using AGUIChatClient
    • Hybrid tool execution (client-side + server-side)
    • Thread management and conversation continuity
  • Examples - Complete examples for AG-UI features

Features

This integration supports all 7 AG-UI features:

  1. Agentic Chat: Basic streaming chat with tool calling support
  2. Backend Tool Rendering: Tools executed on backend with results streamed to client
  3. Human in the Loop: Function approval requests for user confirmation before tool execution
  4. Agentic Generative UI: Async tools for long-running operations with progress updates
  5. Tool-based Generative UI: Custom UI components rendered on frontend based on tool calls
  6. Shared State: Bidirectional state sync between client and server
  7. Predictive State Updates: Stream tool arguments as optimistic state updates during execution

Additional compatibility and draft support:

  • Native Workflow endpoint registration via add_agent_framework_fastapi_endpoint(...)
  • Workflow-to-AG-UI event mapping (run/step/activity/tool/custom events)
  • Custom event compatibility for inbound CUSTOM, CUSTOM_EVENT, and custom_event
  • Pragmatic multimodal input parsing for both legacy (binary) and draft media-part shapes
  • Pragmatic interrupt/resume handling (availableInterrupts, resume, and RUN_FINISHED.interrupt)

Security: Authentication & Authorization

The AG-UI endpoint does not enforce authentication by default. For production deployments, you should add authentication using FastAPI's dependency injection system via the dependencies parameter.

API Key Authentication Example

import os
from fastapi import Depends, FastAPI, HTTPException, Security
from fastapi.security import APIKeyHeader
from agent_framework import Agent
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint

# Configure API key authentication
API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False)
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."""
    if not api_key or api_key != EXPECTED_API_KEY:
        raise HTTPException(status_code=401, detail="Invalid or missing API key")

# Create agent and app
agent = Agent(name="my_agent", instructions="...", client=...)
app = FastAPI()

# Register endpoint WITH authentication
add_agent_framework_fastapi_endpoint(
    app,
    agent,
    "/",
    dependencies=[Depends(verify_api_key)],  # Authentication enforced here
)

Other Authentication Options

The dependencies parameter accepts any FastAPI dependency, enabling integration with:

  • OAuth 2.0 / OpenID Connect - Use fastapi.security.OAuth2PasswordBearer
  • JWT Tokens - Validate tokens with libraries like python-jose
  • Azure AD / Entra ID - Use azure-identity for Microsoft identity platform
  • Rate Limiting - Add request throttling dependencies
  • Custom Authentication - Implement your organization's auth requirements

For a complete authentication example, see getting_started/server.py.

Architecture

The package uses a clean, orchestrator-based architecture:

  • AgentFrameworkAgent: Lightweight wrapper that delegates to orchestrators
  • Orchestrators: Handle different execution flows (default, human-in-the-loop, etc.)
  • Confirmation Strategies: Domain-specific confirmation messages (extensible)
  • AgentFrameworkEventBridge: Converts Agent Framework events to AG-UI events
  • Message Adapters: Bidirectional conversion between AG-UI and Agent Framework message formats
  • FastAPI Endpoint: Streaming HTTP endpoint with Server-Sent Events (SSE)

Next Steps

  1. New to AG-UI? Start with the Getting Started Tutorial
  2. Want to see examples? Check out the Examples for AG-UI features

License

MIT