Harrison/pgvector (#1679)

Co-authored-by: Aman Kumar <krsingh.aman@gmail.com>

docs/ecosystem/pgvector.md

@@ -0,0 +1,29 @@
+# PGVector
+
+This page covers how to use the Postgres [PGVector](https://github.com/pgvector/pgvector) ecosystem within LangChain
+It is broken into two parts: installation and setup, and then references to specific PGVector wrappers.
+
+## Installation
+- Install the Python package with `pip install pgvector`
+
+
+## Setup
+1. The first step is to create a database with the `pgvector` extension installed.
+
+    Follow the steps at [PGVector Installation Steps](https://github.com/pgvector/pgvector#installation) to install the database and the extension. The docker image is the easiest way to get started.
+
+## Wrappers
+
+### VectorStore
+
+There exists a wrapper around Postgres vector databases, allowing you to use it as a vectorstore,
+whether for semantic search or example selection.
+
+To import this vectorstore:
+```python
+from langchain.vectorstores.pgvector import PGVector
+```
+
+### Usage
+
+For a more detailed walkthrough of the PGVector Wrapper, see [this notebook](../modules/indexes/vectorstore_examples/pgvector.ipynb)

docs/modules/indexes/how_to_guides.rst

@@ -52,6 +52,8 @@ In the below guides, we cover different types of vectorstores and how to use the
 
 `Weaviate <./vectorstore_examples/weaviate.html>`_: A walkthrough of how to use the Weaviate vectorstore wrapper.
 
+`PGVector <./vectorstore_examples/pgvector.html>`_: A walkthrough of how to use the PGVector (Postgres Vector DB) vectorstore wrapper.
+
 
 .. toctree::
    :maxdepth: 1

docs/modules/indexes/vectorstore_examples/pgvector.ipynb

@@ -0,0 +1,194 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# PGVector\n",
+    "\n",
+    "This notebook shows how to use functionality related to the Postgres vector database (PGVector)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## Loading Environment Variables\n",
+    "from typing import List, Tuple\n",
+    "from dotenv import load_dotenv\n",
+    "load_dotenv()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
+    "from langchain.text_splitter import CharacterTextSplitter\n",
+    "from langchain.vectorstores.pgvector import PGVector\n",
+    "from langchain.document_loaders import TextLoader\n",
+    "from langchain.docstore.document import Document"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "loader = TextLoader('../../state_of_the_union.txt')\n",
+    "documents = loader.load()\n",
+    "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
+    "docs = text_splitter.split_documents(documents)\n",
+    "\n",
+    "embeddings = OpenAIEmbeddings()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## PGVector needs the connection string to the database.\n",
+    "## We will load it from the environment variables.\n",
+    "import os\n",
+    "CONNECTION_STRING = PGVector.connection_string_from_db_params(\n",
+    "    driver=os.environ.get(\"PGVECTOR_DRIVER\", \"psycopg2\"),\n",
+    "    host=os.environ.get(\"PGVECTOR_HOST\", \"localhost\"),\n",
+    "    port=int(os.environ.get(\"PGVECTOR_PORT\", \"5432\")),\n",
+    "    database=os.environ.get(\"PGVECTOR_DATABASE\", \"postgres\"),\n",
+    "    user=os.environ.get(\"PGVECTOR_USER\", \"postgres\"),\n",
+    "    password=os.environ.get(\"PGVECTOR_PASSWORD\", \"postgres\"),\n",
+    ")\n",
+    "\n",
+    "\n",
+    "## Example\n",
+    "# postgresql+psycopg2://username:password@localhost:5432/database_name"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Similarity search with score"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Similarity Search with Euclidean Distance (Default)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# The PGVector Module will try to create a table with the name of the collection. So, make sure that the collection name is unique and the user has the \n",
+    "# permission to create a table.\n",
+    "\n",
+    "db = PGVector.from_documents(\n",
+    "    embedding=embeddings,\n",
+    "    documents=docs,\n",
+    "    collection_name=\"state_of_the_union\",\n",
+    "    connection_string=CONNECTION_STRING,\n",
+    ")\n",
+    "\n",
+    "query = \"What did the president say about Ketanji Brown Jackson\"\n",
+    "docs_with_score: List[Tuple[Document, float]] = db.similarity_search_with_score(query)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "--------------------------------------------------------------------------------\n",
+      "Score:  0.6076628081132506\n",
+      "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
+      "\n",
+      "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n",
+      "\n",
+      "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n",
+      "\n",
+      "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.\n",
+      "--------------------------------------------------------------------------------\n",
+      "--------------------------------------------------------------------------------\n",
+      "Score:  0.6076628081132506\n",
+      "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
+      "\n",
+      "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n",
+      "\n",
+      "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n",
+      "\n",
+      "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.\n",
+      "--------------------------------------------------------------------------------\n",
+      "--------------------------------------------------------------------------------\n",
+      "Score:  0.6076804780049968\n",
+      "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
+      "\n",
+      "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n",
+      "\n",
+      "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n",
+      "\n",
+      "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.\n",
+      "--------------------------------------------------------------------------------\n",
+      "--------------------------------------------------------------------------------\n",
+      "Score:  0.6076804780049968\n",
+      "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
+      "\n",
+      "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n",
+      "\n",
+      "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n",
+      "\n",
+      "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.\n",
+      "--------------------------------------------------------------------------------\n"
+     ]
+    }
+   ],
+   "source": [
+    "for doc, score in docs_with_score:\n",
+    "    print(\"-\" * 80)\n",
+    "    print(\"Score: \", score)\n",
+    "    print(doc.page_content)\n",
+    "    print(\"-\" * 80)\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.10"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

langchain/vectorstores/pgvector.py

@@ -0,0 +1,417 @@
+import enum
+import logging
+import uuid
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import sqlalchemy
+from pgvector.sqlalchemy import Vector
+from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy.orm import Mapped, Session, declarative_base, relationship
+
+from langchain.docstore.document import Document
+from langchain.embeddings.base import Embeddings
+from langchain.utils import get_from_dict_or_env
+from langchain.vectorstores.base import VectorStore
+
+Base = declarative_base()  # type: Any
+
+
+ADA_TOKEN_COUNT = 1536
+_LANGCHAIN_DEFAULT_COLLECTION_NAME = "langchain"
+
+
+class BaseModel(Base):
+    __abstract__ = True
+    uuid = sqlalchemy.Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+
+
+class CollectionStore(BaseModel):
+    __tablename__ = "langchain_pg_collection"
+
+    name = sqlalchemy.Column(sqlalchemy.String)
+    cmetadata = sqlalchemy.Column(sqlalchemy.JSON)
+
+    embeddings = relationship(
+        "EmbeddingStore",
+        back_populates="collection",
+        passive_deletes=True,
+    )
+
+    @classmethod
+    def get_by_name(cls, session: Session, name: str) -> Optional["CollectionStore"]:
+        return session.query(cls).filter(cls.name == name).first()
+
+    @classmethod
+    def get_or_create(
+        cls,
+        session: Session,
+        name: str,
+        cmetadata: Optional[dict] = None,
+    ) -> Tuple["CollectionStore", bool]:
+        """
+        Get or create a collection.
+        Returns [Collection, bool] where the bool is True if the collection was created.
+        """
+        created = False
+        collection = cls.get_by_name(session, name)
+        if collection:
+            return collection, created
+
+        collection = cls(name=name, metadata=cmetadata)
+        session.add(collection)
+        session.commit()
+        created = True
+        return collection, created
+
+
+class EmbeddingStore(BaseModel):
+    __tablename__ = "langchain_pg_embedding"
+
+    collection_id: Mapped[UUID] = sqlalchemy.Column(
+        UUID(as_uuid=True),
+        sqlalchemy.ForeignKey(
+            f"{CollectionStore.__tablename__}.uuid",
+            ondelete="CASCADE",
+        ),
+    )
+    collection = relationship(CollectionStore, back_populates="embeddings")
+
+    embedding: Vector = sqlalchemy.Column(Vector(ADA_TOKEN_COUNT))
+    document = sqlalchemy.Column(sqlalchemy.String, nullable=True)
+    cmetadata = sqlalchemy.Column(sqlalchemy.JSON, nullable=True)
+
+    # custom_id : any user defined id
+    custom_id = sqlalchemy.Column(sqlalchemy.String, nullable=True)
+
+
+class QueryResult:
+    EmbeddingStore: EmbeddingStore
+    distance: float
+
+
+class DistanceStrategy(str, enum.Enum):
+    EUCLIDEAN = EmbeddingStore.embedding.l2_distance
+    COSINE = EmbeddingStore.embedding.cosine_distance
+    MAX_INNER_PRODUCT = EmbeddingStore.embedding.max_inner_product
+
+
+DEFAULT_DISTANCE_STRATEGY = DistanceStrategy.EUCLIDEAN
+
+
+class PGVector(VectorStore):
+    """
+    VectorStore implementation using Postgres and pgvector.
+    - `connection_string` is a postgres connection string.
+    - `embedding_function` any embedding function implementing
+        `langchain.embeddings.base.Embeddings` interface.
+    - `collection_name` is the name of the collection to use. (default: langchain)
+        - NOTE: This is not the name of the table, but the name of the collection.
+            The tables will be created when initializing the store (if not exists)
+            So, make sure the user has the right permissions to create tables.
+    - `distance_strategy` is the distance strategy to use. (default: EUCLIDEAN)
+        - `EUCLIDEAN` is the euclidean distance.
+        - `COSINE` is the cosine distance.
+    - `pre_delete_collection` if True, will delete the collection if it exists.
+        (default: False)
+        - Useful for testing.
+    """
+
+    def __init__(
+        self,
+        connection_string: str,
+        embedding_function: Embeddings,
+        collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
+        distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY,
+        pre_delete_collection: bool = False,
+        logger: Optional[logging.Logger] = None,
+    ) -> None:
+        self.connection_string = connection_string
+        self.embedding_function = embedding_function
+        self.collection_name = collection_name
+        self.distance_strategy = distance_strategy
+        self.pre_delete_collection = pre_delete_collection
+        self.logger = logger or logging.getLogger(__name__)
+        self.__post_init__()
+
+    def __post_init__(
+        self,
+    ) -> None:
+        """
+        Initialize the store.
+        """
+        self._conn = self.connect()
+        # self.create_vector_extension()
+        self.create_tables_if_not_exists()
+        self.create_collection()
+
+    def connect(self) -> sqlalchemy.engine.Connection:
+        engine = sqlalchemy.create_engine(self.connection_string)
+        conn = engine.connect()
+        return conn
+
+    def create_vector_extension(self) -> None:
+        try:
+            with Session(self._conn) as session:
+                statement = sqlalchemy.text("CREATE EXTENSION IF NOT EXISTS vector")
+                session.execute(statement)
+                session.commit()
+        except Exception as e:
+            self.logger.exception(e)
+
+    def create_tables_if_not_exists(self) -> None:
+        Base.metadata.create_all(self._conn)
+
+    def drop_tables(self) -> None:
+        Base.metadata.drop_all(self._conn)
+
+    def create_collection(self) -> None:
+        if self.pre_delete_collection:
+            self.delete_collection()
+        with Session(self._conn) as session:
+            CollectionStore.get_or_create(session, self.collection_name)
+
+    def delete_collection(self) -> None:
+        self.logger.debug("Trying to delete collection")
+        with Session(self._conn) as session:
+            collection = self.get_collection(session)
+            if not collection:
+                self.logger.error("Collection not found")
+                return
+            session.delete(collection)
+            session.commit()
+
+    def get_collection(self, session: Session) -> Optional["CollectionStore"]:
+        return CollectionStore.get_by_name(session, self.collection_name)
+
+    def add_texts(
+        self,
+        texts: Iterable[str],
+        metadatas: Optional[List[dict]] = None,
+        ids: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> List[str]:
+        """Run more texts through the embeddings and add to the vectorstore.
+
+        Args:
+            texts: Iterable of strings to add to the vectorstore.
+            metadatas: Optional list of metadatas associated with the texts.
+            kwargs: vectorstore specific parameters
+
+        Returns:
+            List of ids from adding the texts into the vectorstore.
+        """
+        if ids is None:
+            ids = [str(uuid.uuid1()) for _ in texts]
+
+        embeddings = self.embedding_function.embed_documents(list(texts))
+
+        if not metadatas:
+            metadatas = [{} for _ in texts]
+
+        with Session(self._conn) as session:
+            collection = self.get_collection(session)
+            if not collection:
+                raise ValueError("Collection not found")
+            for text, metadata, embedding, id in zip(texts, metadatas, embeddings, ids):
+                embedding_store = EmbeddingStore(
+                    embedding=embedding,
+                    document=text,
+                    cmetadata=metadata,
+                    custom_id=id,
+                )
+                collection.embeddings.append(embedding_store)
+                session.add(embedding_store)
+            session.commit()
+
+        return ids
+
+    def similarity_search(
+        self,
+        query: str,
+        k: int = 4,
+        **kwargs: Any,
+    ) -> List[Document]:
+        """Run similarity search with PGVector with distance.
+
+        Args:
+            query (str): Query text to search for.
+            k (int): Number of results to return. Defaults to 4.
+            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.
+
+        Returns:
+            List of Documents most similar to the query.
+        """
+        embedding = self.embedding_function.embed_query(text=query)
+        return self.similarity_search_by_vector(
+            embedding=embedding,
+            k=k,
+        )
+
+    def similarity_search_with_score(
+        self, query: str, k: int = 4
+    ) -> List[Tuple[Document, float]]:
+        """Return docs most similar to query.
+
+        Args:
+            query: Text to look up documents similar to.
+            k: Number of Documents to return. Defaults to 4.
+
+        Returns:
+            List of Documents most similar to the query and score for each
+        """
+        embedding = self.embedding_function.embed_query(query)
+        docs = self.similarity_search_with_score_by_vector(embedding, k)
+        return docs
+
+    def similarity_search_with_score_by_vector(
+        self,
+        embedding: List[float],
+        k: int = 4,
+    ) -> List[Tuple[Document, float]]:
+        with Session(self._conn) as session:
+            collection = self.get_collection(session)
+            if not collection:
+                raise ValueError("Collection not found")
+
+            results: List[QueryResult] = (
+                session.query(
+                    EmbeddingStore,
+                    self.distance_strategy(embedding).label("distance"),  # type: ignore
+                )
+                .filter(EmbeddingStore.collection_id == collection.uuid)
+                .order_by(sqlalchemy.asc("distance"))
+                .join(
+                    CollectionStore,
+                    EmbeddingStore.collection_id == CollectionStore.uuid,
+                )
+                .limit(k)
+                .all()
+            )
+            docs = [
+                (
+                    Document(
+                        page_content=result.EmbeddingStore.document,
+                        metadata=result.EmbeddingStore.cmetadata,
+                    ),
+                    result.distance if self.embedding_function is not None else None,
+                )
+                for result in results
+            ]
+            return docs
+
+    def similarity_search_by_vector(
+        self,
+        embedding: List[float],
+        k: int = 4,
+        **kwargs: Any,
+    ) -> List[Document]:
+        """Return docs most similar to embedding vector.
+
+        Args:
+            embedding: Embedding to look up documents similar to.
+            k: Number of Documents to return. Defaults to 4.
+
+        Returns:
+            List of Documents most similar to the query vector.
+        """
+        docs_and_scores = self.similarity_search_with_score_by_vector(
+            embedding=embedding, k=k
+        )
+        return [doc for doc, _ in docs_and_scores]
+
+    @classmethod
+    def from_texts(
+        cls,
+        texts: List[str],
+        embedding: Embeddings,
+        metadatas: Optional[List[dict]] = None,
+        collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
+        distance_strategy: DistanceStrategy = DistanceStrategy.COSINE,
+        ids: Optional[List[str]] = None,
+        pre_delete_collection: bool = False,
+        **kwargs: Any,
+    ) -> "PGVector":
+        """
+        Return VectorStore initialized from texts and embeddings.
+        Postgres connection string is required
+        "Either pass it as a parameter
+        or set the PGVECTOR_CONNECTION_STRING environment variable.
+        """
+
+        connection_string = cls.get_connection_string(kwargs)
+
+        store = cls(
+            connection_string=connection_string,
+            collection_name=collection_name,
+            embedding_function=embedding,
+            distance_strategy=distance_strategy,
+            pre_delete_collection=pre_delete_collection,
+        )
+
+        store.add_texts(texts=texts, metadatas=metadatas, ids=ids, **kwargs)
+        return store
+
+    @classmethod
+    def get_connection_string(cls, kwargs: Dict[str, Any]) -> str:
+        connection_string: str = get_from_dict_or_env(
+            data=kwargs,
+            key="connection_string",
+            env_key="PGVECTOR_CONNECTION_STRING",
+        )
+
+        if not connection_string:
+            raise ValueError(
+                "Postgres connection string is required"
+                "Either pass it as a parameter"
+                "or set the PGVECTOR_CONNECTION_STRING environment variable."
+            )
+
+        return connection_string
+
+    @classmethod
+    def from_documents(
+        cls,
+        documents: List[Document],
+        embedding: Embeddings,
+        collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
+        distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY,
+        ids: Optional[List[str]] = None,
+        pre_delete_collection: bool = False,
+        **kwargs: Any,
+    ) -> "PGVector":
+        """
+        Return VectorStore initialized from documents and embeddings.
+        Postgres connection string is required
+        "Either pass it as a parameter
+        or set the PGVECTOR_CONNECTION_STRING environment variable.
+        """
+
+        texts = [d.page_content for d in documents]
+        metadatas = [d.metadata for d in documents]
+        connection_string = cls.get_connection_string(kwargs)
+
+        kwargs["connection_string"] = connection_string
+
+        return cls.from_texts(
+            texts=texts,
+            pre_delete_collection=pre_delete_collection,
+            embedding=embedding,
+            distance_strategy=distance_strategy,
+            metadatas=metadatas,
+            ids=ids,
+            collection_name=collection_name,
+            **kwargs,
+        )
+
+    @classmethod
+    def connection_string_from_db_params(
+        cls,
+        driver: str,
+        host: str,
+        port: int,
+        database: str,
+        user: str,
+        password: str,
+    ) -> str:
+        """Return connection string from database parameters."""
+        return f"postgresql+{driver}://{user}:{password}@{host}:{port}/{database}"

pyproject.toml

@@ -51,6 +51,9 @@ pypdf = {version = "^3.4.0", optional = true}
 networkx = {version="^2.6.3", optional = true}
 aleph-alpha-client = {version="^2.15.0", optional = true}
 deeplake = {version = "^3.2.9", optional = true}
+pgvector = {version = "^0.1.6", optional = true}
+psycopg2-binary = {version = "^2.9.5", optional = true}
+
 
 [tool.poetry.group.docs.dependencies]
 autodoc_pydantic = "^1.8.0"
@@ -97,7 +100,7 @@ playwright = "^1.28.0"
 
 [tool.poetry.extras]
 llms = ["anthropic", "cohere", "openai", "nlpcloud", "huggingface_hub", "manifest-ml", "torch", "transformers"]
-all = ["anthropic", "cohere", "openai", "nlpcloud", "huggingface_hub", "manifest-ml", "elasticsearch", "opensearch-py", "google-search-results", "faiss-cpu", "sentence_transformers", "transformers", "spacy", "nltk", "wikipedia", "beautifulsoup4", "tiktoken", "torch", "jinja2", "pinecone-client", "weaviate-client", "redis", "google-api-python-client", "wolframalpha", "qdrant-client", "tensorflow-text", "pypdf", "networkx", "nomic", "aleph-alpha-client", "deeplake", ]
+all = ["anthropic", "cohere", "openai", "nlpcloud", "huggingface_hub", "manifest-ml", "elasticsearch", "opensearch-py", "google-search-results", "faiss-cpu", "sentence_transformers", "transformers", "spacy", "nltk", "wikipedia", "beautifulsoup4", "tiktoken", "torch", "jinja2", "pinecone-client", "weaviate-client", "redis", "google-api-python-client", "wolframalpha", "qdrant-client", "tensorflow-text", "pypdf", "networkx", "nomic", "aleph-alpha-client", "deeplake", "pgvector", "psycopg2-binary"]
 
 [tool.ruff]
 select = [

tests/integration_tests/vectorstores/test_pgvector.py

@@ -0,0 +1,81 @@
+"""Test PGVector functionality."""
+import os
+from typing import List
+
+from langchain.docstore.document import Document
+from langchain.vectorstores.pgvector import PGVector
+from tests.integration_tests.vectorstores.fake_embeddings import (
+    FakeEmbeddings,
+)
+
+CONNECTION_STRING = PGVector.connection_string_from_db_params(
+    driver=os.environ.get("TEST_PGVECTOR_DRIVER", "psycopg2"),
+    host=os.environ.get("TEST_PGVECTOR_HOST", "localhost"),
+    port=int(os.environ.get("TEST_PGVECTOR_PORT", "5432")),
+    database=os.environ.get("TEST_PGVECTOR_DATABASE", "postgres"),
+    user=os.environ.get("TEST_PGVECTOR_USER", "postgres"),
+    password=os.environ.get("TEST_PGVECTOR_PASSWORD", "postgres"),
+)
+
+
+ADA_TOKEN_COUNT = 1536
+
+
+class FakeEmbeddingsWithAdaDimension(FakeEmbeddings):
+    """Fake embeddings functionality for testing."""
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Return simple embeddings."""
+        return [
+            [float(1.0)] * (ADA_TOKEN_COUNT - 1) + [float(i)] for i in range(len(texts))
+        ]
+
+    def embed_query(self, text: str) -> List[float]:
+        """Return simple embeddings."""
+        return [float(1.0)] * (ADA_TOKEN_COUNT - 1) + [float(0.0)]
+
+
+def test_pgvector() -> None:
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    docsearch = PGVector.from_texts(
+        texts=texts,
+        collection_name="test_collection",
+        embedding=FakeEmbeddingsWithAdaDimension(),
+        connection_string=CONNECTION_STRING,
+        pre_delete_collection=True,
+    )
+    output = docsearch.similarity_search("foo", k=1)
+    assert output == [Document(page_content="foo")]
+
+
+def test_pgvector_with_metadatas() -> None:
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    metadatas = [{"page": str(i)} for i in range(len(texts))]
+    docsearch = PGVector.from_texts(
+        texts=texts,
+        collection_name="test_collection",
+        embedding=FakeEmbeddingsWithAdaDimension(),
+        metadatas=metadatas,
+        connection_string=CONNECTION_STRING,
+        pre_delete_collection=True,
+    )
+    output = docsearch.similarity_search("foo", k=1)
+    assert output == [Document(page_content="foo", metadata={"page": "0"})]
+
+
+def test_pgvector_with_metadatas_with_scores() -> None:
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    metadatas = [{"page": str(i)} for i in range(len(texts))]
+    docsearch = PGVector.from_texts(
+        texts=texts,
+        collection_name="test_collection",
+        embedding=FakeEmbeddingsWithAdaDimension(),
+        metadatas=metadatas,
+        connection_string=CONNECTION_STRING,
+        pre_delete_collection=True,
+    )
+    output = docsearch.similarity_search_with_score("foo", k=1)
+    assert output == [(Document(page_content="foo", metadata={"page": "0"}), 0.0)]