agent-framework-lib 0.8.8.post9__tar.gz → 0.8.9__tar.gz

{agent_framework_lib-0.8.8.post9/agent_framework_lib.egg-info → agent_framework_lib-0.8.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agent-framework-lib
-Version: 0.8.8.post9
+Version: 0.8.9
 Summary: A comprehensive Python framework for building and serving conversational AI agents with FastAPI
 Author-email: Sebastian Pavel <sebastian@cinco.ai>, Elliott Girard <elliott.girard@icloud.com>
 Maintainer-email: Sebastian Pavel <sebastian@cinco.ai>
@@ -63,7 +63,7 @@ Requires-Dist: llama-index-llms-openai>=0.7.9
 Requires-Dist: llama-index-llms-anthropic>=0.11.5
 Requires-Dist: llama-index-llms-google-genai>=0.1.0
 Requires-Dist: graphiti-core>=0.24.3
-Requires-Dist: tiktoken>=0.7.0
+Requires-Dist: tiktoken>=0.8.0
 Requires-Dist: falkordb>=1.0.0
 Requires-Dist: grpcio-status>=1.71.2
 Requires-Dist: nodeenv>=1.8.0
@@ -84,6 +84,8 @@ Requires-Dist: llama-index-workflows>=2.16.0; extra == "llamaindex"
 Requires-Dist: llama-index-llms-openai>=0.7.9; extra == "llamaindex"
 Requires-Dist: llama-index-llms-google-genai>=0.1.0; extra == "llamaindex"
 Requires-Dist: llama-index-llms-anthropic>=0.11.5; extra == "llamaindex"
+Provides-Extra: langgraph
+Requires-Dist: langgraph>=0.2.0; extra == "langgraph"
 Provides-Extra: mcp
 Requires-Dist: llama-index-tools-mcp>=0.4.0; extra == "mcp"
 Provides-Extra: websearch
@@ -151,7 +153,7 @@ Requires-Dist: traceloop-sdk>=0.30.0; extra == "observability"
 Provides-Extra: monitoring
 Requires-Dist: psutil>=5.9.0; extra == "monitoring"
 Provides-Extra: all
-Requires-Dist: agent-framework-lib[dev,drawio,elasticsearch,excel,gcp,graphiti-all,llamaindex,mcp,memory,microsoft,minio,mongodb,monitoring,multimodal,observability,postgresql,s3,websearch]; extra == "all"
+Requires-Dist: agent-framework-lib[dev,drawio,elasticsearch,excel,gcp,graphiti-all,langgraph,llamaindex,mcp,memory,microsoft,minio,mongodb,monitoring,multimodal,observability,postgresql,s3,websearch]; extra == "all"
 Dynamic: license-file
 
 # Agent Framework Library

{agent_framework_lib-0.8.8.post9 → agent_framework_lib-0.8.9}/agent_framework/core/model_clients.py

@@ -8,7 +8,7 @@ This module is framework-agnostic and uses native client libraries.
 import json
 import logging
 import re
-from typing import Any
+from typing import Any, cast
 
 from .agent_interface import AgentConfig
 from .model_config import ModelConfigManager, ModelProvider, model_config
@@ -1095,3 +1095,572 @@ class ModelClientFactory:
 
 # Global factory instance
 client_factory = ModelClientFactory()
+
+
+# ===========================================================================
+# Provider-execution layer (additive)
+# ---------------------------------------------------------------------------
+# This is the reusable "execute a chat completion" capability that does NOT
+# exist elsewhere in the codebase today: ``ModelRouter`` only *selects* a model
+# and ``ModelClientFactory`` only *creates* a raw provider SDK client. These
+# additive module-level functions own ALL provider branching (dispatching on the
+# concrete client type to the correct SDK method) and normalize every provider's
+# native response shape onto a single vendor-neutral structure, so callers such
+# as ``workflow/core`` stay provider-neutral.
+#
+# Token usage is taken from the provider's reported counts whenever available
+# (OpenAI ``usage.prompt_tokens``/``completion_tokens``; Anthropic
+# ``usage.input_tokens``/``output_tokens``; Gemini ``usage_metadata``), plus
+# ``thinking_tokens`` where the provider reports them. ``LLMMetricsCollector``
+# estimation is used ONLY as a fallback when the provider reports no usage
+# (notably mid-stream or a stream that ends without a usage chunk) (Req 2.7, 2.12).
+# ===========================================================================
+
+from collections.abc import AsyncIterator  # noqa: E402
+from dataclasses import dataclass  # noqa: E402
+
+
+# Default max_tokens for providers (Anthropic) that REQUIRE the parameter when
+# the caller and model_params supply none.
+_DEFAULT_MAX_TOKENS = 4096
+
+
+@dataclass
+class ChatExecutionResult:
+    """Normalized, provider-agnostic result of a non-streaming chat completion.
+
+    The orchestration layer (``workflow/core/chat_completion.py``) maps this onto
+    the public ``ChatCompletionResult`` and builds ``LLMMetrics`` from it. Token
+    counts are the EXACT provider-reported values for a completed non-streaming
+    call; ``usage_estimated`` is ``True`` only when the provider reported no usage
+    and the counts were estimated via ``LLMMetricsCollector`` (Req 2.12).
+    """
+
+    content: str
+    model: str
+    input_tokens: int
+    output_tokens: int
+    thinking_tokens: int
+    finish_reason: str | None = None
+    usage_estimated: bool = False
+
+
+@dataclass
+class ChatExecutionChunk:
+    """A single normalized incremental chunk yielded by a streaming completion.
+
+    ``delta`` is the text fragment since the previous chunk. The final chunk sets
+    ``done=True`` and carries the best-effort token counts: provider-reported when
+    the stream emitted a usage chunk, otherwise estimated via
+    ``LLMMetricsCollector`` (token counts mid-stream are best-effort, Req 2.7).
+    """
+
+    delta: str
+    model: str
+    done: bool = False
+    finish_reason: str | None = None
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+    thinking_tokens: int | None = None
+    usage_estimated: bool = False
+
+
+def _split_system_messages(
+    messages: list[dict[str, Any]],
+) -> tuple[str | None, list[dict[str, Any]]]:
+    """Separate system-role messages from the rest.
+
+    Returns ``(system_text, non_system_messages)``. Multiple system messages are
+    joined with newlines. Used by ``_execute_anthropic`` because Anthropic takes
+    no system role inside ``messages`` and instead requires the top-level
+    ``system`` parameter (Req 2.11).
+    """
+    system_parts: list[str] = []
+    remaining: list[dict[str, Any]] = []
+    for msg in messages:
+        if msg.get("role") == "system":
+            content = msg.get("content")
+            if content:
+                system_parts.append(content if isinstance(content, str) else str(content))
+        else:
+            remaining.append(msg)
+    system_text = "\n".join(system_parts) if system_parts else None
+    return system_text, remaining
+
+
+def _join_message_text(messages: list[dict[str, Any]]) -> str:
+    """Concatenate message contents into a single string for token estimation."""
+    parts: list[str] = []
+    for msg in messages:
+        content = msg.get("content")
+        if isinstance(content, str):
+            parts.append(content)
+        elif content is not None:
+            parts.append(str(content))
+    return "\n".join(parts)
+
+
+def _estimate_token_counts(model: str | None, input_text: str, output_text: str) -> tuple[int, int]:
+    """Estimate (input_tokens, output_tokens) via ``LLMMetricsCollector``.
+
+    This is the fallback used ONLY when a provider reports no usage. Imported
+    lazily to avoid a hard import-time dependency on the monitoring package.
+    """
+    try:
+        from agent_framework.monitoring.llm_metrics_collector import LLMMetricsCollector
+
+        collector = LLMMetricsCollector(model_name=model, enabled=True)
+        input_tokens = collector.count_input(input_text)
+        output_tokens = collector.count_output(output_text)
+        return input_tokens, output_tokens
+    except Exception as e:  # pragma: no cover - defensive: estimation must not crash a call
+        logger.warning(f"[execute_chat_completion] Token estimation fallback failed: {e}")
+        return 0, 0
+
+
+def _resolve_call_params(
+    model_params: dict[str, Any],
+    max_tokens: int | None,
+    temperature: float | None,
+) -> tuple[str, float, int | None]:
+    """Resolve effective (model, temperature, max_tokens).
+
+    Explicit call arguments take precedence over the client's ``model_params``.
+    """
+    model = model_params.get("model")
+    effective_temperature = (
+        temperature if temperature is not None else model_params.get("temperature", 0.7)
+    )
+    effective_max_tokens = max_tokens if max_tokens is not None else model_params.get("max_tokens")
+    return cast(str, model), effective_temperature, effective_max_tokens
+
+
+async def execute_chat_completion(
+    client: Any,
+    model_params: dict[str, Any],
+    messages: list[dict[str, Any]],
+    *,
+    stream: bool = False,
+    max_tokens: int | None = None,
+    temperature: float | None = None,
+) -> "ChatExecutionResult | AsyncIterator[ChatExecutionChunk]":
+    """Execute a chat completion against a raw provider client and normalize it.
+
+    Dispatches on the concrete client type (``AsyncOpenAI`` →
+    ``client.chat.completions.create``; ``AsyncAnthropic`` →
+    ``client.messages.create``; Gemini ``GenerativeModel`` →
+    ``generate_content_async``) and returns a provider-agnostic result. All
+    provider branching lives in this layer so ``workflow/core`` stays neutral.
+
+    Args:
+        client: A raw provider SDK client created by ``ModelClientFactory``.
+        model_params: The ``client._model_params`` dict (model name, temperature,
+            and any provider-specific extras attached by the factory).
+        messages: OpenAI-style ``[{"role", "content"}]`` list. The system-role
+            message is ALWAYS preserved and placed in each provider's expected
+            slot by its helper (Req 2.11).
+        stream: When ``True``, returns an async iterator of ``ChatExecutionChunk``.
+        max_tokens: Optional max output tokens (overrides ``model_params``).
+        temperature: Optional sampling temperature (overrides ``model_params``).
+
+    Returns:
+        A ``ChatExecutionResult`` for non-streaming calls, or an
+        ``AsyncIterator[ChatExecutionChunk]`` when ``stream=True``.
+
+    Raises:
+        TypeError: If the client type is not a recognized provider client.
+    """
+    # Dispatch on the concrete client type. Availability flags guard the isinstance
+    # checks so a missing optional SDK never raises at dispatch time.
+    if OPENAI_AVAILABLE and isinstance(client, AsyncOpenAI):
+        return await _execute_openai(
+            client,
+            model_params,
+            messages,
+            stream=stream,
+            max_tokens=max_tokens,
+            temperature=temperature,
+        )
+    if ANTHROPIC_AVAILABLE and isinstance(client, AsyncAnthropic):
+        return await _execute_anthropic(
+            client,
+            model_params,
+            messages,
+            stream=stream,
+            max_tokens=max_tokens,
+            temperature=temperature,
+        )
+    if GEMINI_AVAILABLE and isinstance(client, genai.GenerativeModel):
+        return await _execute_gemini(
+            client,
+            model_params,
+            messages,
+            stream=stream,
+            max_tokens=max_tokens,
+            temperature=temperature,
+        )
+    # Fallback: dispatch by duck-typing on the SDK surface when isinstance fails
+    # (e.g., a subclass or a test double). This keeps dispatch robust without
+    # leaking provider types out of this layer.
+    if hasattr(client, "chat") and hasattr(client.chat, "completions"):
+        return await _execute_openai(
+            client,
+            model_params,
+            messages,
+            stream=stream,
+            max_tokens=max_tokens,
+            temperature=temperature,
+        )
+    if hasattr(client, "messages") and hasattr(client.messages, "create"):
+        return await _execute_anthropic(
+            client,
+            model_params,
+            messages,
+            stream=stream,
+            max_tokens=max_tokens,
+            temperature=temperature,
+        )
+    if hasattr(client, "generate_content_async"):
+        return await _execute_gemini(
+            client,
+            model_params,
+            messages,
+            stream=stream,
+            max_tokens=max_tokens,
+            temperature=temperature,
+        )
+    raise TypeError(
+        f"execute_chat_completion: unrecognized client type {type(client).__name__!r}; "
+        "expected an AsyncOpenAI, AsyncAnthropic, or Gemini GenerativeModel client"
+    )
+
+
+async def _execute_openai(
+    client: Any,
+    model_params: dict[str, Any],
+    messages: list[dict[str, Any]],
+    *,
+    stream: bool = False,
+    max_tokens: int | None = None,
+    temperature: float | None = None,
+) -> "ChatExecutionResult | AsyncIterator[ChatExecutionChunk]":
+    """Execute via ``client.chat.completions.create``.
+
+    OpenAI accepts the system role INSIDE the messages list, so the messages are
+    passed through unchanged (Req 2.11). Token usage is read from
+    ``usage.prompt_tokens``/``usage.completion_tokens`` with reasoning tokens (if
+    reported) mapped to ``thinking_tokens``.
+    """
+    model, effective_temperature, effective_max_tokens = _resolve_call_params(
+        model_params, max_tokens, temperature
+    )
+
+    create_kwargs: dict[str, Any] = {
+        "model": model,
+        "messages": messages,
+        "temperature": effective_temperature,
+    }
+    if effective_max_tokens is not None:
+        create_kwargs["max_tokens"] = effective_max_tokens
+
+    if stream:
+        # Ask OpenAI to include the usage chunk at the end of the stream so token
+        # counts are provider-reported rather than estimated when possible.
+        create_kwargs["stream"] = True
+        create_kwargs["stream_options"] = {"include_usage": True}
+
+        async def _openai_stream() -> AsyncIterator[ChatExecutionChunk]:
+            response = await client.chat.completions.create(**create_kwargs)
+            collected: list[str] = []
+            finish_reason: str | None = None
+            usage = None
+            async for chunk in response:
+                if getattr(chunk, "usage", None):
+                    usage = chunk.usage
+                choices = getattr(chunk, "choices", None)
+                if choices:
+                    choice = choices[0]
+                    delta_obj = getattr(choice, "delta", None)
+                    delta_text = getattr(delta_obj, "content", None) if delta_obj else None
+                    if getattr(choice, "finish_reason", None):
+                        finish_reason = choice.finish_reason
+                    if delta_text:
+                        collected.append(delta_text)
+                        yield ChatExecutionChunk(delta=delta_text, model=model)
+            full_text = "".join(collected)
+            input_tokens, output_tokens, thinking_tokens, estimated = _normalize_openai_usage(
+                usage, model, messages, full_text
+            )
+            yield ChatExecutionChunk(
+                delta="",
+                model=model,
+                done=True,
+                finish_reason=finish_reason,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                thinking_tokens=thinking_tokens,
+                usage_estimated=estimated,
+            )
+
+        return _openai_stream()
+
+    response = await client.chat.completions.create(**create_kwargs)
+    choice = response.choices[0]
+    content = getattr(choice.message, "content", None) or ""
+    finish_reason = getattr(choice, "finish_reason", None)
+    input_tokens, output_tokens, thinking_tokens, estimated = _normalize_openai_usage(
+        getattr(response, "usage", None), model, messages, content
+    )
+    return ChatExecutionResult(
+        content=content,
+        model=model,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        thinking_tokens=thinking_tokens,
+        finish_reason=finish_reason,
+        usage_estimated=estimated,
+    )
+
+
+def _normalize_openai_usage(
+    usage: Any, model: str | None, messages: list[dict[str, Any]], output_text: str
+) -> tuple[int, int, int, bool]:
+    """Normalize OpenAI usage to (input, output, thinking, estimated)."""
+    if usage is not None:
+        input_tokens = getattr(usage, "prompt_tokens", 0) or 0
+        output_tokens = getattr(usage, "completion_tokens", 0) or 0
+        thinking_tokens = 0
+        details = getattr(usage, "completion_tokens_details", None)
+        if details is not None:
+            thinking_tokens = getattr(details, "reasoning_tokens", 0) or 0
+        return input_tokens, output_tokens, thinking_tokens, False
+    # Fallback: provider reported no usage -> estimate (Req 2.12).
+    input_tokens, output_tokens = _estimate_token_counts(
+        model, _join_message_text(messages), output_text
+    )
+    return input_tokens, output_tokens, 0, True
+
+
+async def _execute_anthropic(
+    client: Any,
+    model_params: dict[str, Any],
+    messages: list[dict[str, Any]],
+    *,
+    stream: bool = False,
+    max_tokens: int | None = None,
+    temperature: float | None = None,
+) -> "ChatExecutionResult | AsyncIterator[ChatExecutionChunk]":
+    """Execute via ``client.messages.create``.
+
+    Anthropic takes NO system role inside ``messages`` and instead requires the
+    top-level ``system`` parameter, so the system-role message is lifted out of
+    the messages list into ``system`` (Req 2.11). Token usage is read from
+    ``usage.input_tokens``/``usage.output_tokens``.
+    """
+    model, effective_temperature, effective_max_tokens = _resolve_call_params(
+        model_params, max_tokens, temperature
+    )
+    # Anthropic REQUIRES max_tokens; fall back to a default when none was provided.
+    if effective_max_tokens is None:
+        effective_max_tokens = model_params.get("max_tokens", _DEFAULT_MAX_TOKENS)
+
+    system_text, non_system_messages = _split_system_messages(messages)
+
+    create_kwargs: dict[str, Any] = {
+        "model": model,
+        "messages": non_system_messages,
+        "max_tokens": effective_max_tokens,
+        "temperature": effective_temperature,
+    }
+    if system_text is not None:
+        create_kwargs["system"] = system_text
+
+    if stream:
+
+        async def _anthropic_stream() -> AsyncIterator[ChatExecutionChunk]:
+            collected: list[str] = []
+            usage = None
+            finish_reason: str | None = None
+            # The streaming helper yields text deltas and exposes the final message.
+            async with client.messages.stream(**create_kwargs) as stream_mgr:
+                async for text in stream_mgr.text_stream:
+                    if text:
+                        collected.append(text)
+                        yield ChatExecutionChunk(delta=text, model=model)
+                final_message = await stream_mgr.get_final_message()
+                usage = getattr(final_message, "usage", None)
+                finish_reason = getattr(final_message, "stop_reason", None)
+            full_text = "".join(collected)
+            input_tokens, output_tokens, thinking_tokens, estimated = _normalize_anthropic_usage(
+                usage, model, messages, full_text
+            )
+            yield ChatExecutionChunk(
+                delta="",
+                model=model,
+                done=True,
+                finish_reason=finish_reason,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                thinking_tokens=thinking_tokens,
+                usage_estimated=estimated,
+            )
+
+        return _anthropic_stream()
+
+    response = await client.messages.create(**create_kwargs)
+    content = _extract_anthropic_text(response)
+    finish_reason = getattr(response, "stop_reason", None)
+    input_tokens, output_tokens, thinking_tokens, estimated = _normalize_anthropic_usage(
+        getattr(response, "usage", None), model, messages, content
+    )
+    return ChatExecutionResult(
+        content=content,
+        model=model,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        thinking_tokens=thinking_tokens,
+        finish_reason=finish_reason,
+        usage_estimated=estimated,
+    )
+
+
+def _extract_anthropic_text(response: Any) -> str:
+    """Concatenate text from Anthropic content blocks."""
+    blocks = getattr(response, "content", None)
+    if not blocks:
+        return ""
+    parts: list[str] = []
+    for block in blocks:
+        text = getattr(block, "text", None)
+        if text:
+            parts.append(text)
+    return "".join(parts)
+
+
+def _normalize_anthropic_usage(
+    usage: Any, model: str | None, messages: list[dict[str, Any]], output_text: str
+) -> tuple[int, int, int, bool]:
+    """Normalize Anthropic usage to (input, output, thinking, estimated)."""
+    if usage is not None:
+        input_tokens = getattr(usage, "input_tokens", 0) or 0
+        output_tokens = getattr(usage, "output_tokens", 0) or 0
+        # Anthropic does not separately report thinking tokens in usage today.
+        thinking_tokens = 0
+        return input_tokens, output_tokens, thinking_tokens, False
+    input_tokens, output_tokens = _estimate_token_counts(
+        model, _join_message_text(messages), output_text
+    )
+    return input_tokens, output_tokens, 0, True
+
+
+def _messages_to_gemini_contents(messages: list[dict[str, Any]]) -> str:
+    """Flatten messages into a single prompt for the Gemini SDK.
+
+    Gemini sets a system instruction at model construction time; the model passed
+    here is already constructed, so the system-role message is kept INSIDE the
+    prompt content (it is NOT dropped) (Req 2.11). Roles are prefixed so the model
+    retains the system/user distinction.
+    """
+    parts: list[str] = []
+    for msg in messages:
+        role = msg.get("role", "user")
+        content = msg.get("content")
+        if content is None:
+            continue
+        text = content if isinstance(content, str) else str(content)
+        if role == "system":
+            parts.append(f"System: {text}")
+        elif role == "assistant":
+            parts.append(f"Assistant: {text}")
+        else:
+            parts.append(text)
+    return "\n\n".join(parts)
+
+
+async def _execute_gemini(
+    client: Any,
+    model_params: dict[str, Any],
+    messages: list[dict[str, Any]],
+    *,
+    stream: bool = False,
+    max_tokens: int | None = None,
+    temperature: float | None = None,
+) -> "ChatExecutionResult | AsyncIterator[ChatExecutionChunk]":
+    """Execute via Gemini ``GenerativeModel.generate_content_async``.
+
+    Gemini keeps the system role inside the prompt / system instruction (Req 2.11).
+    Token usage is read from ``usage_metadata`` (``prompt_token_count`` /
+    ``candidates_token_count`` / ``thoughts_token_count``).
+    """
+    model, effective_temperature, effective_max_tokens = _resolve_call_params(
+        model_params, max_tokens, temperature
+    )
+    contents = _messages_to_gemini_contents(messages)
+
+    generation_config: dict[str, Any] = {"temperature": effective_temperature}
+    if effective_max_tokens is not None:
+        generation_config["max_output_tokens"] = effective_max_tokens
+
+    if stream:
+
+        async def _gemini_stream() -> AsyncIterator[ChatExecutionChunk]:
+            collected: list[str] = []
+            usage_metadata = None
+            finish_reason: str | None = None
+            response = await client.generate_content_async(
+                contents, generation_config=generation_config, stream=True
+            )
+            async for chunk in response:
+                text = getattr(chunk, "text", None)
+                if getattr(chunk, "usage_metadata", None):
+                    usage_metadata = chunk.usage_metadata
+                if text:
+                    collected.append(text)
+                    yield ChatExecutionChunk(delta=text, model=model)
+            full_text = "".join(collected)
+            input_tokens, output_tokens, thinking_tokens, estimated = _normalize_gemini_usage(
+                usage_metadata, model, messages, full_text
+            )
+            yield ChatExecutionChunk(
+                delta="",
+                model=model,
+                done=True,
+                finish_reason=finish_reason,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                thinking_tokens=thinking_tokens,
+                usage_estimated=estimated,
+            )
+
+        return _gemini_stream()
+
+    response = await client.generate_content_async(contents, generation_config=generation_config)
+    content = getattr(response, "text", None) or ""
+    finish_reason = None
+    input_tokens, output_tokens, thinking_tokens, estimated = _normalize_gemini_usage(
+        getattr(response, "usage_metadata", None), model, messages, content
+    )
+    return ChatExecutionResult(
+        content=content,
+        model=model,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        thinking_tokens=thinking_tokens,
+        finish_reason=finish_reason,
+        usage_estimated=estimated,
+    )
+
+
+def _normalize_gemini_usage(
+    usage_metadata: Any, model: str | None, messages: list[dict[str, Any]], output_text: str
+) -> tuple[int, int, int, bool]:
+    """Normalize Gemini ``usage_metadata`` to (input, output, thinking, estimated)."""
+    if usage_metadata is not None:
+        input_tokens = getattr(usage_metadata, "prompt_token_count", 0) or 0
+        output_tokens = getattr(usage_metadata, "candidates_token_count", 0) or 0
+        thinking_tokens = getattr(usage_metadata, "thoughts_token_count", 0) or 0
+        return input_tokens, output_tokens, thinking_tokens, False
+    input_tokens, output_tokens = _estimate_token_counts(
+        model, _join_message_text(messages), output_text
+    )
+    return input_tokens, output_tokens, 0, True

{agent_framework_lib-0.8.8.post9 → agent_framework_lib-0.8.9}/agent_framework/monitoring/__init__.py

@@ -42,6 +42,7 @@ from agent_framework.monitoring.metrics_aggregator import (
     MetricsAggregator,
     UnifiedMetricsDocument,
     ensure_unified_index_template,
+    write_unified_metrics_document,
 )
 from agent_framework.monitoring.metrics_config import (
     MetricsConfig,
@@ -161,6 +162,7 @@ __all__ = [
     "MetricsAggregator",
     "UnifiedMetricsDocument",
     "ensure_unified_index_template",
+    "write_unified_metrics_document",
     # LLM Metrics Extraction
     "LLMMetricsExtractor",
     "TraceMetricsAccumulator",