feat: allow non-uuid data types for vectorstore primary key

Author: dishaprakash

src/langchain_google_cloud_sql_pg/engine.py

@@ -410,12 +410,13 @@ async def _ainit_vectorstore_table(
         embedding_column: str = "embedding",
         metadata_columns: List[Column] = [],
         metadata_json_column: str = "langchain_metadata",
-        id_column: str = "langchain_id",
+        id_column: Union[str, Column] = "langchain_id",
         overwrite_existing: bool = False,
         store_metadata: bool = True,
     ) -> None:
         """
         Create a table for saving of vectors to be used with PostgresVectorStore.
+        If id column data type is invalid, an UNDEFINED_OBJECT_ERROR is thrown.
 
         Args:
             table_name (str): The Postgres database table name.
@@ -430,14 +431,14 @@ async def _ainit_vectorstore_table(
                 metadata. Default: []. Optional.
             metadata_json_column (str): The column to store extra metadata in JSON format.
                 Default: "langchain_metadata". Optional.
-            id_column (str):  Name of the column to store ids.
-                Default: "langchain_id". Optional,
+            id_column (Union[str, Column]) :  Column to store ids.
+                Default: "langchain_id" column name with data type UUID. Optional,
             overwrite_existing (bool): Whether to drop existing table. Default: False.
             store_metadata (bool): Whether to store metadata in the table.
                 Default: True.
-
         Raises:
             :class:`DuplicateTableError <asyncpg.exceptions.DuplicateTableError>`: if table already exists and overwrite flag is not set.
+            :class:`UndefinedObjectError <asyncpg.exceptions.UndefinedObjectError>`: if id column data type is invalid.
         """
         async with self._pool.connect() as conn:
             await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
@@ -450,8 +451,11 @@ async def _ainit_vectorstore_table(
                 )
                 await conn.commit()
 
+        id_data_type = "UUID" if isinstance(id_column, str) else id_column.data_type
+        id_column_name = id_column if isinstance(id_column, str) else id_column.name
+
         query = f"""CREATE TABLE "{schema_name}"."{table_name}"(
-            "{id_column}" UUID PRIMARY KEY,
+            "{id_column_name}" {id_data_type} PRIMARY KEY,
             "{content_column}" TEXT NOT NULL,
             "{embedding_column}" vector({vector_size}) NOT NULL"""
         for column in metadata_columns:
@@ -524,12 +528,13 @@ def init_vectorstore_table(
         embedding_column: str = "embedding",
         metadata_columns: List[Column] = [],
         metadata_json_column: str = "langchain_metadata",
-        id_column: str = "langchain_id",
+        id_column: Union[str, Column] = "langchain_id",
         overwrite_existing: bool = False,
         store_metadata: bool = True,
     ) -> None:
         """
         Create a table for saving of vectors to be used with PostgresVectorStore.
+        If id column data type is invalid, an UNDEFINED_OBJECT_ERROR is thrown.
 
         Args:
             table_name (str): The Postgres database table name.
@@ -544,11 +549,13 @@ def init_vectorstore_table(
                 metadata. Default: []. Optional.
             metadata_json_column (str): The column to store extra metadata in JSON format.
                 Default: "langchain_metadata". Optional.
-            id_column (str):  Name of the column to store ids.
+            id_column (Union[str, Column]):  Name of the column to store ids.
                 Default: "langchain_id". Optional,
             overwrite_existing (bool): Whether to drop existing table. Default: False.
             store_metadata (bool): Whether to store metadata in the table.
                 Default: True.
+        Raises:
+            :class:`UndefinedObjectError <asyncpg.exceptions.UndefinedObjectError>`: if id column data type is invalid.
         """
         self._run_as_sync(
             self._ainit_vectorstore_table(

src/langchain_google_cloud_sql_pg/vectorstore.py

@@ -187,10 +187,10 @@ async def aadd_texts(
         self,
         texts: Iterable[str],
         metadatas: Optional[List[dict]] = None,
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         **kwargs: Any,
     ) -> List[str]:
-        """Embed texts and add to the table."""
+        """Embed texts and add to the table. Throws an error if Id column data type doesn't match with the ids."""
         return await self._engine._run_as_async(
             self.__vs.aadd_texts(texts, metadatas, ids, **kwargs)
         )
@@ -199,50 +199,50 @@ def add_texts(
         self,
         texts: Iterable[str],
         metadatas: Optional[List[dict]] = None,
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         **kwargs: Any,
     ) -> List[str]:
-        """Embed texts and add to the table."""
+        """Embed texts and add to the table. Throws an error if Id column data type doesn't match with the ids."""
         return self._engine._run_as_sync(
             self.__vs.aadd_texts(texts, metadatas, ids, **kwargs)
         )
 
     async def aadd_documents(
         self,
         documents: List[Document],
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         **kwargs: Any,
     ) -> List[str]:
-        """Embed documents and add to the table"""
+        """Embed documents and add to the table. Throws an error if Id column data type doesn't match with the ids."""
         return await self._engine._run_as_async(
             self.__vs.aadd_documents(documents, ids, **kwargs)
         )
 
     def add_documents(
         self,
         documents: List[Document],
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         **kwargs: Any,
     ) -> List[str]:
-        """Embed documents and add to the table."""
+        """Embed documents and add to the table. Throws an error if Id column data type doesn't match with the ids."""
         return self._engine._run_as_sync(
             self.__vs.aadd_documents(documents, ids, **kwargs)
         )
 
     async def adelete(
         self,
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         **kwargs: Any,
     ) -> Optional[bool]:
-        """Delete records from the table."""
+        """Delete records from the table. Throws an error if Id column data type doesn't match with the ids."""
         return await self._engine._run_as_async(self.__vs.adelete(ids, **kwargs))
 
     def delete(
         self,
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         **kwargs: Any,
     ) -> Optional[bool]:
-        """Delete records from the table."""
+        """Delete records from the table. Throws an error if Id column data type doesn't match with the ids."""
         return self._engine._run_as_sync(self.__vs.adelete(ids, **kwargs))
 
     @classmethod
@@ -254,7 +254,7 @@ async def afrom_texts(  # type: ignore[override]
         table_name: str,
         schema_name: str = "public",
         metadatas: Optional[List[dict]] = None,
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         content_column: str = "content",
         embedding_column: str = "embedding",
         metadata_columns: List[str] = [],
@@ -268,14 +268,16 @@ async def afrom_texts(  # type: ignore[override]
         index_query_options: Optional[QueryOptions] = None,
     ) -> PostgresVectorStore:
         """Create an PostgresVectorStore instance from texts.
+        Throws an error if Id column data type doesn't match with the ids.
+
         Args:
             texts (List[str]): Texts to add to the vector store.
             embedding (Embeddings): Text embedding model to use.
             engine (PostgresEngine): Connection pool engine for managing connections to Postgres database.
             table_name (str): Name of the existing table or the table to be created.
             schema_name (str, optional): Database schema name of the table. Defaults to "public".
             metadatas (Optional[List[dict]]): List of metadatas to add to table records.
-            ids: (Optional[List[str]]): List of IDs to add to table records.
+            ids: (Optional[List]): List of IDs to add to table records.
             content_column (str): Column that represent a Document’s page_content. Defaults to "content".
             embedding_column (str): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding".
             metadata_columns (List[str]): Column(s) that represent a document's metadata.
@@ -319,7 +321,7 @@ async def afrom_documents(  # type: ignore[override]
         engine: PostgresEngine,
         table_name: str,
         schema_name: str = "public",
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         content_column: str = "content",
         embedding_column: str = "embedding",
         metadata_columns: List[str] = [],
@@ -333,6 +335,7 @@ async def afrom_documents(  # type: ignore[override]
         index_query_options: Optional[QueryOptions] = None,
     ) -> PostgresVectorStore:
         """Create an PostgresVectorStore instance from documents.
+        Throws an error if Id column data type doesn't match with the ids.
 
         Args:
             documents (List[Document]): Documents to add to the vector store.
@@ -341,7 +344,7 @@ async def afrom_documents(  # type: ignore[override]
             table_name (str): Name of the existing table or the table to be created.
             schema_name (str, optional): Database schema name of the table. Defaults to "public".
             metadatas (Optional[List[dict]]): List of metadatas to add to table records.
-            ids: (Optional[List[str]]): List of IDs to add to table records.
+            ids: (Optional[List]): List of IDs to add to table records.
             content_column (str): Column that represent a Document’s page_content. Defaults to "content".
             embedding_column (str): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding".
             metadata_columns (List[str]): Column(s) that represent a document's metadata.
@@ -386,7 +389,7 @@ def from_texts(  # type: ignore[override]
         table_name: str,
         schema_name: str = "public",
         metadatas: Optional[List[dict]] = None,
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         content_column: str = "content",
         embedding_column: str = "embedding",
         metadata_columns: List[str] = [],
@@ -400,14 +403,16 @@ def from_texts(  # type: ignore[override]
         index_query_options: Optional[QueryOptions] = None,
     ) -> PostgresVectorStore:
         """Create an PostgresVectorStore instance from texts.
+        Throws an error if Id column data type doesn't match with the ids.
+
         Args:
             texts (List[str]): Texts to add to the vector store.
             embedding (Embeddings): Text embedding model to use.
             engine (PostgresEngine): Connection pool engine for managing connections to Postgres database.
             table_name (str): Name of the existing table or the table to be created.
             schema_name (str, optional): Database schema name of the table. Defaults to "public".
             metadatas (Optional[List[dict]]): List of metadatas to add to table records.
-            ids: (Optional[List[str]]): List of IDs to add to table records.
+            ids: (Optional[List]): List of IDs to add to table records.
             content_column (str): Column that represent a Document’s page_content. Defaults to "content".
             embedding_column (str): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding".
             metadata_columns (List[str]): Column(s) that represent a document's metadata.
@@ -451,7 +456,7 @@ def from_documents(  # type: ignore[override]
         engine: PostgresEngine,
         table_name: str,
         schema_name: str = "public",
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         content_column: str = "content",
         embedding_column: str = "embedding",
         metadata_columns: List[str] = [],
@@ -465,6 +470,7 @@ def from_documents(  # type: ignore[override]
         index_query_options: Optional[QueryOptions] = None,
     ) -> PostgresVectorStore:
         """Create an PostgresVectorStore instance from documents.
+        Throws an error if Id column data type doesn't match with the ids.
 
         Args:
             documents (List[Document]): Documents to add to the vector store.
@@ -473,7 +479,7 @@ def from_documents(  # type: ignore[override]
             table_name (str): Name of the existing table or the table to be created.
             schema_name (str, optional): Database schema name of the table. Defaults to "public".
             metadatas (Optional[List[dict]]): List of metadatas to add to table records.
-            ids: (Optional[List[str]]): List of IDs to add to table records.
+            ids: (Optional[List]): List of IDs to add to table records.
             content_column (str): Column that represent a Document’s page_content. Defaults to "content".
             embedding_column (str): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding".
             metadata_columns (List[str]): Column(s) that represent a document's metadata.

tests/test_engine.py

@@ -31,8 +31,10 @@
 
 DEFAULT_TABLE = "test_table" + str(uuid.uuid4()).replace("-", "_")
 CUSTOM_TABLE = "test_table_custom" + str(uuid.uuid4()).replace("-", "_")
+INT_ID_CUSTOM_TABLE = "test_table_custom_int_id" + str(uuid.uuid4()).replace("-", "_")
 DEFAULT_TABLE_SYNC = "test_table" + str(uuid.uuid4()).replace("-", "_")
 CUSTOM_TABLE_SYNC = "test_table_custom" + str(uuid.uuid4()).replace("-", "_")
+INT_ID_CUSTOM_TABLE_SYNC = "test_table_custom_int_id" + str(uuid.uuid4()).replace("-", "_")
 VECTOR_SIZE = 768
 
 embeddings_service = DeterministicFakeEmbedding(size=VECTOR_SIZE)
@@ -110,6 +112,7 @@ async def engine(self, db_project, db_region, db_instance, db_name):
         yield engine
         await aexecute(engine, f'DROP TABLE "{CUSTOM_TABLE}"')
         await aexecute(engine, f'DROP TABLE "{DEFAULT_TABLE}"')
+        await aexecute(engine, f'DROP TABLE "{INT_ID_CUSTOM_TABLE}"')
         await engine.close()
 
     async def test_init_table(self, engine):
@@ -143,6 +146,29 @@ async def test_init_table_custom(self, engine):
         for row in results:
             assert row in expected
 
+    async def test_init_table_with_int_id(self, engine):
+        await engine.ainit_vectorstore_table(
+            INT_ID_CUSTOM_TABLE,
+            VECTOR_SIZE,
+            id_column=Column(name="integer_id", data_type="INTEGER", nullable="False"),
+            content_column="my-content",
+            embedding_column="my_embedding",
+            metadata_columns=[Column("page", "TEXT"), Column("source", "TEXT")],
+            store_metadata=True,
+        )
+        stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{INT_ID_CUSTOM_TABLE}';"
+        results = await afetch(engine, stmt)
+        expected = [
+            {"column_name": "integer_id", "data_type": "integer"},
+            {"column_name": "my_embedding", "data_type": "USER-DEFINED"},
+            {"column_name": "langchain_metadata", "data_type": "json"},
+            {"column_name": "my-content", "data_type": "text"},
+            {"column_name": "page", "data_type": "text"},
+            {"column_name": "source", "data_type": "text"},
+        ]
+        for row in results:
+            assert row in expected
+
     async def test_password(
         self,
         db_project,
@@ -306,6 +332,7 @@ async def engine(self, db_project, db_region, db_instance, db_name):
         yield engine
         await aexecute(engine, f'DROP TABLE "{CUSTOM_TABLE_SYNC}"')
         await aexecute(engine, f'DROP TABLE "{DEFAULT_TABLE_SYNC}"')
+        await aexecute(engine, f'DROP TABLE "{INT_ID_CUSTOM_TABLE_SYNC}"')
         await engine.close()
 
     async def test_init_table(self, engine):
@@ -339,6 +366,29 @@ async def test_init_table_custom(self, engine):
         for row in results:
             assert row in expected
 
+    async def test_init_table_with_int_id(self, engine):
+        engine.init_vectorstore_table(
+            INT_ID_CUSTOM_TABLE_SYNC,
+            VECTOR_SIZE,
+            id_column=Column(name="integer_id", data_type="INTEGER", nullable=False),
+            content_column="my-content",
+            embedding_column="my_embedding",
+            metadata_columns=[Column("page", "TEXT"), Column("source", "TEXT")],
+            store_metadata=True,
+        )
+        stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{INT_ID_CUSTOM_TABLE_SYNC}';"
+        results = await afetch(engine, stmt)
+        expected = [
+            {"column_name": "integer_id", "data_type": "integer"},
+            {"column_name": "my_embedding", "data_type": "USER-DEFINED"},
+            {"column_name": "langchain_metadata", "data_type": "json"},
+            {"column_name": "my-content", "data_type": "text"},
+            {"column_name": "page", "data_type": "text"},
+            {"column_name": "source", "data_type": "text"},
+        ]
+        for row in results:
+            assert row in expected
+
     async def test_password(
         self,
         db_project,