diff --git a/src/backend/base/langflow/components/vectorstores/pgvector.py b/src/backend/base/langflow/components/vectorstores/pgvector.py index 95808c70c..a078d67cf 100644 --- a/src/backend/base/langflow/components/vectorstores/pgvector.py +++ b/src/backend/base/langflow/components/vectorstores/pgvector.py @@ -11,24 +11,19 @@ from langflow.schema import Data class PGVectorStoreComponent(LCVectorStoreComponent): display_name = "PGVector" description = "PGVector Vector Store with search capabilities" - documentation = "https://python.langchain.com/docs/modules/data_connection/vectorstores/integrations/pgvector" + documentation = "https://python.langchain.com/v0.2/docs/integrations/vectorstores/pgvector/" icon = "PGVector" inputs = [ SecretStrInput(name="pg_server_url", display_name="PostgreSQL Server Connection String", required=True), StrInput(name="collection_name", display_name="Table", required=True), - HandleInput(name="embedding", display_name="Embedding", input_types=["Embeddings"]), + MultilineInput(name="search_query", display_name="Search Query"), DataInput( - name="vector_store_inputs", - display_name="Vector Store Inputs", + name="ingest_data", + display_name="Ingestion Data", is_list=True, ), - BoolInput( - name="add_to_vector_store", - display_name="Add to Vector Store", - info="If true, the Vector Store Inputs will be added to the Vector Store.", - ), - MultilineInput(name="search_input", display_name="Search Input"), + HandleInput(name="embedding", display_name="Embedding", input_types=["Embeddings"]), IntInput( name="number_of_results", display_name="Number of Results", @@ -36,33 +31,27 @@ class PGVectorStoreComponent(LCVectorStoreComponent): value=4, advanced=True, ), + HandleInput(name="embedding", display_name="Embedding", input_types=["Embeddings"]), ] def build_vector_store(self) -> PGVector: return self._build_pgvector() def _build_pgvector(self) -> PGVector: - if self.add_to_vector_store: - documents = [] - for _input in self.vector_store_inputs or []: - if isinstance(_input, Data): - documents.append(_input.to_lc_document()) - else: - documents.append(_input) - - if documents: - pgvector = PGVector.from_documents( - embedding=self.embedding, - documents=documents, - collection_name=self.collection_name, - connection_string=self.pg_server_url, - ) + documents = [] + for _input in self.ingest_data or []: + if isinstance(_input, Data): + documents.append(_input.to_lc_document()) else: - pgvector = PGVector.from_existing_index( - embedding=self.embedding, - collection_name=self.collection_name, - connection_string=self.pg_server_url, - ) + documents.append(_input) + + if documents: + pgvector = PGVector.from_documents( + embedding=self.embedding, + documents=documents, + collection_name=self.collection_name, + connection_string=self.pg_server_url, + ) else: pgvector = PGVector.from_existing_index( embedding=self.embedding, @@ -75,9 +64,9 @@ class PGVectorStoreComponent(LCVectorStoreComponent): def search_documents(self) -> List[Data]: vector_store = self._build_pgvector() - if self.search_input and isinstance(self.search_input, str) and self.search_input.strip(): + if self.search_query and isinstance(self.search_query, str) and self.search_query.strip(): docs = vector_store.similarity_search( - query=self.search_input, + query=self.search_query, k=self.number_of_results, )