mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-28 17:48:13 +00:00
2485ed93d6
- update imports and include handling for chromadb v1.1.0 - fix mypy and typing_compat issues (required, typeddict, voyageai) - refine embedderconfig typing and allow base provider instances - handle mem0 as special case for external memory storage - bump tools and clean up redundant deps
54 lines
1.6 KiB
Python
54 lines
1.6 KiB
Python
"""Amazon Bedrock embeddings provider."""
|
|
|
|
from typing import Any
|
|
|
|
from chromadb.utils.embedding_functions.amazon_bedrock_embedding_function import (
|
|
AmazonBedrockEmbeddingFunction,
|
|
)
|
|
from pydantic import Field
|
|
|
|
from crewai.rag.core.base_embeddings_provider import BaseEmbeddingsProvider
|
|
|
|
|
|
def create_aws_session() -> Any:
|
|
"""Create an AWS session for Bedrock.
|
|
|
|
Returns:
|
|
boto3.Session: AWS session object
|
|
|
|
Raises:
|
|
ImportError: If boto3 is not installed
|
|
ValueError: If AWS session creation fails
|
|
"""
|
|
try:
|
|
import boto3 # type: ignore[import]
|
|
|
|
return boto3.Session()
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"boto3 is required for amazon-bedrock embeddings. "
|
|
"Install it with: uv add boto3"
|
|
) from e
|
|
except Exception as e:
|
|
raise ValueError(
|
|
f"Failed to create AWS session for amazon-bedrock. "
|
|
f"Ensure AWS credentials are configured. Error: {e}"
|
|
) from e
|
|
|
|
|
|
class BedrockProvider(BaseEmbeddingsProvider[AmazonBedrockEmbeddingFunction]):
|
|
"""Amazon Bedrock embeddings provider."""
|
|
|
|
embedding_callable: type[AmazonBedrockEmbeddingFunction] = Field(
|
|
default=AmazonBedrockEmbeddingFunction,
|
|
description="Amazon Bedrock embedding function class",
|
|
)
|
|
model_name: str = Field(
|
|
default="amazon.titan-embed-text-v1",
|
|
description="Model name to use for embeddings",
|
|
validation_alias="BEDROCK_MODEL_NAME",
|
|
)
|
|
session: Any = Field(
|
|
default_factory=create_aws_session, description="AWS session object"
|
|
)
|