feat: add contextualized embeddings support for voyage ai

cookbook/agent_concepts/knowledge/embedders/voyageai_embedder.py

@@ -2,13 +2,29 @@
 from agno.embedder.voyageai import VoyageAIEmbedder
 from agno.vectordb.pgvector import PgVector
 
-embeddings = VoyageAIEmbedder().get_embedding(
+embedder = VoyageAIEmbedder(
+    id="voyage-2",
+)
+
+embeddings = embedder.get_standard_embedding(
     "The quick brown fox jumps over the lazy dog."
 )
 
-# Print the embeddings and their dimensions
-print(f"Embeddings: {embeddings[:5]}")
-print(f"Dimensions: {len(embeddings)}")
+print(f"Standard Embeddings: {embeddings['embeddings'][:5]}")
+print(f"Usage: {embeddings['usage']}")
+
+# Contextualized embeddings
+# document_chunks = [
+#     ["This is the SEC filing on Leafy Inc.'s Q2 2024 performance."],
+#     ["The company's revenue increased by 15% compared to the previous quarter."],
+#     ["Operating expenses decreased by 3% due to cost optimization initiatives."],
+#     ["The company expects continued growth in Q3 2024."]
+# ]
+
+# contextualized_result = embedder.get_contextualized_embeddings(document_chunks)
+# print(f"Contextualized Embeddings count: {len(contextualized_result['embeddings'])}")
+# print(f"First embedding: {contextualized_result['embeddings'][0][:5]}")
+# print(f"Usage: {contextualized_result['usage']}")
 
 # Example usage:
 knowledge_base = AgentKnowledge(

libs/agno/agno/embedder/voyageai.py

@@ -1,23 +1,22 @@
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
 from agno.embedder.base import Embedder
 from agno.utils.log import logger
 
 try:
     from voyageai import Client as VoyageClient
-    from voyageai.object import EmbeddingsObject
+    from voyageai.object import EmbeddingsObject, ContextualizedEmbeddingsObject
 except ImportError:
     raise ImportError("`voyageai` not installed. Please install using `pip install voyageai`")
 
 
 @dataclass
 class VoyageAIEmbedder(Embedder):
-    id: str = "voyage-2"
+    id: str = "voyage-context-3"
     dimensions: int = 1024
     request_params: Optional[Dict[str, Any]] = None
     api_key: Optional[str] = None
-    base_url: str = "https://api.voyageai.com/v1/embeddings"
     max_retries: Optional[int] = None
     timeout: Optional[float] = None
     client_params: Optional[Dict[str, Any]] = None
@@ -39,26 +38,54 @@ def client(self) -> VoyageClient:
         self.voyage_client = VoyageClient(**_client_params)
         return self.voyage_client
 
-    def _response(self, text: str) -> EmbeddingsObject:
+    def _contextualized_response(self, documents_chunks: List[List[str]], input_type: str = "document") -> ContextualizedEmbeddingsObject:
+        """Get contextualized embedding response for multiple documents with chunks"""
         _request_params: Dict[str, Any] = {
-            "texts": [text],
+            "inputs": documents_chunks,  # List[List[str]] format
+            "model": self.id,
+            "input_type": input_type,
+        }
+        if self.request_params:
+            _request_params.update(self.request_params)
+        return self.client.contextualized_embed(**_request_params)
+
+    def _standard_response(self, text: str) -> EmbeddingsObject:
+        """Get standard embedding response for a single text"""
+        _request_params: Dict[str, Any] = {
+            "texts": [text], 
             "model": self.id,
         }
         if self.request_params:
             _request_params.update(self.request_params)
         return self.client.embed(**_request_params)
 
-    def get_embedding(self, text: str) -> List[float]:
-        response: EmbeddingsObject = self._response(text=text)
+    def get_standard_embedding(self, text: str) -> Dict[str, Any]:
+        """Get standard embedding for a single text with usage info"""
         try:
-            return response.embeddings[0]
+            response = self._standard_response(text=text)
+            return {
+                "embeddings": response.embeddings[0],
+                "total_tokens": response.total_tokens
+            }
         except Exception as e:
-            logger.warning(e)
-            return []
-
-    def get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict]]:
-        response: EmbeddingsObject = self._response(text=text)
+            logger.warning(f"Error getting standard embedding: {e}")
+            return {"embeddings": [], "total_tokens": None}
 
-        embedding = response.embeddings[0]
-        usage = {"total_tokens": response.total_tokens}
-        return embedding, usage
+    def get_contextualized_embeddings(
+        self, 
+        documents_chunks: List[List[str]], 
+        input_type: str = "document"
+    ) -> Dict[str, Any]:
+        """
+         Get contextualized embeddings for multiple documents, each with their own chunks.
+        """
+        try:
+            response = self._contextualized_response(documents_chunks, input_type)
+            embeddings = [emb for r in response.results for emb in r.embeddings]
+            return {
+                "embeddings": embeddings,
+                "total_tokens": response.total_tokens
+            }
+        except Exception as e:
+            logger.warning(f"Error getting contextualized embeddings: {e}")
+            return {"embeddings": [], "total_tokens": None}