mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Merge and move scripts (#4308)
* .NET: Add Microsoft Fabric sample #3674 (#4230) Co-authored-by: Chris <66376200+crickman@users.noreply.github.com> * 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> * [Purview] Mark responses as responses and fix epoch bug for python long overflow (#4225) * .NET: Support InvokeMcpTool for declarative workflows (#4204) * Initial implementation of InvokeMcpTool in declarative workflow * Cleaned up sample implementation * Updated sample comments. * Added missing executor routing attribute * Fix PR comments. * Updated based on PR comments. * Updated based on PR comments. * Removed unnecessary using statement. * Update Python package versions to rc2 (#4258) - Bump core and azure-ai to 1.0.0rc2 - Bump preview packages to 1.0.0b260225 - Update dependencies to >=1.0.0rc2 - Add CHANGELOG entries for changes since rc1 - Update uv.lock Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * .NET: Fixing issue where OpenTelemetry span is never exported in .NET in-process workflow execution (#4196) * 1. Add reproduction test for issue #4155: workflow.run Activity never stopped in streaming OffThread path The WorkflowRunActivity_IsStopped_Streaming_OffThread test demonstrates that the workflow.run OpenTelemetry Activity created in StreamingRunEventStream.RunLoopAsync is started but never stopped when using the OffThread/Default streaming execution. The background run loop keeps running after event consumption completes, so the using Activity? declaration never disposes until explicit StopAsync() is called. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> 2. Fix workflow.run Activity never stopped in streaming OffThread execution (#4155) The workflow.run OpenTelemetry Activity in StreamingRunEventStream.RunLoopAsync was scoped to the method lifetime via 'using'. Since the run loop only exits on cancellation, the Activity was never stopped/exported until explicit disposal. Fix: Remove 'using' and explicitly dispose the Activity when the workflow reaches Idle status (all supersteps complete). A safety-net disposal in the finally block handles cancellation and error paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Add root-level workflow.session activity spanning run loop lifetime\n\nImplements two-level telemetry hierarchy per PR feedback from lokitoth:\n- workflow.session: spans the entire run loop / stream lifetime\n- workflow_invoke: per input-to-halt cycle, nested within the session\n\nThis ensures the session activity stays open across multiple turns,\nwhile individual run activities are created and disposed per cycle.\n\nAlso fixes linkedSource CancellationTokenSource disposal leak in\nStreamingRunEventStream (added using declaration)." * Address Copilot review: fix Activity/CTS disposal, rename activity, add error tag\n\n1. LockstepRunEventStream: Remove 'using' from Activity in async iterator\n and manually dispose in finally block (fixes #4155 pattern). Also dispose\n linkedSource CTS in finally to prevent leak.\n2. Tags.cs: Add ErrorMessage (\"error.message\") tag for runtime errors,\n distinct from BuildErrorMessage (\"build.error.message\").\n3. ActivityNames: Rename WorkflowRun from \"workflow_invoke\" to \"workflow.run\"\n for cross-language consistency.\n4. WorkflowTelemetryContext: Fix XML doc to say \"outer/parent span\" instead\n of \"root-level span\".\n5. ObservabilityTests: Assert WorkflowSession absence when DisableWorkflowRun\n is true.\n6. WorkflowRunActivityStopTests: Fix streaming test race by disposing\n StreamingRun before asserting activities are stopped.\n7. StreamingRunEventStream/LockstepRunEventStream: Use Tags.ErrorMessage\n instead of Tags.BuildErrorMessage for runtime error events." * Review fixes: revert workflow_invoke rename, use 'using' for linkedSource, move SessionStarted earlier\n\n- Revert ActivityNames.WorkflowRun back to \"workflow_invoke\" (OTEL semantic convention contract)\n- Use 'using' declaration for linkedSource CTS in LockstepRunEventStream (no timing sensitivity)\n- Move SessionStarted event before WaitForInputAsync in StreamingRunEventStream to match Lockstep behavior" * Improve naming and comments in WorkflowRunActivityStopTests" * Prevent session Activity.Current leak in lockstep mode, add nesting test Save and restore Activity.Current in LockstepRunEventStream.Start() so the session activity doesn't leak into caller code via AsyncLocal. Re-establish Activity.Current = sessionActivity before creating the run activity in TakeEventStreamAsync to preserve parent-child nesting. Add test verifying app activities after RunAsync are not parented under the session, and that the workflow_invoke activity nests under the session." * Fix stale XML doc: WorkflowRun -> WorkflowInvoke in ObservabilityTests --------- Co-authored-by: alliscode <bentho@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Python / .NET Samples - Restructure and Improve Samples (Feature Branc… (#4092) * Python: .NET Samples - Restructure and Improve Samples (Feature Branch) (#4091) * Moved by agent (#4094) * Fix readme links * .NET Samples - Create `04-hosting` learning path step (#4098) * Agent move * Agent reorderd * Remove A2A section from README Removed A2A section from the Getting Started README. * Agent fixed links * Fix broken sample links in durable-agents README (#4101) * Initial plan * Fix broken internal links in documentation Co-authored-by: crickman <66376200+crickman@users.noreply.github.com> * Revert template link changes; keep only durable-agents README fix Co-authored-by: crickman <66376200+crickman@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: crickman <66376200+crickman@users.noreply.github.com> * .NET Samples - Create `03-workflows` learning path step (#4102) * Fix solution project path * Python: Fix broken markdown links to repo resources (outside /docs) (#4105) * Initial plan * Fix broken markdown links to repo resources Co-authored-by: crickman <66376200+crickman@users.noreply.github.com> * Update README to rename .NET Workflows Samples section --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: crickman <66376200+crickman@users.noreply.github.com> * .NET Samples - Create `02-agents` learning path step (#4107) * .NET: Fix broken relative link in GroupChatToolApproval README (#4108) * Initial plan * Fix broken link in GroupChatToolApproval README Co-authored-by: crickman <66376200+crickman@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: crickman <66376200+crickman@users.noreply.github.com> * Update labeler configuration for workflow samples * .NET - Reorder Agents samples to start from Step01 instead of Step04 (#4110) * Fix solution * Resolve new sample paths * Move new AgentSkills and AgentWithMemory_Step04 samples * Fix link * Fix readme path * fix: update stale dotnet/samples/Durable path reference in AGENTS.md Co-authored-by: crickman <66376200+crickman@users.noreply.github.com> * Moved new sample * Update solution * Resolve merge (new sample) * Sync to new sample - FoundryAgents_Step21_BingCustomSearch * Updated README * .NET Samples - Configuration Naming Update (#4149) * .NET: Restore AzureFunctions index parity with ConsoleApps under DurableAgents samples (#4221) * Clean-up `05_host_your_agent` * Config setting consistency * Refine samples * AGENTS.md * Move new samples * Re-order samples * Move new project and fixup solution * Fixup model config * Fix up new UT project --------- Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> * Python: Fix Bedrock embedding test stub missing meta attribute (#4287) * Fix Bedrock embedding test stub missing meta attribute * Increase test coverage so gate passes * Python: (ag-ui): fix approval payloads being re-processed on subsequent conversation turns (#4232) * Fix ag-ui tool call issue * Safe json fix * Python: Update workflow orchestration samples to use AzureOpenAIResponsesClient (#4285) * Update workflow orchestration samples to use AzureOpenAIResponsesClient * Fix broken link * Move scripts to scripts folder --------- Co-authored-by: Roger Barreto <19890735+rogerbarreto@users.noreply.github.com> Co-authored-by: Chris <66376200+crickman@users.noreply.github.com> Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Rishabh Chawla <rishabhchawla1995@gmail.com> Co-authored-by: Peter Ibekwe <109177538+peibekwe@users.noreply.github.com> Co-authored-by: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> Co-authored-by: Ben Thomas <ben.thomas@microsoft.com> Co-authored-by: alliscode <bentho@microsoft.com> Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: Evan Mattson <35585003+moonbox3@users.noreply.github.com>
This commit is contained in:
westey
committed by
GitHub
Unverified
GitHub
parent
cc1ef730e3
commit
8b191de936
@@ -3,6 +3,7 @@
|
||||
import importlib.metadata
|
||||
|
||||
from ._chat_client import BedrockChatClient, BedrockChatOptions, BedrockGuardrailConfig, BedrockSettings
|
||||
from ._embedding_client import BedrockEmbeddingClient, BedrockEmbeddingOptions, BedrockEmbeddingSettings
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
@@ -12,6 +13,9 @@ except importlib.metadata.PackageNotFoundError:
|
||||
__all__ = [
|
||||
"BedrockChatClient",
|
||||
"BedrockChatOptions",
|
||||
"BedrockEmbeddingClient",
|
||||
"BedrockEmbeddingOptions",
|
||||
"BedrockEmbeddingSettings",
|
||||
"BedrockGuardrailConfig",
|
||||
"BedrockSettings",
|
||||
"__version__",
|
||||
|
||||
@@ -0,0 +1,292 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, ClassVar, Generic, TypedDict
|
||||
|
||||
from agent_framework import (
|
||||
AGENT_FRAMEWORK_USER_AGENT,
|
||||
BaseEmbeddingClient,
|
||||
Embedding,
|
||||
EmbeddingGenerationOptions,
|
||||
GeneratedEmbeddings,
|
||||
SecretString,
|
||||
UsageDetails,
|
||||
load_settings,
|
||||
)
|
||||
from agent_framework.observability import EmbeddingTelemetryLayer
|
||||
from boto3.session import Session as Boto3Session
|
||||
from botocore.client import BaseClient
|
||||
from botocore.config import Config as BotoConfig
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # type: ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypeVar # type: ignore # pragma: no cover
|
||||
|
||||
|
||||
logger = logging.getLogger("agent_framework.bedrock")
|
||||
DEFAULT_REGION = "us-east-1"
|
||||
|
||||
|
||||
class BedrockEmbeddingSettings(TypedDict, total=False):
|
||||
"""Bedrock embedding settings."""
|
||||
|
||||
region: str | None
|
||||
embedding_model_id: str | None
|
||||
access_key: SecretString | None
|
||||
secret_key: SecretString | None
|
||||
session_token: SecretString | None
|
||||
|
||||
|
||||
class BedrockEmbeddingOptions(EmbeddingGenerationOptions, total=False):
|
||||
"""Bedrock-specific embedding options.
|
||||
|
||||
Extends EmbeddingGenerationOptions with Bedrock-specific fields.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework_bedrock import BedrockEmbeddingOptions
|
||||
|
||||
options: BedrockEmbeddingOptions = {
|
||||
"model_id": "amazon.titan-embed-text-v2:0",
|
||||
"dimensions": 1024,
|
||||
"normalize": True,
|
||||
}
|
||||
"""
|
||||
|
||||
normalize: bool
|
||||
|
||||
|
||||
BedrockEmbeddingOptionsT = TypeVar(
|
||||
"BedrockEmbeddingOptionsT",
|
||||
bound=TypedDict, # type: ignore[valid-type]
|
||||
default="BedrockEmbeddingOptions",
|
||||
covariant=True,
|
||||
)
|
||||
|
||||
|
||||
class RawBedrockEmbeddingClient(
|
||||
BaseEmbeddingClient[str, list[float], BedrockEmbeddingOptionsT],
|
||||
Generic[BedrockEmbeddingOptionsT],
|
||||
):
|
||||
"""Raw Bedrock embedding client without telemetry.
|
||||
|
||||
Keyword Args:
|
||||
model_id: The Bedrock embedding model ID (e.g. "amazon.titan-embed-text-v2:0").
|
||||
Can also be set via environment variable BEDROCK_EMBEDDING_MODEL_ID.
|
||||
region: AWS region. Will try to load from BEDROCK_REGION env var,
|
||||
if not set, the regular Boto3 configuration/loading applies
|
||||
(which may include other env vars, config files, or instance metadata).
|
||||
access_key: AWS access key for manual credential injection.
|
||||
secret_key: AWS secret key paired with access_key.
|
||||
session_token: AWS session token for temporary credentials.
|
||||
client: Preconfigured Bedrock runtime client.
|
||||
boto3_session: Custom boto3 session used to build the runtime client.
|
||||
env_file_path: Path to .env file for settings.
|
||||
env_file_encoding: Encoding for .env file.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
region: str | None = None,
|
||||
model_id: str | None = None,
|
||||
access_key: str | None = None,
|
||||
secret_key: str | None = None,
|
||||
session_token: str | None = None,
|
||||
client: BaseClient | None = None,
|
||||
boto3_session: Boto3Session | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize a raw Bedrock embedding client."""
|
||||
settings = load_settings(
|
||||
BedrockEmbeddingSettings,
|
||||
env_prefix="BEDROCK_",
|
||||
required_fields=["embedding_model_id"],
|
||||
region=region,
|
||||
embedding_model_id=model_id,
|
||||
access_key=access_key,
|
||||
secret_key=secret_key,
|
||||
session_token=session_token,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
resolved_region = settings.get("region") or DEFAULT_REGION
|
||||
|
||||
if client is None:
|
||||
if not boto3_session:
|
||||
session_kwargs: dict[str, Any] = {}
|
||||
if region := settings.get("region"):
|
||||
session_kwargs["region_name"] = region
|
||||
if (access_key := settings.get("access_key")) and (secret_key := settings.get("secret_key")):
|
||||
session_kwargs["aws_access_key_id"] = access_key.get_secret_value() # type: ignore[union-attr]
|
||||
session_kwargs["aws_secret_access_key"] = secret_key.get_secret_value() # type: ignore[union-attr]
|
||||
if session_token := settings.get("session_token"):
|
||||
session_kwargs["aws_session_token"] = session_token.get_secret_value() # type: ignore[union-attr]
|
||||
boto3_session = Boto3Session(**session_kwargs)
|
||||
client = boto3_session.client(
|
||||
"bedrock-runtime",
|
||||
region_name=boto3_session.region_name or resolved_region,
|
||||
config=BotoConfig(user_agent_extra=AGENT_FRAMEWORK_USER_AGENT),
|
||||
)
|
||||
|
||||
self._bedrock_client = client
|
||||
self.model_id = settings["embedding_model_id"] # type: ignore[assignment]
|
||||
self.region = resolved_region
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def service_url(self) -> str:
|
||||
"""Get the URL of the service."""
|
||||
return str(self._bedrock_client.meta.endpoint_url)
|
||||
|
||||
async def get_embeddings(
|
||||
self,
|
||||
values: Sequence[str],
|
||||
*,
|
||||
options: BedrockEmbeddingOptionsT | None = None,
|
||||
) -> GeneratedEmbeddings[list[float]]:
|
||||
"""Call the Bedrock invoke_model API for embeddings.
|
||||
|
||||
Uses the Amazon Titan Embeddings model format. Each value is embedded
|
||||
individually since Titan's invoke_model API accepts one input at a time.
|
||||
|
||||
Args:
|
||||
values: The text values to generate embeddings for.
|
||||
options: Optional embedding generation options.
|
||||
|
||||
Returns:
|
||||
Generated embeddings with usage metadata.
|
||||
|
||||
Raises:
|
||||
ValueError: If model_id is not provided or values is empty.
|
||||
"""
|
||||
if not values:
|
||||
return GeneratedEmbeddings([], options=options)
|
||||
|
||||
opts: dict[str, Any] = dict(options) if options else {}
|
||||
model = opts.get("model_id") or self.model_id
|
||||
if not model:
|
||||
raise ValueError("model_id is required")
|
||||
|
||||
embedding_results = await asyncio.gather(
|
||||
*(self._generate_embedding_for_text(opts, model, text) for text in values)
|
||||
)
|
||||
embeddings: list[Embedding[list[float]]] = []
|
||||
total_input_tokens = 0
|
||||
for embedding, input_tokens in embedding_results:
|
||||
embeddings.append(embedding)
|
||||
total_input_tokens += input_tokens
|
||||
|
||||
usage_dict: UsageDetails | None = None
|
||||
if total_input_tokens > 0:
|
||||
usage_dict = {"input_token_count": total_input_tokens}
|
||||
|
||||
return GeneratedEmbeddings(embeddings, options=options, usage=usage_dict)
|
||||
|
||||
async def _generate_embedding_for_text(
|
||||
self,
|
||||
opts: dict[str, Any],
|
||||
model: str,
|
||||
text: str,
|
||||
) -> tuple[Embedding[list[float]], int]:
|
||||
body: dict[str, Any] = {"inputText": text}
|
||||
if dimensions := opts.get("dimensions"):
|
||||
body["dimensions"] = dimensions
|
||||
if (normalize := opts.get("normalize")) is not None:
|
||||
body["normalize"] = normalize
|
||||
|
||||
response = await asyncio.to_thread(
|
||||
self._bedrock_client.invoke_model,
|
||||
modelId=model,
|
||||
contentType="application/json",
|
||||
accept="application/json",
|
||||
body=json.dumps(body),
|
||||
)
|
||||
|
||||
response_body = json.loads(response["body"].read())
|
||||
embedding = Embedding(
|
||||
vector=response_body["embedding"],
|
||||
dimensions=len(response_body["embedding"]),
|
||||
model_id=model,
|
||||
)
|
||||
input_tokens = int(response_body.get("inputTextTokenCount", 0))
|
||||
return embedding, input_tokens
|
||||
|
||||
|
||||
class BedrockEmbeddingClient(
|
||||
EmbeddingTelemetryLayer[str, list[float], BedrockEmbeddingOptionsT],
|
||||
RawBedrockEmbeddingClient[BedrockEmbeddingOptionsT],
|
||||
Generic[BedrockEmbeddingOptionsT],
|
||||
):
|
||||
"""Bedrock embedding client with telemetry support.
|
||||
|
||||
Uses the Amazon Titan Embeddings model via Bedrock's invoke_model API.
|
||||
|
||||
Keyword Args:
|
||||
model_id: The Bedrock embedding model ID (e.g. "amazon.titan-embed-text-v2:0").
|
||||
Can also be set via environment variable BEDROCK_EMBEDDING_MODEL_ID.
|
||||
region: AWS region. Defaults to "us-east-1".
|
||||
Can also be set via environment variable BEDROCK_REGION.
|
||||
access_key: AWS access key for manual credential injection.
|
||||
secret_key: AWS secret key paired with access_key.
|
||||
session_token: AWS session token for temporary credentials.
|
||||
client: Preconfigured Bedrock runtime client.
|
||||
boto3_session: Custom boto3 session used to build the runtime client.
|
||||
env_file_path: Path to .env file for settings.
|
||||
env_file_encoding: Encoding for .env file.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework_bedrock import BedrockEmbeddingClient
|
||||
|
||||
# Using default AWS credentials
|
||||
client = BedrockEmbeddingClient(
|
||||
model_id="amazon.titan-embed-text-v2:0",
|
||||
)
|
||||
|
||||
# Generate embeddings
|
||||
result = await client.get_embeddings(["Hello, world!"])
|
||||
print(result[0].vector)
|
||||
"""
|
||||
|
||||
OTEL_PROVIDER_NAME: ClassVar[str] = "aws.bedrock" # type: ignore[reportIncompatibleVariableOverride, misc]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
region: str | None = None,
|
||||
model_id: str | None = None,
|
||||
access_key: str | None = None,
|
||||
secret_key: str | None = None,
|
||||
session_token: str | None = None,
|
||||
client: BaseClient | None = None,
|
||||
boto3_session: Boto3Session | None = None,
|
||||
otel_provider_name: str | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize a Bedrock embedding client."""
|
||||
super().__init__(
|
||||
region=region,
|
||||
model_id=model_id,
|
||||
access_key=access_key,
|
||||
secret_key=secret_key,
|
||||
session_token=session_token,
|
||||
client=client,
|
||||
boto3_session=boto3_session,
|
||||
otel_provider_name=otel_provider_name,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -4,7 +4,7 @@ description = "Amazon Bedrock integration for Microsoft Agent Framework."
|
||||
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
version = "1.0.0b260219"
|
||||
version = "1.0.0b260225"
|
||||
license-files = ["LICENSE"]
|
||||
urls.homepage = "https://aka.ms/agent-framework"
|
||||
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
|
||||
@@ -23,12 +23,11 @@ classifiers = [
|
||||
"Typing :: Typed",
|
||||
]
|
||||
dependencies = [
|
||||
"agent-framework-core>=1.0.0rc1",
|
||||
"agent-framework-core>=1.0.0rc2",
|
||||
"boto3>=1.35.0,<2.0.0",
|
||||
"botocore>=1.35.0,<2.0.0",
|
||||
]
|
||||
|
||||
|
||||
[tool.uv]
|
||||
prerelease = "if-necessary-or-explicit"
|
||||
environments = [
|
||||
@@ -46,6 +45,9 @@ addopts = "-ra -q -r fEX"
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
filterwarnings = []
|
||||
markers = [
|
||||
"integration: marks tests as integration tests that require external services",
|
||||
]
|
||||
timeout = 120
|
||||
|
||||
[tool.ruff]
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from agent_framework import Embedding, GeneratedEmbeddings
|
||||
|
||||
from agent_framework_bedrock import BedrockEmbeddingClient, BedrockEmbeddingOptions
|
||||
|
||||
|
||||
class _StubBedrockEmbeddingRuntime:
|
||||
"""Stub for the Bedrock runtime client that handles invoke_model for embeddings."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
self.meta = MagicMock(endpoint_url="https://bedrock-runtime.us-west-2.amazonaws.com")
|
||||
|
||||
def invoke_model(self, **kwargs: Any) -> dict[str, Any]:
|
||||
self.calls.append(kwargs)
|
||||
body = json.loads(kwargs.get("body", "{}"))
|
||||
# Simulate Titan embedding response
|
||||
dimensions = body.get("dimensions", 3)
|
||||
return {
|
||||
"body": MagicMock(
|
||||
read=lambda: json.dumps({
|
||||
"embedding": [0.1 * (i + 1) for i in range(dimensions)],
|
||||
"inputTextTokenCount": 5,
|
||||
}).encode()
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
async def test_bedrock_embedding_construction() -> None:
|
||||
"""Test construction with explicit parameters."""
|
||||
stub = _StubBedrockEmbeddingRuntime()
|
||||
client = BedrockEmbeddingClient(
|
||||
model_id="amazon.titan-embed-text-v2:0",
|
||||
region="us-west-2",
|
||||
client=stub,
|
||||
)
|
||||
assert client.model_id == "amazon.titan-embed-text-v2:0"
|
||||
assert client.region == "us-west-2"
|
||||
|
||||
|
||||
async def test_bedrock_embedding_construction_missing_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Test that missing model_id raises an error."""
|
||||
monkeypatch.delenv("BEDROCK_EMBEDDING_MODEL_ID", raising=False)
|
||||
from agent_framework.exceptions import SettingNotFoundError
|
||||
|
||||
with pytest.raises(SettingNotFoundError):
|
||||
BedrockEmbeddingClient(region="us-west-2")
|
||||
|
||||
|
||||
async def test_bedrock_embedding_get_embeddings() -> None:
|
||||
"""Test generating embeddings via the Bedrock invoke_model API."""
|
||||
stub = _StubBedrockEmbeddingRuntime()
|
||||
client = BedrockEmbeddingClient(
|
||||
model_id="amazon.titan-embed-text-v2:0",
|
||||
region="us-west-2",
|
||||
client=stub,
|
||||
)
|
||||
|
||||
result = await client.get_embeddings(["hello", "world"])
|
||||
|
||||
assert isinstance(result, GeneratedEmbeddings)
|
||||
assert len(result) == 2
|
||||
assert len(result[0].vector) == 3
|
||||
assert len(result[1].vector) == 3
|
||||
assert result[0].model_id == "amazon.titan-embed-text-v2:0"
|
||||
assert result.usage == {"input_token_count": 10}
|
||||
|
||||
# Two calls since Titan processes one input at a time
|
||||
assert len(stub.calls) == 2
|
||||
call_texts = {json.loads(call["body"])["inputText"] for call in stub.calls}
|
||||
assert call_texts == {"hello", "world"}
|
||||
|
||||
|
||||
async def test_bedrock_embedding_get_embeddings_empty_input() -> None:
|
||||
"""Test generating embeddings with empty input."""
|
||||
stub = _StubBedrockEmbeddingRuntime()
|
||||
client = BedrockEmbeddingClient(
|
||||
model_id="amazon.titan-embed-text-v2:0",
|
||||
region="us-west-2",
|
||||
client=stub,
|
||||
)
|
||||
|
||||
result = await client.get_embeddings([])
|
||||
|
||||
assert isinstance(result, GeneratedEmbeddings)
|
||||
assert len(result) == 0
|
||||
assert len(stub.calls) == 0
|
||||
|
||||
|
||||
async def test_bedrock_embedding_get_embeddings_with_options() -> None:
|
||||
"""Test generating embeddings with custom options."""
|
||||
stub = _StubBedrockEmbeddingRuntime()
|
||||
client = BedrockEmbeddingClient(
|
||||
model_id="amazon.titan-embed-text-v2:0",
|
||||
region="us-west-2",
|
||||
client=stub,
|
||||
)
|
||||
|
||||
options: BedrockEmbeddingOptions = {
|
||||
"dimensions": 5,
|
||||
"normalize": True,
|
||||
}
|
||||
result = await client.get_embeddings(["hello"], options=options)
|
||||
|
||||
assert len(result) == 1
|
||||
assert len(result[0].vector) == 5
|
||||
|
||||
body = json.loads(stub.calls[0]["body"])
|
||||
assert body["dimensions"] == 5
|
||||
assert body["normalize"] is True
|
||||
|
||||
|
||||
async def test_bedrock_embedding_get_embeddings_no_model_raises() -> None:
|
||||
"""Test that missing model_id at call time raises ValueError."""
|
||||
stub = _StubBedrockEmbeddingRuntime()
|
||||
client = BedrockEmbeddingClient(
|
||||
model_id="amazon.titan-embed-text-v2:0",
|
||||
region="us-west-2",
|
||||
client=stub,
|
||||
)
|
||||
client.model_id = None # type: ignore[assignment]
|
||||
|
||||
with pytest.raises(ValueError, match="model_id is required"):
|
||||
await client.get_embeddings(["hello"])
|
||||
|
||||
|
||||
async def test_bedrock_embedding_default_region() -> None:
|
||||
"""Test that default region is us-east-1."""
|
||||
stub = _StubBedrockEmbeddingRuntime()
|
||||
client = BedrockEmbeddingClient(
|
||||
model_id="amazon.titan-embed-text-v2:0",
|
||||
client=stub,
|
||||
)
|
||||
assert client.region == "us-east-1"
|
||||
|
||||
|
||||
# region: Integration Tests
|
||||
|
||||
skip_if_bedrock_embedding_integration_tests_disabled = pytest.mark.skipif(
|
||||
os.getenv("BEDROCK_EMBEDDING_MODEL_ID", "") in ("", "test-model")
|
||||
or not (os.getenv("AWS_ACCESS_KEY_ID") or os.getenv("BEDROCK_ACCESS_KEY")),
|
||||
reason="No real Bedrock embedding model or AWS credentials provided; skipping integration tests.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.flaky
|
||||
@pytest.mark.integration
|
||||
@skip_if_bedrock_embedding_integration_tests_disabled
|
||||
async def test_bedrock_embedding_integration() -> None:
|
||||
"""Integration test for Bedrock embedding client."""
|
||||
client = BedrockEmbeddingClient()
|
||||
result = await client.get_embeddings(["Hello, world!", "How are you?"])
|
||||
|
||||
assert isinstance(result, GeneratedEmbeddings)
|
||||
assert len(result) == 2
|
||||
for embedding in result:
|
||||
assert isinstance(embedding, Embedding)
|
||||
assert isinstance(embedding.vector, list)
|
||||
assert len(embedding.vector) > 0
|
||||
assert all(isinstance(v, float) for v in embedding.vector)
|
||||
Reference in New Issue
Block a user