feat: allow non-uuid data types for vectorstore primary key (#209)

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

* Update src/langchain_google_cloud_sql_pg/engine.py

* Update src/langchain_google_cloud_sql_pg/engine.py

* Update src/langchain_google_cloud_sql_pg/engine.py

---------

Co-authored-by: Averi Kitsch <akitsch@google.com>

Author: dishaprakash

src/langchain_google_cloud_sql_pg/async_vectorstore.py

@@ -226,10 +226,14 @@ async def __aadd_embeddings(
         texts: Iterable[str],
         embeddings: List[List[float]],
         metadatas: Optional[List[dict]] = None,
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         **kwargs: Any,
     ) -> List[str]:
-        """Add embeddings to the table."""
+        """Add embeddings to the table.
+
+        Raises:
+            :class:`InvalidTextRepresentationError <asyncpg.exceptions.InvalidTextRepresentationError>`: if the `ids` data type does not match that of the `id_column`.
+        """
         if not ids:
             ids = [str(uuid.uuid4()) for _ in texts]
         if not metadatas:
@@ -276,10 +280,14 @@ 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.
+
+        Raises:
+            :class:`InvalidTextRepresentationError <asyncpg.exceptions.InvalidTextRepresentationError>`: if the `ids` data type does not match that of the `id_column`.
+        """
         embeddings = self.embedding_service.embed_documents(list(texts))
         ids = await self.__aadd_embeddings(
             texts, embeddings, metadatas=metadatas, ids=ids, **kwargs
@@ -289,21 +297,29 @@ async def aadd_texts(
     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.
+
+        Raises:
+            :class:`InvalidTextRepresentationError <asyncpg.exceptions.InvalidTextRepresentationError>`: if the `ids` data type does not match that of the `id_column`.
+        """
         texts = [doc.page_content for doc in documents]
         metadatas = [doc.metadata for doc in documents]
         ids = await self.aadd_texts(texts, metadatas=metadatas, ids=ids, **kwargs)
         return ids
 
     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.
+
+        Raises:
+            :class:`InvalidTextRepresentationError <asyncpg.exceptions.InvalidTextRepresentationError>`: if the `ids` data type does not match that of the `id_column`.
+        """
         if not ids:
             return False
 
@@ -323,7 +339,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] = [],
@@ -338,6 +354,7 @@ async def afrom_texts(  # type: ignore[override]
         **kwargs: Any,
     ) -> AsyncPostgresVectorStore:
         """Create an AsyncPostgresVectorStore instance from texts.
+
         Args:
             texts (List[str]): Texts to add to the vector store.
             embedding (Embeddings): Text embedding model to use.
@@ -358,6 +375,9 @@ async def afrom_texts(  # type: ignore[override]
             lambda_mult (float): Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding to maximum diversity and 1 to minimum diversity. Defaults to 0.5.
             index_query_options (QueryOptions): Index query option.
 
+        Raises:
+            :class:`InvalidTextRepresentationError <asyncpg.exceptions.InvalidTextRepresentationError>`: if the `ids` data type does not match that of the `id_column`.
+
         Returns:
             AsyncPostgresVectorStore
         """
@@ -389,7 +409,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] = [],
@@ -425,6 +445,9 @@ async def afrom_documents(  # type: ignore[override]
             lambda_mult (float): Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding to maximum diversity and 1 to minimum diversity. Defaults to 0.5.
             index_query_options (QueryOptions): Index query option.
 
+        Raises:
+            :class:`InvalidTextRepresentationError <asyncpg.exceptions.InvalidTextRepresentationError>`: if the `ids` data type does not match that of the `id_column`.
+
         Returns:
             AsyncPostgresVectorStore
         """
@@ -735,7 +758,7 @@ 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]:
         raise NotImplementedError(
@@ -745,7 +768,7 @@ def add_texts(
     def add_documents(
         self,
         documents: List[Document],
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         **kwargs: Any,
     ) -> List[str]:
         raise NotImplementedError(
@@ -754,7 +777,7 @@ def add_documents(
 
     def delete(
         self,
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         **kwargs: Any,
     ) -> Optional[bool]:
         raise NotImplementedError(
@@ -769,7 +792,7 @@ def from_texts(  # type: ignore[override]
         engine: PostgresEngine,
         table_name: str,
         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] = [],
@@ -789,7 +812,7 @@ def from_documents(  # type: ignore[override]
         embedding: Embeddings,
         engine: PostgresEngine,
         table_name: str,
-        ids: Optional[List[str]] = None,
+        ids: Optional[List] = None,
         content_column: str = "content",
         embedding_column: str = "embedding",
         metadata_columns: List[str] = [],

src/langchain_google_cloud_sql_pg/engine.py

@@ -410,7 +410,7 @@ 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:
@@ -430,14 +430,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 the data type of the id column is not a postgreSQL data type.
         """
         async with self._pool.connect() as conn:
             await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
@@ -450,8 +450,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:
@@ -474,7 +477,7 @@ 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:
@@ -494,8 +497,8 @@ 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.
@@ -524,7 +527,7 @@ 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:
@@ -544,11 +547,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.
-                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:`UndefinedObjectError <asyncpg.exceptions.UndefinedObjectError>`: if the `ids` data type does not match that of the `id_column`.
         """
         self._run_as_sync(
             self._ainit_vectorstore_table(