agno 2.0.0rc2__py3-none-any.whl → 2.3.0__py3-none-any.whl

agno/knowledge/embedder/google.py

@@ -3,7 +3,7 @@ from os import getenv
 from typing import Any, Dict, List, Optional, Tuple
 
 from agno.knowledge.embedder.base import Embedder
-from agno.utils.log import log_error, log_info
+from agno.utils.log import log_error, log_info, log_warning
 
 try:
     from google import genai
@@ -178,3 +178,81 @@ class GeminiEmbedder(Embedder):
         except Exception as e:
             log_error(f"Error extracting embeddings: {e}")
             return [], usage
+
+    async def async_get_embeddings_batch_and_usage(
+        self, texts: List[str]
+    ) -> Tuple[List[List[float]], List[Optional[Dict[str, Any]]]]:
+        """
+        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: List[List[float]] = []
+        all_usage: List[Optional[Dict[str, Any]]] = []
+        log_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]
+
+            # If a user provides a model id with the `models/` prefix, we need to remove it
+            _id = self.id
+            if _id.startswith("models/"):
+                _id = _id.split("/")[-1]
+
+            _request_params: Dict[str, Any] = {"contents": batch_texts, "model": _id, "config": {}}
+            if self.dimensions:
+                _request_params["config"]["output_dimensionality"] = self.dimensions
+            if self.task_type:
+                _request_params["config"]["task_type"] = self.task_type
+            if self.title:
+                _request_params["config"]["title"] = self.title
+            if not _request_params["config"]:
+                del _request_params["config"]
+
+            if self.request_params:
+                _request_params.update(self.request_params)
+
+            try:
+                response = await self.aclient.aio.models.embed_content(**_request_params)
+
+                # Extract embeddings from batch response
+                if response.embeddings:
+                    batch_embeddings = []
+                    for embedding in response.embeddings:
+                        if embedding.values is not None:
+                            batch_embeddings.append(embedding.values)
+                        else:
+                            batch_embeddings.append([])
+                    all_embeddings.extend(batch_embeddings)
+                else:
+                    # If no embeddings, add empty lists for each text in batch
+                    all_embeddings.extend([[] for _ in batch_texts])
+
+                # Extract usage information
+                usage_dict = None
+                if response.metadata and hasattr(response.metadata, "billable_character_count"):
+                    usage_dict = {"billable_character_count": response.metadata.billable_character_count}
+
+                # Add same usage info for each embedding in the batch
+                all_usage.extend([usage_dict] * len(batch_texts))
+
+            except Exception as e:
+                log_warning(f"Error in async batch embedding: {e}")
+                # Fallback to individual calls for this batch
+                for text in batch_texts:
+                    try:
+                        text_embedding: List[float]
+                        text_usage: Optional[Dict[str, Any]]
+                        text_embedding, text_usage = await self.async_get_embedding_and_usage(text)
+                        all_embeddings.append(text_embedding)
+                        all_usage.append(text_usage)
+                    except Exception as e2:
+                        log_warning(f"Error in individual async embedding fallback: {e2}")
+                        all_embeddings.append([])
+                        all_usage.append(None)
+
+        return all_embeddings, all_usage

agno/knowledge/embedder/huggingface.py

@@ -3,12 +3,12 @@ from os import getenv
 from typing import Any, Dict, List, Optional, Tuple
 
 from agno.knowledge.embedder.base import Embedder
-from agno.utils.log import logger
+from agno.utils.log import log_error, log_warning
 
 try:
     from huggingface_hub import AsyncInferenceClient, InferenceClient
 except ImportError:
-    logger.error("`huggingface-hub` not installed, please run `pip install huggingface-hub`")
+    log_error("`huggingface-hub` not installed, please run `pip install huggingface-hub`")
     raise
 
 
@@ -22,6 +22,11 @@ class HuggingfaceCustomEmbedder(Embedder):
     huggingface_client: Optional[InferenceClient] = None
     async_client: Optional[AsyncInferenceClient] = None
 
+    def __post_init__(self):
+        if self.enable_batch:
+            log_warning("HuggingfaceEmbedder does not support batch embeddings, setting enable_batch to False")
+            self.enable_batch = False
+
     @property
     def client(self) -> InferenceClient:
         if self.huggingface_client:
@@ -61,7 +66,7 @@ class HuggingfaceCustomEmbedder(Embedder):
             else:
                 return list(response)
         except Exception as e:
-            logger.warning(f"Failed to process embeddings: {e}")
+            log_warning(f"Failed to process embeddings: {e}")
             return []
 
     def get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict]]:
@@ -80,7 +85,7 @@ class HuggingfaceCustomEmbedder(Embedder):
             else:
                 return list(response)
         except Exception as e:
-            logger.warning(f"Failed to process embeddings: {e}")
+            log_warning(f"Failed to process embeddings: {e}")
             return []
 
     async def async_get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict]]:

agno/knowledge/embedder/jina.py

@@ -117,3 +117,66 @@ class JinaEmbedder(Embedder):
         except Exception as e:
             logger.warning(f"Failed to get embedding and usage: {e}")
             return [], None
+
+    async def _async_batch_response(self, texts: List[str]) -> Dict[str, Any]:
+        """Async batch version of _response using aiohttp."""
+        data = {
+            "model": self.id,
+            "late_chunking": self.late_chunking,
+            "dimensions": self.dimensions,
+            "embedding_type": self.embedding_type,
+            "input": texts,  # Jina API expects a list of texts for batch processing
+        }
+        if self.user is not None:
+            data["user"] = self.user
+        if self.request_params:
+            data.update(self.request_params)
+
+        timeout = aiohttp.ClientTimeout(total=self.timeout) if self.timeout else None
+
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.post(self.base_url, headers=self._get_headers(), json=data) as response:
+                response.raise_for_status()
+                return await response.json()
+
+    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]
+
+            try:
+                result = await self._async_batch_response(batch_texts)
+                batch_embeddings = [data["embedding"] for data in result["data"]]
+                all_embeddings.extend(batch_embeddings)
+
+                # For each embedding in the batch, add the same usage information
+                usage_dict = result.get("usage")
+                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/knowledge/embedder/mistral.py

@@ -3,13 +3,13 @@ from os import getenv
 from typing import Any, Dict, List, Optional, Tuple
 
 from agno.knowledge.embedder.base import Embedder
-from agno.utils.log import logger
+from agno.utils.log import log_error, log_info, log_warning
 
 try:
     from mistralai import Mistral  # type: ignore
     from mistralai.models.embeddingresponse import EmbeddingResponse  # type: ignore
 except ImportError:
-    logger.error("`mistralai` not installed")
+    log_error("`mistralai` not installed")
     raise
 
 
@@ -50,7 +50,7 @@ class MistralEmbedder(Embedder):
 
     def _response(self, text: str) -> EmbeddingResponse:
         _request_params: Dict[str, Any] = {
-            "inputs": text,
+            "inputs": [text],  # Mistral API expects a list
             "model": self.id,
         }
         if self.request_params:
@@ -67,7 +67,7 @@ class MistralEmbedder(Embedder):
                 return response.data[0].embedding
             return []
         except Exception as e:
-            logger.warning(f"Error getting embedding: {e}")
+            log_warning(f"Error getting embedding: {e}")
             return []
 
     def get_embedding_and_usage(self, text: str) -> Tuple[List[float], Dict[str, Any]]:
@@ -79,7 +79,7 @@ class MistralEmbedder(Embedder):
             usage: Dict[str, Any] = response.usage.model_dump() if response.usage else {}
             return embedding, usage
         except Exception as e:
-            logger.warning(f"Error getting embedding and usage: {e}")
+            log_warning(f"Error getting embedding and usage: {e}")
             return [], {}
 
     async def async_get_embedding(self, text: str) -> List[float]:
@@ -88,7 +88,7 @@ class MistralEmbedder(Embedder):
             # Check if the client has an async version of embeddings.create
             if hasattr(self.client.embeddings, "create_async"):
                 response: EmbeddingResponse = await self.client.embeddings.create_async(
-                    inputs=text, model=self.id, **self.request_params if self.request_params else {}
+                    inputs=[text], model=self.id, **self.request_params if self.request_params else {}
                 )
             else:
                 # Fallback to running sync method in thread executor
@@ -98,7 +98,7 @@ class MistralEmbedder(Embedder):
                 response: EmbeddingResponse = await loop.run_in_executor(  # type: ignore
                     None,
                     lambda: self.client.embeddings.create(
-                        inputs=text, model=self.id, **self.request_params if self.request_params else {}
+                        inputs=[text], model=self.id, **self.request_params if self.request_params else {}
                     ),
                 )
 
@@ -106,7 +106,7 @@ class MistralEmbedder(Embedder):
                 return response.data[0].embedding
             return []
         except Exception as e:
-            logger.warning(f"Error getting embedding: {e}")
+            log_warning(f"Error getting embedding: {e}")
             return []
 
     async def async_get_embedding_and_usage(self, text: str) -> Tuple[List[float], Dict[str, Any]]:
@@ -115,7 +115,7 @@ class MistralEmbedder(Embedder):
             # Check if the client has an async version of embeddings.create
             if hasattr(self.client.embeddings, "create_async"):
                 response: EmbeddingResponse = await self.client.embeddings.create_async(
-                    inputs=text, model=self.id, **self.request_params if self.request_params else {}
+                    inputs=[text], model=self.id, **self.request_params if self.request_params else {}
                 )
             else:
                 # Fallback to running sync method in thread executor
@@ -125,7 +125,7 @@ class MistralEmbedder(Embedder):
                 response: EmbeddingResponse = await loop.run_in_executor(  # type: ignore
                     None,
                     lambda: self.client.embeddings.create(
-                        inputs=text, model=self.id, **self.request_params if self.request_params else {}
+                        inputs=[text], model=self.id, **self.request_params if self.request_params else {}
                     ),
                 )
 
@@ -135,5 +135,72 @@ class MistralEmbedder(Embedder):
             usage: Dict[str, Any] = response.usage.model_dump() if response.usage else {}
             return embedding, usage
         except Exception as e:
-            logger.warning(f"Error getting embedding and usage: {e}")
+            log_warning(f"Error getting embedding and usage: {e}")
             return [], {}
+
+    async def async_get_embeddings_batch_and_usage(
+        self, texts: List[str]
+    ) -> Tuple[List[List[float]], List[Optional[Dict[str, Any]]]]:
+        """
+        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 = []
+        log_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]
+
+            _request_params: Dict[str, Any] = {
+                "inputs": batch_texts,  # Mistral API expects a list for batch processing
+                "model": self.id,
+            }
+            if self.request_params:
+                _request_params.update(self.request_params)
+
+            try:
+                # Check if the client has an async version of embeddings.create
+                if hasattr(self.client.embeddings, "create_async"):
+                    response: EmbeddingResponse = await self.client.embeddings.create_async(**_request_params)
+                else:
+                    # Fallback to running sync method in thread executor
+                    import asyncio
+
+                    loop = asyncio.get_running_loop()
+                    response: EmbeddingResponse = await loop.run_in_executor(  # type: ignore
+                        None, lambda: self.client.embeddings.create(**_request_params)
+                    )
+
+                # Extract embeddings from batch response
+                if response.data:
+                    batch_embeddings = [data.embedding for data in response.data if data.embedding]
+                    all_embeddings.extend(batch_embeddings)
+                else:
+                    # If no embeddings, add empty lists for each text in batch
+                    all_embeddings.extend([[] for _ in batch_texts])
+
+                # Extract usage information
+                usage_dict = response.usage.model_dump() if response.usage else None
+                # Add same usage info for each embedding in the batch
+                all_usage.extend([usage_dict] * len(batch_texts))
+
+            except Exception as e:
+                log_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:
+                        log_warning(f"Error in individual async embedding fallback: {e2}")
+                        all_embeddings.append([])
+                        all_usage.append(None)
+
+        return all_embeddings, all_usage

agno/knowledge/embedder/nebius.py

@@ -10,4 +10,4 @@ class NebiusEmbedder(OpenAIEmbedder):
     id: str = "BAAI/bge-en-icl"
     dimensions: int = 1024
     api_key: Optional[str] = getenv("NEBIUS_API_KEY")
-    base_url: str = "https://api.studio.nebius.com/v1/"
+    base_url: str = "https://api.tokenfactory.nebius.com/v1/"

agno/knowledge/embedder/ollama.py

@@ -45,6 +45,11 @@ class OllamaEmbedder(Embedder):
     ollama_client: Optional[OllamaClient] = None
     async_client: Optional[AsyncOllamaClient] = None
 
+    def __post_init__(self):
+        if self.enable_batch:
+            logger.warning("OllamaEmbedder does not support batch embeddings, setting enable_batch to False")
+            self.enable_batch = False
+
     @property
     def client(self) -> OllamaClient:
         if self.ollama_client:
@@ -80,6 +85,10 @@ class OllamaEmbedder(Embedder):
         if self.options is not None:
             kwargs["options"] = self.options
 
+        # Add dimensions parameter for models that support it
+        if self.dimensions is not None:
+            kwargs["dimensions"] = self.dimensions
+
         response = self.client.embed(input=text, model=self.id, **kwargs)
         if response and "embeddings" in response:
             embeddings = response["embeddings"]
@@ -112,6 +121,10 @@ class OllamaEmbedder(Embedder):
         if self.options is not None:
             kwargs["options"] = self.options
 
+        # Add dimensions parameter for models that support it
+        if self.dimensions is not None:
+            kwargs["dimensions"] = self.dimensions
+
         response = await self.aclient.embed(input=text, model=self.id, **kwargs)
         if response and "embeddings" in response:
             embeddings = response["embeddings"]

agno/knowledge/embedder/openai.py

@@ -78,21 +78,25 @@ class OpenAIEmbedder(Embedder):
         return self.client.embeddings.create(**_request_params)
 
     def get_embedding(self, text: str) -> List[float]:
-        response: CreateEmbeddingResponse = self.response(text=text)
         try:
+            response: CreateEmbeddingResponse = self.response(text=text)
             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)
+        try:
+            response: CreateEmbeddingResponse = self.response(text=text)
 
-        embedding = response.data[0].embedding
-        usage = response.usage
-        if usage:
-            return embedding, usage.model_dump()
-        return embedding, None
+            embedding = response.data[0].embedding
+            usage = response.usage
+            if usage:
+                return embedding, usage.model_dump()
+            return embedding, None
+        except Exception as e:
+            logger.warning(e)
+            return [], None
 
     async def async_get_embedding(self, text: str) -> List[float]:
         req: Dict[str, Any] = {
@@ -127,71 +131,33 @@ class OpenAIEmbedder(Embedder):
         if self.request_params:
             req.update(self.request_params)
 
-        response = await self.aclient.embeddings.create(**req)
-        embedding = response.data[0].embedding
-        usage = response.usage
-        return embedding, usage.model_dump() if usage else None
-
-    def get_embeddings_batch(self, texts: List[str], batch_size: int = 100) -> List[List[float]]:
-        """
-        Get embeddings for multiple texts in batches.
-
-        Args:
-            texts: List of text strings to embed
-            batch_size: Number of texts to process in each API call (max ~2048)
-
-        Returns:
-            List of embedding vectors
-        """
-        all_embeddings = []
-
-        for i in range(0, len(texts), batch_size):
-            batch_texts = texts[i : i + 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 = self.client.embeddings.create(**req)
-                batch_embeddings = [data.embedding for data in response.data]
-                all_embeddings.extend(batch_embeddings)
-            except Exception as e:
-                logger.warning(f"Error in batch embedding: {e}")
-                # Fallback to individual calls for this batch
-                for text in batch_texts:
-                    try:
-                        embedding = self.get_embedding(text)
-                        all_embeddings.append(embedding)
-                    except Exception as e2:
-                        logger.warning(f"Error in individual embedding fallback: {e2}")
-                        all_embeddings.append([])
-
-        return all_embeddings
+        try:
+            response = await self.aclient.embeddings.create(**req)
+            embedding = response.data[0].embedding
+            usage = response.usage
+            return embedding, usage.model_dump() if usage else None
+        except Exception as e:
+            logger.warning(e)
+            return [], None
 
-    async def async_get_embeddings_batch(self, texts: List[str], batch_size: int = 100) -> List[List[float]]:
+    async def async_get_embeddings_batch_and_usage(
+        self, texts: List[str]
+    ) -> Tuple[List[List[float]], List[Optional[Dict]]]:
         """
-        Get embeddings for multiple texts in batches (async version).
+        Get embeddings and usage for multiple texts in batches (async version).
 
         Args:
             texts: List of text strings to embed
-            batch_size: Number of texts to process in each API call (max ~2048)
 
         Returns:
-            List of embedding vectors
+            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} (async)")
 
-        for i in range(0, len(texts), batch_size):
-            batch_texts = texts[i : i + 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,
@@ -209,15 +175,21 @@ class OpenAIEmbedder(Embedder):
                 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 async calls for this batch
+                # Fallback to individual calls for this batch
                 for text in batch_texts:
                     try:
-                        embedding = await self.async_get_embedding(text)
+                        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
+        return all_embeddings, all_usage

agno/knowledge/embedder/sentence_transformer.py

@@ -25,11 +25,15 @@ class SentenceTransformerEmbedder(Embedder):
     prompt: Optional[str] = None
     normalize_embeddings: bool = False
 
+    def __post_init__(self):
+        # Initialize the SentenceTransformer model eagerly to avoid race conditions in async contexts
+        if self.sentence_transformer_client is None:
+            self.sentence_transformer_client = SentenceTransformer(model_name_or_path=self.id)
+
     def get_embedding(self, text: Union[str, List[str]]) -> List[float]:
-        if not self.sentence_transformer_client:
-            model = SentenceTransformer(model_name_or_path=self.id)
-        else:
-            model = self.sentence_transformer_client
+        if self.sentence_transformer_client is None:
+            raise RuntimeError("SentenceTransformer model not initialized")
+        model = self.sentence_transformer_client
         embedding = model.encode(text, prompt=self.prompt, normalize_embeddings=self.normalize_embeddings)
         try:
             if isinstance(embedding, np.ndarray):