Upload 105 files

src/rag_pipelines/vectordb/milvus_retriever.py

@@ -1,8 +1,7 @@
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 import weave
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_voyageai import VoyageAIEmbeddings
 from langchain_milvus.retrievers import MilvusCollectionHybridSearchRetriever
 from pymilvus import (
     Collection,
@@ -25,7 +24,7 @@ class MilvusRetriever(weave.Model):
     """
 
     collection: Optional[Collection] = None
-    dense_embedding_model: Optional[Union[VoyageAIEmbeddings, HuggingFaceEmbeddings]] = None
+    dense_embedding_model: Optional[HuggingFaceEmbeddings] = None
     sparse_embedding_model: Optional[SparseEmbeddings] = None
     anns_fields: Optional[list[str]] = None
     field_search_params: Optional[list[dict[str, Any]]] = None
@@ -37,7 +36,7 @@ class MilvusRetriever(weave.Model):
     def __init__(
         self,
         collection: Collection,
-        dense_embedding_model: Union[VoyageAIEmbeddings, HuggingFaceEmbeddings],
+        dense_embedding_model: HuggingFaceEmbeddings,
         sparse_embedding_model: SparseEmbeddings,
         anns_fields: Optional[list[str]] = None,
         field_search_params: Optional[list[dict[str, Any]]] = None,
@@ -52,7 +51,7 @@ class MilvusRetriever(weave.Model):
 
         Args:
             collection (Collection): The Milvus collection instance for document retrieval.
-            dense_embedding_model (HuggingFaceEmbeddings/ VoyageAIEmbeddings): Model for generating dense embeddings.
+            dense_embedding_model (HuggingFaceEmbeddings): Model for generating dense embeddings.
             sparse_embedding_model (SparseEmbeddings): Model for generating sparse embeddings.
             anns_fields (Optional[list[str]]): List of fields to search in the collection.
             field_search_params (Optional[dict[str, Any]]): Parameters for field-specific search.