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

chore: Updating PGVector Vector Store parameters format

Browse source
This commit is contained in:
joaoguilhermeS 2024-06-22 17:01:38 -03:00 • committed by Gabriel Luiz Freitas Almeida
commit be9539f420
1 changed files with 21 additions and 32 deletions
Download patch file Download diff file Expand all files Collapse all files

53
src/backend/base/langflow/components/vectorstores/pgvector.py
View file

@ -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,
)