agno 0.1.2__py3-none-any.whl → 2.3.13__py3-none-any.whl

agno/knowledge/embedder/aws_bedrock.py

@@ -0,0 +1,343 @@
+import json
+from dataclasses import dataclass
+from os import getenv
+from typing import Any, Dict, List, Optional, Tuple
+
+from agno.exceptions import AgnoError, ModelProviderError
+from agno.knowledge.embedder.base import Embedder
+from agno.utils.log import log_error, log_warning
+
+try:
+    from boto3 import client as AwsClient
+    from boto3.session import Session
+    from botocore.exceptions import ClientError
+except ImportError:
+    log_error("`boto3` not installed. Please install it via `pip install boto3`.")
+    raise
+
+try:
+    import aioboto3
+    from aioboto3.session import Session as AioSession
+except ImportError:
+    log_error("`aioboto3` not installed. Please install it via `pip install aioboto3`.")
+    aioboto3 = None
+    AioSession = None
+
+
+@dataclass
+class AwsBedrockEmbedder(Embedder):
+    """
+    AWS Bedrock embedder.
+
+    To use this embedder, you need to either:
+    1. Set the following environment variables:
+       - AWS_ACCESS_KEY_ID
+       - AWS_SECRET_ACCESS_KEY
+       - AWS_REGION
+    2. Or provide a boto3 Session object
+
+    Args:
+        id (str): The model ID to use. Default is 'cohere.embed-multilingual-v3'.
+        dimensions (Optional[int]): The dimensions of the embeddings. Default is 1024.
+        input_type (str): Prepends special tokens to differentiate types. Options:
+            'search_document', 'search_query', 'classification', 'clustering'. Default is 'search_query'.
+        truncate (Optional[str]): How to handle inputs longer than the maximum token length.
+            Options: 'NONE', 'START', 'END'. Default is 'NONE'.
+        embedding_types (Optional[List[str]]): Types of embeddings to return. Options:
+            'float', 'int8', 'uint8', 'binary', 'ubinary'. Default is ['float'].
+        aws_region (Optional[str]): The AWS region to use.
+        aws_access_key_id (Optional[str]): The AWS access key ID to use.
+        aws_secret_access_key (Optional[str]): The AWS secret access key to use.
+        session (Optional[Session]): A boto3 Session object to use for authentication.
+        request_params (Optional[Dict[str, Any]]): Additional parameters to pass to the API requests.
+        client_params (Optional[Dict[str, Any]]): Additional parameters to pass to the boto3 client.
+    """
+
+    id: str = "cohere.embed-multilingual-v3"
+    dimensions: int = 1024  # Cohere models have 1024 dimensions by default
+    input_type: str = "search_query"
+    truncate: Optional[str] = None  # 'NONE', 'START', or 'END'
+    # 'float', 'int8', 'uint8', etc.
+    embedding_types: Optional[List[str]] = None
+
+    aws_region: Optional[str] = None
+    aws_access_key_id: Optional[str] = None
+    aws_secret_access_key: Optional[str] = None
+    session: Optional[Session] = None
+
+    request_params: Optional[Dict[str, Any]] = None
+    client_params: Optional[Dict[str, Any]] = None
+    client: Optional[AwsClient] = None
+
+    def __post_init__(self):
+        if self.enable_batch:
+            log_warning("AwsBedrockEmbedder does not support batch embeddings, setting enable_batch to False")
+            self.enable_batch = False
+
+    def get_client(self) -> AwsClient:
+        """
+        Returns an AWS Bedrock client.
+
+        Returns:
+            AwsClient: An instance of the AWS Bedrock client.
+        """
+        if self.client is not None:
+            return self.client
+
+        if self.session:
+            self.client = self.session.client("bedrock-runtime")
+            return self.client
+
+        self.aws_access_key_id = self.aws_access_key_id or getenv("AWS_ACCESS_KEY_ID")
+        self.aws_secret_access_key = self.aws_secret_access_key or getenv("AWS_SECRET_ACCESS_KEY")
+        self.aws_region = self.aws_region or getenv("AWS_REGION")
+
+        if not self.aws_access_key_id or not self.aws_secret_access_key:
+            raise AgnoError(
+                message="AWS credentials not found. Please set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables or provide a boto3 session.",
+                status_code=400,
+            )
+
+        self.client = AwsClient(
+            service_name="bedrock-runtime",
+            region_name=self.aws_region,
+            aws_access_key_id=self.aws_access_key_id,
+            aws_secret_access_key=self.aws_secret_access_key,
+            **(self.client_params or {}),
+        )
+        return self.client
+
+    def get_async_client(self):
+        """
+        Returns an async AWS Bedrock client using aioboto3.
+
+        Returns:
+            An aioboto3 bedrock-runtime client context manager.
+        """
+        if aioboto3 is None:
+            raise AgnoError(
+                message="aioboto3 not installed. Please install it via `pip install aioboto3`.",
+                status_code=400,
+            )
+
+        if self.session:
+            # Convert boto3 session to aioboto3 session
+            aio_session = aioboto3.Session(
+                aws_access_key_id=self.session.get_credentials().access_key,
+                aws_secret_access_key=self.session.get_credentials().secret_key,
+                aws_session_token=self.session.get_credentials().token,
+                region_name=self.session.region_name,
+            )
+        else:
+            self.aws_access_key_id = self.aws_access_key_id or getenv("AWS_ACCESS_KEY_ID")
+            self.aws_secret_access_key = self.aws_secret_access_key or getenv("AWS_SECRET_ACCESS_KEY")
+            self.aws_region = self.aws_region or getenv("AWS_REGION")
+
+            if not self.aws_access_key_id or not self.aws_secret_access_key:
+                raise AgnoError(
+                    message="AWS credentials not found. Please set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables or provide a boto3 session.",
+                    status_code=400,
+                )
+
+            aio_session = aioboto3.Session(
+                aws_access_key_id=self.aws_access_key_id,
+                aws_secret_access_key=self.aws_secret_access_key,
+                region_name=self.aws_region,
+            )
+
+        return aio_session.client("bedrock-runtime", **(self.client_params or {}))
+
+    def _format_request_body(self, text: str) -> str:
+        """
+        Format the request body for the embedder.
+
+        Args:
+            text (str): The text to embed.
+
+        Returns:
+            str: The formatted request body as a JSON string.
+        """
+        request_body = {
+            "texts": [text],
+            "input_type": self.input_type,
+        }
+
+        if self.truncate:
+            request_body["truncate"] = self.truncate
+
+        if self.embedding_types:
+            request_body["embedding_types"] = self.embedding_types
+
+        # Add additional request parameters if provided
+        if self.request_params:
+            request_body.update(self.request_params)
+
+        return json.dumps(request_body)
+
+    def response(self, text: str) -> Dict[str, Any]:
+        """
+        Get embeddings from AWS Bedrock for the given text.
+
+        Args:
+            text (str): The text to embed.
+
+        Returns:
+            Dict[str, Any]: The response from the API.
+        """
+        try:
+            body = self._format_request_body(text)
+            response = self.get_client().invoke_model(
+                modelId=self.id,
+                body=body,
+                contentType="application/json",
+                accept="application/json",
+            )
+            response_body = json.loads(response["body"].read().decode("utf-8"))
+            return response_body
+        except ClientError as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e.response), model_name="AwsBedrockEmbedder", model_id=self.id) from e
+        except Exception as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e), model_name="AwsBedrockEmbedder", model_id=self.id) from e
+
+    def get_embedding(self, text: str) -> List[float]:
+        """
+        Get embeddings for the given text.
+
+        Args:
+            text (str): The text to embed.
+
+        Returns:
+            List[float]: The embedding vector.
+        """
+        response = self.response(text=text)
+        try:
+            # Check if response contains embeddings or embeddings by type
+            if "embeddings" in response:
+                if isinstance(response["embeddings"], list):
+                    # Default 'float' embeddings response format
+                    return response["embeddings"][0]
+                elif isinstance(response["embeddings"], dict):
+                    # If embeddings_types parameter was used, select float embeddings
+                    if "float" in response["embeddings"]:
+                        return response["embeddings"]["float"][0]
+                    # Fallback to the first available embedding type
+                    for embedding_type in response["embeddings"]:
+                        return response["embeddings"][embedding_type][0]
+            log_warning("No embeddings found in response")
+            return []
+        except Exception as e:
+            log_warning(f"Error extracting embeddings: {e}")
+            return []
+
+    def get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict[str, Any]]]:
+        """
+        Get embeddings and usage information for the given text.
+
+        Args:
+            text (str): The text to embed.
+
+        Returns:
+            Tuple[List[float], Optional[Dict[str, Any]]]: The embedding vector and usage information.
+        """
+        response = self.response(text=text)
+
+        embedding: List[float] = []
+        # Extract embeddings
+        if "embeddings" in response:
+            if isinstance(response["embeddings"], list):
+                embedding = response["embeddings"][0]
+            elif isinstance(response["embeddings"], dict):
+                if "float" in response["embeddings"]:
+                    embedding = response["embeddings"]["float"][0]
+                # Fallback to the first available embedding type
+                else:
+                    for embedding_type in response["embeddings"]:
+                        embedding = response["embeddings"][embedding_type][0]
+                        break
+
+        # Extract usage metrics if available
+        usage = None
+        if "usage" in response:
+            usage = response["usage"]
+
+        return embedding, usage
+
+    async def async_get_embedding(self, text: str) -> List[float]:
+        """
+        Async version of get_embedding() using native aioboto3 async client.
+        """
+        try:
+            body = self._format_request_body(text)
+            async with self.get_async_client() as client:
+                response = await client.invoke_model(
+                    modelId=self.id,
+                    body=body,
+                    contentType="application/json",
+                    accept="application/json",
+                )
+                response_body = json.loads((await response["body"].read()).decode("utf-8"))
+
+                # Extract embeddings using the same logic as get_embedding
+                if "embeddings" in response_body:
+                    if isinstance(response_body["embeddings"], list):
+                        # Default 'float' embeddings response format
+                        return response_body["embeddings"][0]
+                    elif isinstance(response_body["embeddings"], dict):
+                        # If embeddings_types parameter was used, select float embeddings
+                        if "float" in response_body["embeddings"]:
+                            return response_body["embeddings"]["float"][0]
+                        # Fallback to the first available embedding type
+                        for embedding_type in response_body["embeddings"]:
+                            return response_body["embeddings"][embedding_type][0]
+                log_warning("No embeddings found in response")
+                return []
+        except ClientError as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e.response), model_name="AwsBedrockEmbedder", model_id=self.id) from e
+        except Exception as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e), model_name="AwsBedrockEmbedder", model_id=self.id) from e
+
+    async def async_get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict[str, Any]]]:
+        """
+        Async version of get_embedding_and_usage() using native aioboto3 async client.
+        """
+        try:
+            body = self._format_request_body(text)
+            async with self.get_async_client() as client:
+                response = await client.invoke_model(
+                    modelId=self.id,
+                    body=body,
+                    contentType="application/json",
+                    accept="application/json",
+                )
+                response_body = json.loads((await response["body"].read()).decode("utf-8"))
+
+                embedding: List[float] = []
+                # Extract embeddings using the same logic as get_embedding_and_usage
+                if "embeddings" in response_body:
+                    if isinstance(response_body["embeddings"], list):
+                        embedding = response_body["embeddings"][0]
+                    elif isinstance(response_body["embeddings"], dict):
+                        if "float" in response_body["embeddings"]:
+                            embedding = response_body["embeddings"]["float"][0]
+                        # Fallback to the first available embedding type
+                        else:
+                            for embedding_type in response_body["embeddings"]:
+                                embedding = response_body["embeddings"][embedding_type][0]
+                                break
+
+                # Extract usage metrics if available
+                usage = None
+                if "usage" in response_body:
+                    usage = response_body["usage"]
+
+                return embedding, usage
+        except ClientError as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e.response), model_name="AwsBedrockEmbedder", model_id=self.id) from e
+        except Exception as e:
+            log_error(f"Unexpected error calling Bedrock API: {str(e)}")
+            raise ModelProviderError(message=str(e), model_name="AwsBedrockEmbedder", model_id=self.id) from e

agno/knowledge/embedder/azure_openai.py

@@ -0,0 +1,210 @@
+from dataclasses import dataclass
+from os import getenv
+from typing import Any, Dict, List, Optional, Tuple
+
+from typing_extensions import Literal
+
+from agno.knowledge.embedder.base import Embedder
+from agno.utils.log import logger
+
+try:
+    from openai import AsyncAzureOpenAI as AsyncAzureOpenAIClient
+    from openai import AzureOpenAI as AzureOpenAIClient
+    from openai.types.create_embedding_response import CreateEmbeddingResponse
+except ImportError:
+    raise ImportError("`openai` not installed")
+
+
+@dataclass
+class AzureOpenAIEmbedder(Embedder):
+    id: str = "text-embedding-3-small"  # This has to match the model that you deployed at the provided URL
+
+    dimensions: int = 1536
+    encoding_format: Literal["float", "base64"] = "float"
+    user: Optional[str] = None
+    api_key: Optional[str] = getenv("AZURE_EMBEDDER_OPENAI_API_KEY")
+    api_version: str = getenv("AZURE_EMBEDDER_OPENAI_API_VERSION", "2024-10-21")
+    azure_endpoint: Optional[str] = getenv("AZURE_EMBEDDER_OPENAI_ENDPOINT")
+    azure_deployment: Optional[str] = getenv("AZURE_EMBEDDER_DEPLOYMENT")
+    base_url: Optional[str] = None
+    azure_ad_token: Optional[str] = None
+    azure_ad_token_provider: Optional[Any] = None
+    organization: Optional[str] = None
+    request_params: Optional[Dict[str, Any]] = None
+    client_params: Optional[Dict[str, Any]] = None
+    openai_client: Optional[AzureOpenAIClient] = None
+    async_client: Optional[AsyncAzureOpenAIClient] = None
+
+    @property
+    def client(self) -> AzureOpenAIClient:
+        if self.openai_client:
+            return self.openai_client
+
+        _client_params: Dict[str, Any] = {}
+        if self.api_key:
+            _client_params["api_key"] = self.api_key
+        if self.api_version:
+            _client_params["api_version"] = self.api_version
+        if self.organization:
+            _client_params["organization"] = self.organization
+        if self.azure_endpoint:
+            _client_params["azure_endpoint"] = self.azure_endpoint
+        if self.azure_deployment:
+            _client_params["azure_deployment"] = self.azure_deployment
+        if self.base_url:
+            _client_params["base_url"] = self.base_url
+        if self.azure_ad_token:
+            _client_params["azure_ad_token"] = self.azure_ad_token
+        if self.azure_ad_token_provider:
+            _client_params["azure_ad_token_provider"] = self.azure_ad_token_provider
+
+        if self.client_params:
+            _client_params.update(self.client_params)
+
+        return AzureOpenAIClient(**_client_params)
+
+    @property
+    def aclient(self) -> AsyncAzureOpenAIClient:
+        if self.async_client:
+            return self.async_client
+
+        _client_params: Dict[str, Any] = {}
+        if self.api_key:
+            _client_params["api_key"] = self.api_key
+        if self.api_version:
+            _client_params["api_version"] = self.api_version
+        if self.organization:
+            _client_params["organization"] = self.organization
+        if self.azure_endpoint:
+            _client_params["azure_endpoint"] = self.azure_endpoint
+        if self.azure_deployment:
+            _client_params["azure_deployment"] = self.azure_deployment
+        if self.base_url:
+            _client_params["base_url"] = self.base_url
+        if self.azure_ad_token:
+            _client_params["azure_ad_token"] = self.azure_ad_token
+        if self.azure_ad_token_provider:
+            _client_params["azure_ad_token_provider"] = self.azure_ad_token_provider
+
+        if self.client_params:
+            _client_params.update(self.client_params)
+
+        self.async_client = AsyncAzureOpenAIClient(**_client_params)
+        return self.async_client
+
+    def _response(self, text: str) -> CreateEmbeddingResponse:
+        _request_params: Dict[str, Any] = {
+            "input": text,
+            "model": self.id,
+            "encoding_format": self.encoding_format,
+        }
+        if self.user is not None:
+            _request_params["user"] = self.user
+        if self.id.startswith("text-embedding-3"):
+            _request_params["dimensions"] = self.dimensions
+        if self.request_params:
+            _request_params.update(self.request_params)
+
+        return self.client.embeddings.create(**_request_params)
+
+    def get_embedding(self, text: str) -> List[float]:
+        response: CreateEmbeddingResponse = self._response(text=text)
+        try:
+            return response.data[0].embedding
+        except Exception as e:
+            logger.warning(e)
+            return []
+
+    def get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict]]:
+        response: CreateEmbeddingResponse = self._response(text=text)
+
+        embedding = response.data[0].embedding
+        usage = response.usage
+        return embedding, usage.model_dump()
+
+    async def _aresponse(self, text: str) -> CreateEmbeddingResponse:
+        """Async version of _response method."""
+        _request_params: Dict[str, Any] = {
+            "input": text,
+            "model": self.id,
+            "encoding_format": self.encoding_format,
+        }
+        if self.user is not None:
+            _request_params["user"] = self.user
+        if self.id.startswith("text-embedding-3"):
+            _request_params["dimensions"] = self.dimensions
+        if self.request_params:
+            _request_params.update(self.request_params)
+
+        return await self.aclient.embeddings.create(**_request_params)
+
+    async def async_get_embedding(self, text: str) -> List[float]:
+        """Async version of get_embedding using the native Azure OpenAI async client."""
+        response: CreateEmbeddingResponse = await self._aresponse(text=text)
+        try:
+            return response.data[0].embedding
+        except Exception as e:
+            logger.warning(e)
+            return []
+
+    async def async_get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict]]:
+        """Async version of get_embedding_and_usage using the native Azure OpenAI async client."""
+        response: CreateEmbeddingResponse = await self._aresponse(text=text)
+
+        embedding = response.data[0].embedding
+        usage = response.usage
+        return embedding, usage.model_dump()
+
+    async def async_get_embeddings_batch_and_usage(
+        self, texts: List[str]
+    ) -> Tuple[List[List[float]], List[Optional[Dict]]]:
+        """
+        Get embeddings and usage for multiple texts in batches.
+
+        Args:
+            texts: List of text strings to embed
+
+        Returns:
+            Tuple of (List of embedding vectors, List of usage dictionaries)
+        """
+        all_embeddings = []
+        all_usage = []
+        logger.info(f"Getting embeddings and usage for {len(texts)} texts in batches of {self.batch_size}")
+
+        for i in range(0, len(texts), self.batch_size):
+            batch_texts = texts[i : i + self.batch_size]
+
+            req: Dict[str, Any] = {
+                "input": batch_texts,
+                "model": self.id,
+                "encoding_format": self.encoding_format,
+            }
+            if self.user is not None:
+                req["user"] = self.user
+            if self.id.startswith("text-embedding-3"):
+                req["dimensions"] = self.dimensions
+            if self.request_params:
+                req.update(self.request_params)
+
+            try:
+                response: CreateEmbeddingResponse = await self.aclient.embeddings.create(**req)
+                batch_embeddings = [data.embedding for data in response.data]
+                all_embeddings.extend(batch_embeddings)
+
+                # For each embedding in the batch, add the same usage information
+                usage_dict = response.usage.model_dump() if response.usage else None
+                all_usage.extend([usage_dict] * len(batch_embeddings))
+            except Exception as e:
+                logger.warning(f"Error in async batch embedding: {e}")
+                # Fallback to individual calls for this batch
+                for text in batch_texts:
+                    try:
+                        embedding, usage = await self.async_get_embedding_and_usage(text)
+                        all_embeddings.append(embedding)
+                        all_usage.append(usage)
+                    except Exception as e2:
+                        logger.warning(f"Error in individual async embedding fallback: {e2}")
+                        all_embeddings.append([])
+                        all_usage.append(None)
+
+        return all_embeddings, all_usage

agno/{embedder → knowledge/embedder}/base.py

@@ -7,9 +7,17 @@ class Embedder:
     """Base class for managing embedders"""
 
     dimensions: Optional[int] = 1536
+    enable_batch: bool = False
+    batch_size: int = 100  # Number of texts to process in each API call
 
     def get_embedding(self, text: str) -> List[float]:
         raise NotImplementedError
 
     def get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict]]:
         raise NotImplementedError
+
+    async def async_get_embedding(self, text: str) -> List[float]:
+        raise NotImplementedError
+
+    async def async_get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict]]:
+        raise NotImplementedError