jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Add number_of_results parameter to vector search components

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-04-02 02:23:15 -03:00
commit de3cce67ab
10 changed files with 70 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

8
src/backend/base/langflow/components/vectorsearch/ChromaSearch.py
View file

@ -48,6 +48,11 @@ class ChromaSearchComponent(LCVectorStoreComponent):
"display_name": "Server SSL Enabled",
"advanced": True,
},
"number_of_results": {
"display_name": "Number of Results",
"info": "Number of results to return.",
"advanced": True,
},
}
def build(
@ -57,6 +62,7 @@ class ChromaSearchComponent(LCVectorStoreComponent):
collection_name: str,
embedding: Embeddings,
chroma_server_ssl_enabled: bool,
number_of_results: int = 4,
index_directory: Optional[str] = None,
chroma_server_cors_allow_origins: Optional[str] = None,
chroma_server_host: Optional[str] = None,
@ -102,4 +108,4 @@ class ChromaSearchComponent(LCVectorStoreComponent):
client_settings=chroma_settings,
)
return self.search_with_vector_store(input_value, search_type, vector_store)
return self.search_with_vector_store(input_value, search_type, vector_store, k=number_of_results)

8
src/backend/base/langflow/components/vectorsearch/FAISSSearch.py
View file

@ -21,6 +21,11 @@ class FAISSSearchComponent(LCVectorStoreComponent):
},
"input_value": {"display_name": "Input"},
"index_name": {"display_name": "Index Name"},
"number_of_results": {
"display_name": "Number of Results",
"info": "Number of results to return.",
"advanced": True,
},
}
def build(
@ -28,6 +33,7 @@ class FAISSSearchComponent(LCVectorStoreComponent):
input_value: Text,
embedding: Embeddings,
folder_path: str,
number_of_results: int = 4,
index_name: str = "langflow_index",
) -> List[Record]:
if not folder_path:
@ -38,5 +44,5 @@ class FAISSSearchComponent(LCVectorStoreComponent):
raise ValueError("Failed to load the FAISS index.")
return self.search_with_vector_store(
vector_store=vector_store, input_value=input_value, search_type="similarity"
vector_store=vector_store, input_value=input_value, search_type="similarity", k=number_of_results
)

8
src/backend/base/langflow/components/vectorsearch/MongoDBAtlasVectorSearch.py
View file

@ -23,6 +23,11 @@ class MongoDBAtlasSearchComponent(LCVectorStoreComponent):
"index_name": {"display_name": "Index Name"},
"mongodb_atlas_cluster_uri": {"display_name": "MongoDB Atlas Cluster URI"},
"search_kwargs": {"display_name": "Search Kwargs", "advanced": True},
"number_of_results": {
"display_name": "Number of Results",
"info": "Number of results to return.",
"advanced": True,
},
}
def build( # type: ignore[override]
@ -30,6 +35,7 @@ class MongoDBAtlasSearchComponent(LCVectorStoreComponent):
input_value: Text,
search_type: str,
embedding: Embeddings,
number_of_results: int = 4,
collection_name: str = "",
db_name: str = "",
index_name: str = "",
@ -47,5 +53,5 @@ class MongoDBAtlasSearchComponent(LCVectorStoreComponent):
if not vector_store:
raise ValueError("Failed to create MongoDB Atlas Vector Store")
return self.search_with_vector_store(
vector_store=vector_store, input_value=input_value, search_type=search_type
vector_store=vector_store, input_value=input_value, search_type=search_type, k=number_of_results
)

8
src/backend/base/langflow/components/vectorsearch/PineconeSearch.py
View file

@ -38,6 +38,11 @@ class PineconeSearchComponent(PineconeComponent, LCVectorStoreComponent):
"default": 1,
"advanced": True,
},
"number_of_results": {
"display_name": "Number of Results",
"info": "Number of results to return.",
"advanced": True,
},
}
def build( # type: ignore[override]
@ -46,6 +51,7 @@ class PineconeSearchComponent(PineconeComponent, LCVectorStoreComponent):
embedding: Embeddings,
pinecone_env: str,
text_key: str = "text",
number_of_results: int = 4,
pool_threads: int = 4,
index_name: Optional[str] = None,
pinecone_api_key: Optional[str] = None,
@ -66,5 +72,5 @@ class PineconeSearchComponent(PineconeComponent, LCVectorStoreComponent):
raise ValueError("Failed to load the Pinecone index.")
return self.search_with_vector_store(
vector_store=vector_store, input_value=input_value, search_type=search_type
vector_store=vector_store, input_value=input_value, search_type=search_type, k=number_of_results
)

8
src/backend/base/langflow/components/vectorsearch/QdrantSearch.py
View file

@ -41,6 +41,11 @@ class QdrantSearchComponent(QdrantComponent, LCVectorStoreComponent):
"search_kwargs": {"display_name": "Search Kwargs", "advanced": True},
"timeout": {"display_name": "Timeout", "advanced": True},
"url": {"display_name": "URL", "advanced": True},
"number_of_results": {
"display_name": "Number of Results",
"info": "Number of results to return.",
"advanced": True,
},
}
def build( # type: ignore[override]
@ -48,6 +53,7 @@ class QdrantSearchComponent(QdrantComponent, LCVectorStoreComponent):
input_value: Text,
embedding: Embeddings,
collection_name: str,
number_of_results: int = 4,
search_type: str = "similarity",
api_key: Optional[str] = None,
content_payload_key: str = "page_content",
@ -88,5 +94,5 @@ class QdrantSearchComponent(QdrantComponent, LCVectorStoreComponent):
raise ValueError("Failed to load the Qdrant index.")
return self.search_with_vector_store(
vector_store=vector_store, input_value=input_value, search_type=search_type
vector_store=vector_store, input_value=input_value, search_type=search_type, k=number_of_results
)

8
src/backend/base/langflow/components/vectorsearch/RedisSearch.py
View file

@ -39,6 +39,11 @@ class RedisSearchComponent(RedisComponent, LCVectorStoreComponent):
"advanced": False,
},
"redis_index_name": {"display_name": "Redis Index", "advanced": False},
"number_of_results": {
"display_name": "Number of Results",
"info": "Number of results to return.",
"advanced": True,
},
}
def build( # type: ignore[override]
@ -48,6 +53,7 @@ class RedisSearchComponent(RedisComponent, LCVectorStoreComponent):
embedding: Embeddings,
redis_server_url: str,
redis_index_name: str,
number_of_results: int = 4,
schema: Optional[str] = None,
) -> List[Record]:
"""
@ -72,5 +78,5 @@ class RedisSearchComponent(RedisComponent, LCVectorStoreComponent):
raise ValueError("Failed to load the Redis index.")
return self.search_with_vector_store(
input_value=input_value, search_type=search_type, vector_store=vector_store
input_value=input_value, search_type=search_type, vector_store=vector_store, k=number_of_results
)

8
src/backend/base/langflow/components/vectorsearch/SupabaseVectorStoreSearch.py
View file

@ -26,6 +26,11 @@ class SupabaseSearchComponent(LCVectorStoreComponent):
"supabase_service_key": {"display_name": "Supabase Service Key"},
"supabase_url": {"display_name": "Supabase URL"},
"table_name": {"display_name": "Table Name", "advanced": True},
"number_of_results": {
"display_name": "Number of Results",
"info": "Number of results to return.",
"advanced": True,
},
}
def build(
@ -33,6 +38,7 @@ class SupabaseSearchComponent(LCVectorStoreComponent):
input_value: Text,
search_type: str,
embedding: Embeddings,
number_of_results: int = 4,
query_name: str = "",
supabase_service_key: str = "",
supabase_url: str = "",
@ -45,4 +51,4 @@ class SupabaseSearchComponent(LCVectorStoreComponent):
table_name=table_name,
query_name=query_name,
)
return self.search_with_vector_store(input_value, search_type, vector_store)
return self.search_with_vector_store(input_value, search_type, vector_store, k=number_of_results)

8
src/backend/base/langflow/components/vectorsearch/VectaraSearch.py
View file

@ -34,6 +34,11 @@ class VectaraSearchComponent(VectaraComponent, LCVectorStoreComponent):
"display_name": "Files Url",
"info": "Make vectara object using url of files (optional)",
},
"number_of_results": {
"display_name": "Number of Results",
"info": "Number of results to return.",
"advanced": True,
},
}
def build( # type: ignore[override]
@ -43,6 +48,7 @@ class VectaraSearchComponent(VectaraComponent, LCVectorStoreComponent):
vectara_customer_id: str,
vectara_corpus_id: str,
vectara_api_key: str,
number_of_results: int = 4,
) -> List[Record]:
source = "Langflow"
vector_store = Vectara(
@ -56,5 +62,5 @@ class VectaraSearchComponent(VectaraComponent, LCVectorStoreComponent):
raise ValueError("Failed to create Vectara Vector Store")
return self.search_with_vector_store(
vector_store=vector_store, input_value=input_value, search_type=search_type
vector_store=vector_store, input_value=input_value, search_type=search_type, k=number_of_results
)

9
src/backend/base/langflow/components/vectorsearch/WeaviateSearch.py
View file

@ -49,7 +49,11 @@ class WeaviateSearchVectorStore(WeaviateVectorStoreComponent, LCVectorStoreCompo
"field_type": "bool",
"advanced": True,
},
"code": {"show": False},
"number_of_results": {
"display_name": "Number of Results",
"info": "Number of results to return.",
"advanced": True,
},
}
def build( # type: ignore[override]
@ -57,6 +61,7 @@ class WeaviateSearchVectorStore(WeaviateVectorStoreComponent, LCVectorStoreCompo
input_value: Text,
search_type: str,
url: str,
number_of_results: int = 4,
search_by_text: bool = False,
api_key: Optional[str] = None,
index_name: Optional[str] = None,
@ -77,5 +82,5 @@ class WeaviateSearchVectorStore(WeaviateVectorStoreComponent, LCVectorStoreCompo
raise ValueError("Failed to load the Weaviate index.")
return self.search_with_vector_store(
vector_store=vector_store, input_value=input_value, search_type=search_type
vector_store=vector_store, input_value=input_value, search_type=search_type, k=number_of_results
)

8
src/backend/base/langflow/components/vectorsearch/pgvectorSearch.py
View file

@ -33,6 +33,11 @@ class PGVectorSearchComponent(PGVectorComponent, LCVectorStoreComponent):
},
"collection_name": {"display_name": "Table", "advanced": False},
"input_value": {"display_name": "Input"},
"number_of_results": {
"display_name": "Number of Results",
"info": "Number of results to return.",
"advanced": True,
},
}
def build( # type: ignore[override]
@ -42,6 +47,7 @@ class PGVectorSearchComponent(PGVectorComponent, LCVectorStoreComponent):
search_type: str,
pg_server_url: str,
collection_name: str,
number_of_results: int = 4,
) -> List[Record]:
"""
Builds the Vector Store or BaseRetriever object.
@ -64,5 +70,5 @@ class PGVectorSearchComponent(PGVectorComponent, LCVectorStoreComponent):
except Exception as e:
raise RuntimeError(f"Failed to build PGVector: {e}")
return self.search_with_vector_store(
input_value=input_value, search_type=search_type, vector_store=vector_store
input_value=input_value, search_type=search_type, vector_store=vector_store, k=number_of_results
)