feat: Add the PGVectorStore class

langchain_postgres/__init__.py

@@ -3,7 +3,7 @@
 from langchain_postgres.chat_message_histories import PostgresChatMessageHistory
 from langchain_postgres.engine import Column, PGEngine
 from langchain_postgres.translator import PGVectorTranslator
-from langchain_postgres.vectorstore import PGVectorStore
+from langchain_postgres.vectorstore.v2 import PGVectorStore
 from langchain_postgres.vectorstores import PGVector
 
 try:

langchain_postgres/chat_message_histories.py

@@ -2,7 +2,6 @@
 
 This client provides support for both sync and async via psycopg 3.
 """
-
 from __future__ import annotations
 
 import json

langchain_postgres/engine.py

@@ -2,12 +2,12 @@
 
 import asyncio
 from dataclasses import dataclass
+import re
 from threading import Thread
 from typing import TYPE_CHECKING, Any, Awaitable, Optional, TypeVar, Union
 
-from sqlalchemy import MetaData, Table, text
+from sqlalchemy import text
 from sqlalchemy.engine import URL
-from sqlalchemy.exc import InvalidRequestError
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 
 if TYPE_CHECKING:
@@ -54,7 +54,7 @@ def __init__(
 
         Args:
             key (object): Prevent direct constructor usage.
-            engine (AsyncEngine): Async engine connection pool.
+            pool (AsyncEngine): Async engine connection pool.
             loop (Optional[asyncio.AbstractEventLoop]): Async event loop used to create the engine.
             thread (Optional[Thread]): Thread used to create the engine async.
 
@@ -136,14 +136,18 @@ async def close(self) -> None:
         """Dispose of connection pool"""
         await self._run_as_async(self._pool.dispose())
 
+    def escape_postgres_identifier(self, name: str) -> str:
+        return name.replace('"', '""')
+
     async def _ainit_vectorstore_table(
         self,
         table_name: str,
         vector_size: int,
+        *,
         schema_name: str = "public",
         content_column: str = "content",
         embedding_column: str = "embedding",
-        metadata_columns: list[Column] = [],
+        metadata_columns: Optional[list[Column]] = None,
         metadata_json_column: str = "langchain_metadata",
         id_column: Union[str, Column] = "langchain_id",
         overwrite_existing: bool = False,
@@ -161,8 +165,8 @@ async def _ainit_vectorstore_table(
                 Default: "page_content".
             embedding_column (str) : Name of the column to store vector embeddings.
                 Default: "embedding".
-            metadata_columns (list[Column]): A list of Columns to create for custom
-                metadata. Default: []. Optional.
+            metadata_columns (Optional[list[Column]]): A list of Columns to create for custom
+                metadata. Default: None. Optional.
             metadata_json_column (str): The column to store extra metadata in JSON format.
                 Default: "langchain_metadata". Optional.
             id_column (Union[str, Column]) :  Column to store ids.
@@ -175,6 +179,21 @@ async def _ainit_vectorstore_table(
             :class:`DuplicateTableError <asyncpg.exceptions.DuplicateTableError>`: if table already exists.
             :class:`UndefinedObjectError <asyncpg.exceptions.UndefinedObjectError>`: if the data type of the id column is not a postgreSQL data type.
         """
+
+        schema_name = self.escape_postgres_identifier(schema_name)
+        table_name = self.escape_postgres_identifier(table_name)
+        content_column = self.escape_postgres_identifier(content_column)
+        embedding_column = self.escape_postgres_identifier(embedding_column)
+        if metadata_columns is None:
+            metadata_columns = []
+        else:
+            for col in metadata_columns:
+                col.name = self.escape_postgres_identifier(col.name)
+        if isinstance(id_column, str):
+            id_column = self.escape_postgres_identifier(id_column)
+        else:
+            id_column.name = self.escape_postgres_identifier(id_column.name)
+
         async with self._pool.connect() as conn:
             await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
             await conn.commit()
@@ -208,10 +227,11 @@ async def ainit_vectorstore_table(
         self,
         table_name: str,
         vector_size: int,
+        *,
         schema_name: str = "public",
         content_column: str = "content",
         embedding_column: str = "embedding",
-        metadata_columns: list[Column] = [],
+        metadata_columns: Optional[list[Column]] = None,
         metadata_json_column: str = "langchain_metadata",
         id_column: Union[str, Column] = "langchain_id",
         overwrite_existing: bool = False,
@@ -229,8 +249,8 @@ async def ainit_vectorstore_table(
                 Default: "page_content".
             embedding_column (str) : Name of the column to store vector embeddings.
                 Default: "embedding".
-            metadata_columns (list[Column]): A list of Columns to create for custom
-                metadata. Default: []. Optional.
+            metadata_columns (Optional[list[Column]]): A list of Columns to create for custom
+                metadata. Default: None. Optional.
             metadata_json_column (str): The column to store extra metadata in JSON format.
                 Default: "langchain_metadata". Optional.
             id_column (Union[str, Column]) :  Column to store ids.
@@ -243,25 +263,26 @@ async def ainit_vectorstore_table(
             self._ainit_vectorstore_table(
                 table_name,
                 vector_size,
-                schema_name,
-                content_column,
-                embedding_column,
-                metadata_columns,
-                metadata_json_column,
-                id_column,
-                overwrite_existing,
-                store_metadata,
+                schema_name=schema_name,
+                content_column=content_column,
+                embedding_column=embedding_column,
+                metadata_columns=metadata_columns,
+                metadata_json_column=metadata_json_column,
+                id_column=id_column,
+                overwrite_existing=overwrite_existing,
+                store_metadata=store_metadata,
             )
         )
 
     def init_vectorstore_table(
         self,
         table_name: str,
         vector_size: int,
+        *,
         schema_name: str = "public",
         content_column: str = "content",
         embedding_column: str = "embedding",
-        metadata_columns: list[Column] = [],
+        metadata_columns: Optional[list[Column]] = None,
         metadata_json_column: str = "langchain_metadata",
         id_column: Union[str, Column] = "langchain_id",
         overwrite_existing: bool = False,
@@ -279,8 +300,8 @@ def init_vectorstore_table(
                 Default: "page_content".
             embedding_column (str) : Name of the column to store vector embeddings.
                 Default: "embedding".
-            metadata_columns (list[Column]): A list of Columns to create for custom
-                metadata. Default: []. Optional.
+            metadata_columns (Optional[list[Column]]): A list of Columns to create for custom
+                metadata. Default: None. Optional.
             metadata_json_column (str): The column to store extra metadata in JSON format.
                 Default: "langchain_metadata". Optional.
             id_column (Union[str, Column]) :  Column to store ids.
@@ -293,53 +314,13 @@ def init_vectorstore_table(
             self._ainit_vectorstore_table(
                 table_name,
                 vector_size,
-                schema_name,
-                content_column,
-                embedding_column,
-                metadata_columns,
-                metadata_json_column,
-                id_column,
-                overwrite_existing,
-                store_metadata,
+                schema_name=schema_name,
+                content_column=content_column,
+                embedding_column=embedding_column,
+                metadata_columns=metadata_columns,
+                metadata_json_column=metadata_json_column,
+                id_column=id_column,
+                overwrite_existing=overwrite_existing,
+                store_metadata=store_metadata,
             )
         )
-
-    async def _aload_table_schema(
-        self, table_name: str, schema_name: str = "public"
-    ) -> Table:
-        """
-        Load table schema from an existing table in a PgSQL database, potentially from a specific database schema.
-
-        Args:
-            table_name: The name of the table to load the table schema from.
-            schema_name: The name of the database schema where the table resides.
-                Default: "public".
-
-        Returns:
-            (sqlalchemy.Table): The loaded table, including its table schema information.
-        """
-        metadata = MetaData()
-        async with self._pool.connect() as conn:
-            try:
-                await conn.run_sync(
-                    metadata.reflect, schema=schema_name, only=[table_name]
-                )
-            except InvalidRequestError as e:
-                raise ValueError(
-                    f"Table, '{schema_name}'.'{table_name}', does not exist: " + str(e)
-                )
-
-        table = Table(table_name, metadata, schema=schema_name)
-        # Extract the schema information
-        schema = []
-        for column in table.columns:
-            schema.append(
-                {
-                    "name": column.name,
-                    "type": column.type.python_type,
-                    "max_length": getattr(column.type, "length", None),
-                    "nullable": not column.nullable,
-                }
-            )
-
-        return metadata.tables[f"{schema_name}.{table_name}"]

langchain_postgres/indexes.py

@@ -1,4 +1,10 @@
+"""Index class to add vector indexes on the PGVectorStore.
+
+Learn more about vector indexes at https://github.com/pgvector/pgvector?tab=readme-ov-file#indexing
+"""
+
 import enum
+import re
 import warnings
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
@@ -26,6 +32,18 @@ class DistanceStrategy(StrategyMixin, enum.Enum):
 
 @dataclass
 class BaseIndex(ABC):
+    """
+    Abstract base class for defining vector indexes.
+
+    Attributes:
+        name (Optional[str]): A human-readable name for the index. Defaults to None.
+        index_type (str): A string identifying the type of index. Defaults to "base".
+        distance_strategy (DistanceStrategy): The strategy used to calculate distances
+            between vectors in the index. Defaults to DistanceStrategy.COSINE_DISTANCE.
+        partial_indexes (Optional[list[str]]): A list of names of partial indexes. Defaults to None.
+        extension_name (Optional[str]): The name of the extension to be created for the index, if any. Defaults to None.
+    """
+
     name: Optional[str] = None
     index_type: str = "base"
     distance_strategy: DistanceStrategy = field(
@@ -44,6 +62,20 @@ def index_options(self) -> str:
     def get_index_function(self) -> str:
         return self.distance_strategy.index_function
 
+    def __post_init__(self) -> None:
+        """Check if initialization parameters are valid.
+
+        Raises:
+            ValueError: extension_name is a valid postgreSQL identifier
+        """
+        if (
+            self.extension_name
+            and re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", self.extension_name) is None
+        ):
+            raise ValueError(
+                f"Invalid identifier: {self.extension_name}. Identifiers must start with a letter or underscore, and subsequent characters can be letters, digits, or underscores."
+            )
+
 
 @dataclass
 class ExactNearestNeighbor(BaseIndex):

langchain_postgres/utils/pgvector_migrator.py

@@ -7,7 +7,7 @@
 from sqlalchemy.exc import ProgrammingError, SQLAlchemyError
 
 from ..engine import PGEngine
-from ..vectorstore import PGVectorStore
+from ..vectorstore.v2 import PGVectorStore
 
 COLLECTIONS_TABLE = "langchain_pg_collection"
 EMBEDDINGS_TABLE = "langchain_pg_embedding"