Add support for Couchbase vector store (#1901)

* add couchbase vector store support

* add docs + minor changes

* Fix lint issues

* remove stray lines

* Add required validation and minor changes

* Address Comments

---------

Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
This commit is contained in:
Prajwal Pai 2024-05-23 18:36:21 +05:30 • committed by GitHub
commit 19680bb137
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 281 additions and 4 deletions

View file

@ -0,0 +1,73 @@
from typing import List, Optional
from langflow.components.vectorstores.base.model import LCVectorStoreComponent
from langflow.components.vectorstores.Couchbase import CouchbaseComponent
from langflow.field_typing import Embeddings, NestedDict, Text
from langflow.schema import Record
class CouchbaseSearchComponent(LCVectorStoreComponent):
display_name = "Couchbase Search"
description = "Search a Couchbase Vector Store for similar documents."
documentation = "https://python.langchain.com/docs/integrations/vectorstores/couchbase"
icon = "Couchbase"
field_order = [
"couchbase_connection_string",
"couchbase_username",
"couchbase_password",
"bucket_name",
"scope_name",
"collection_name",
"index_name",
]
def build_config(self):
return {
"input_value": {"display_name": "Input"},
"embedding": {"display_name": "Embedding"},
"couchbase_connection_string": {"display_name": "Couchbase Cluster connection string","required": True},
"couchbase_username": {"display_name": "Couchbase username","required": True},
"couchbase_password": {
"display_name": "Couchbase password",
"password": True,
"required": True
},
"bucket_name": {"display_name": "Bucket Name","required": True},
"scope_name": {"display_name": "Scope Name","required": True},
"collection_name": {"display_name": "Collection Name","required": True},
"index_name": {"display_name": "Index Name","required": True},
"number_of_results": {
"display_name": "Number of Results",
"info": "Number of results to return.",
"advanced": True,
},
}
def build( # type: ignore[override]
self,
input_value: Text,
embedding: Embeddings,
number_of_results: int = 4,
bucket_name: str = "",
scope_name: str = "",
collection_name: str = "",
index_name: str = "",
couchbase_connection_string: str = "",
couchbase_username: str = "",
couchbase_password: str = "",
) -> List[Record]:
vector_store = CouchbaseComponent().build(
couchbase_connection_string=couchbase_connection_string,
couchbase_username=couchbase_username,
couchbase_password=couchbase_password,
bucket_name=bucket_name,
scope_name=scope_name,
collection_name=collection_name,
embedding=embedding,
index_name=index_name,
)
if not vector_store:
raise ValueError("Failed to create Couchbase Vector Store")
return self.search_with_vector_store(
vector_store=vector_store, input_value=input_value, search_type="similarity", k=number_of_results
)

View file

@ -9,10 +9,12 @@ from .SupabaseVectorStoreSearch import SupabaseSearchComponent
from .VectaraSearch import VectaraSearchComponent
from .WeaviateSearch import WeaviateSearchVectorStore
from .pgvectorSearch import PGVectorSearchComponent
from .Couchbase import CouchbaseSearchComponent # type: ignore
__all__ = [
"AstraDBSearchComponent",
"ChromaSearchComponent",
"CouchbaseSearchComponent",
"FAISSSearchComponent",
"MongoDBAtlasSearchComponent",
"PineconeSearchComponent",

View file

@ -0,0 +1,95 @@
from typing import List, Optional, Union
from langchain.schema import BaseRetriever
from langchain_community.vectorstores import CouchbaseVectorStore
from langflow.custom import CustomComponent
from langflow.field_typing import Embeddings, VectorStore
from langflow.schema import Record
from datetime import timedelta
from couchbase.auth import PasswordAuthenticator # type: ignore
from couchbase.cluster import Cluster # type: ignore
from couchbase.options import ClusterOptions # type: ignore
class CouchbaseComponent(CustomComponent):
display_name = "Couchbase"
description = "Construct a `Couchbase Vector Search` vector store from raw documents."
documentation = "https://python.langchain.com/docs/integrations/vectorstores/couchbase"
icon = "Couchbase"
field_order = [
"couchbase_connection_string",
"couchbase_username",
"couchbase_password",
"bucket_name",
"scope_name",
"collection_name",
"index_name",
]
def build_config(self):
return {
"inputs": {"display_name": "Input", "input_types": ["Document", "Record"]},
"embedding": {"display_name": "Embedding"},
"couchbase_connection_string": {"display_name": "Couchbase Cluster connection string","required": True},
"couchbase_username": {"display_name": "Couchbase username","required": True},
"couchbase_password": {
"display_name": "Couchbase password",
"password": True,
"required": True
},
"bucket_name": {"display_name": "Bucket Name","required": True},
"scope_name": {"display_name": "Scope Name","required": True},
"collection_name": {"display_name": "Collection Name","required": True},
"index_name": {"display_name": "Index Name","required": True},
}
def build(
self,
embedding: Embeddings,
inputs: Optional[List[Record]] = None,
bucket_name: str = "",
scope_name: str = "",
collection_name: str = "",
index_name: str = "",
couchbase_connection_string: str = "",
couchbase_username: str = "",
couchbase_password: str = "",
) -> Union[VectorStore, BaseRetriever]:
try:
auth = PasswordAuthenticator(couchbase_username, couchbase_password)
options = ClusterOptions(auth)
cluster = Cluster(couchbase_connection_string, options)
cluster.wait_until_ready(timedelta(seconds=5))
except Exception as e:
raise ValueError(f"Failed to connect to Couchbase: {e}")
documents = []
for _input in inputs or []:
if isinstance(_input, Record):
documents.append(_input.to_lc_document())
else:
documents.append(_input)
if documents:
vector_store = CouchbaseVectorStore.from_documents(
documents=documents,
cluster=cluster,
bucket_name=bucket_name,
scope_name=scope_name,
collection_name=collection_name,
embedding=embedding,
index_name=index_name,
)
else:
vector_store = CouchbaseVectorStore(
cluster=cluster,
bucket_name=bucket_name,
scope_name=scope_name,
collection_name=collection_name,
embedding=embedding,
index_name=index_name,
)
return vector_store

View file

@ -9,10 +9,12 @@ from .SupabaseVectorStore import SupabaseComponent
from .Vectara import VectaraComponent
from .Weaviate import WeaviateVectorStoreComponent
from .pgvector import PGVectorComponent
from .Couchbase import CouchbaseComponent
__all__ = [
"AstraDBVectorStoreComponent",
"ChromaComponent",
"CouchbaseComponent",
"FAISSComponent",
"MongoDBAtlasComponent",
"PineconeComponent",