rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-01-04 02:42:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: broken voyageai embedding API (#2013)

Browse Source 
This PR fixes the broken Embedding API for Voyageai.
This commit is contained in:
Prashant Dixit
2025-01-13 22:22:38 +05:30
committed by GitHub
parent d8d11f48e7
commit b66cd943a7
1 changed files with 16 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
python/python/lancedb/embeddings/voyageai.py
View File

@@ -59,7 +59,7 @@ class VoyageAIEmbeddingFunction(EmbeddingFunction):
.create(name="voyage-3")
class TextModel(LanceModel):
data: str = voyageai.SourceField()
text: str = voyageai.SourceField()
vector: Vector(voyageai.ndims()) = voyageai.VectorField()
data = [ { "text": "hello world" },
@@ -74,6 +74,14 @@ class VoyageAIEmbeddingFunction(EmbeddingFunction):
name: str
client: ClassVar = None
text_embedding_models: list = [
"voyage-3",
"voyage-3-lite",
"voyage-finance-2",
"voyage-law-2",
"voyage-code-2",
]
multimodal_embedding_models: list = ["voyage-multimodal-3"]
def ndims(self):
if self.name == "voyage-3-lite":
@@ -115,13 +123,14 @@ class VoyageAIEmbeddingFunction(EmbeddingFunction):
truncation: Optional[bool]
"""
if self.name in ["voyage-multimodal-3"]:
rs = VoyageAIEmbeddingFunction._get_client().multimodal_embed(
inputs=[[text]], model=self.name, **kwargs
)
client = VoyageAIEmbeddingFunction._get_client()
if self.name in self.text_embedding_models:
rs = client.embed(texts=[text], model=self.name, **kwargs)
elif self.name in self.multimodal_embedding_models:
rs = client.multimodal_embed(inputs=[[text]], model=self.name, **kwargs)
else:
rs = VoyageAIEmbeddingFunction._get_client().embed(
texts=[text], model=self.name, **kwargs
raise ValueError(
f"Model {self.name} not supported to generate text embeddings"
)
return rs.embeddings[0]