diff --git a/src/backend/langflow/components/vectorstores/pgvector.py b/src/backend/langflow/components/vectorstores/pgvector.py index 786629dd0..73844b0da 100644 --- a/src/backend/langflow/components/vectorstores/pgvector.py +++ b/src/backend/langflow/components/vectorstores/pgvector.py @@ -1,13 +1,15 @@ +from typing import Optional, Union +from langflow import CustomComponent from typing import List, Optional from langchain.embeddings.base import Embeddings +from langchain.schema import BaseRetriever from langchain.schema import Document from langchain_community.vectorstores import VectorStore from langchain_community.vectorstores.pgvector import PGVector from langflow import CustomComponent - -class PostgresqlVectorComponent(CustomComponent): +class PGVectorComponent(CustomComponent): """ A custom component for implementing a Vector Store using PostgreSQL. """ @@ -15,7 +17,6 @@ class PostgresqlVectorComponent(CustomComponent): display_name: str = "PGVector" description: str = "Implementation of Vector Store using PostgreSQL" documentation = "https://python.langchain.com/docs/integrations/vectorstores/pgvector" - beta = True def build_config(self): """ @@ -25,8 +26,7 @@ class PostgresqlVectorComponent(CustomComponent): - dict: A dictionary containing the configuration options for the component. """ return { - "index_name": {"display_name": "Index Name", "value": "your_index"}, - "code": {"show": True, "display_name": "Code"}, + "code": {"show": False}, "documents": {"display_name": "Documents", "is_list": True}, "embedding": {"display_name": "Embedding"}, "pg_server_url": { @@ -41,8 +41,8 @@ class PostgresqlVectorComponent(CustomComponent): embedding: Embeddings, pg_server_url: str, collection_name: str, - documents: Optional[List[Document]] = None, - ) -> VectorStore: + documents: Optional[Document] = None, + ) -> Union[VectorStore, BaseRetriever]: """ Builds the Vector Store or BaseRetriever object. @@ -58,13 +58,13 @@ class PostgresqlVectorComponent(CustomComponent): try: if documents is None: - return PGVector.from_existing_index( + vector_store = PGVector.from_existing_index( embedding=embedding, collection_name=collection_name, connection_string=pg_server_url, ) - return PGVector.from_documents( + vector_store = PGVector.from_documents( embedding=embedding, documents=documents, collection_name=collection_name, @@ -72,3 +72,4 @@ class PostgresqlVectorComponent(CustomComponent): ) except Exception as e: raise RuntimeError(f"Failed to build PGVector: {e}") + return vector_store