Skip to content

Commit e92af36

Browse files
authored
fix: Add support for a metadata column named id (#421)
* fix: Add support for a metadata column named id * Add test * remove test to debug * Add test
1 parent 8cb35ca commit e92af36

File tree

2 files changed

+27
-3
lines changed

2 files changed

+27
-3
lines changed

src/langchain_google_alloydb_pg/async_vectorstore.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -281,16 +281,16 @@ async def aadd_embeddings(
281281
)
282282
insert_stmt = f'INSERT INTO "{self.schema_name}"."{self.table_name}"("{self.id_column}", "{self.content_column}", "{self.embedding_column}"{metadata_col_names}'
283283
values = {
284-
"id": id,
284+
"langchain_id": id,
285285
"content": content,
286286
"embedding": str([float(dimension) for dimension in embedding]),
287287
}
288-
values_stmt = "VALUES (:id, :content, :embedding"
288+
values_stmt = "VALUES (:langchain_id, :content, :embedding"
289289
inline_embed_func = getattr(
290290
self.embedding_service, "embed_query_inline", None
291291
)
292292
if not embedding and callable(inline_embed_func):
293-
values_stmt = f"VALUES (:id, :content, {self.embedding_service.embed_query_inline(content)}" # type: ignore
293+
values_stmt = f"VALUES (:langchain_id, :content, {self.embedding_service.embed_query_inline(content)}" # type: ignore
294294

295295
# Add metadata
296296
extra = copy.deepcopy(metadata)

tests/test_async_vectorstore.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
docs = [
4141
Document(page_content=texts[i], metadata=metadatas[i]) for i in range(len(texts))
4242
]
43+
id_column_as_metadata = [{"id": str(i)} for i in range(len(texts))]
4344

4445
embeddings = [embeddings_service.embed_query(texts[i]) for i in range(len(texts))]
4546

@@ -209,6 +210,29 @@ async def test_post_init(self, engine):
209210
metadata_json_column="mymeta",
210211
)
211212

213+
async def test_id_metadata_column(self, engine):
214+
table_name = "id_metadata" + str(uuid.uuid4())
215+
await engine._ainit_vectorstore_table(
216+
table_name,
217+
VECTOR_SIZE,
218+
metadata_columns=[Column("id", "TEXT")],
219+
)
220+
custom_vs = await AsyncAlloyDBVectorStore.create(
221+
engine,
222+
embedding_service=embeddings_service,
223+
table_name=table_name,
224+
metadata_columns=["id"],
225+
)
226+
ids = [str(uuid.uuid4()) for i in range(len(texts))]
227+
await custom_vs.aadd_texts(texts, id_column_as_metadata, ids)
228+
229+
results = await afetch(engine, f'SELECT * FROM "{table_name}"')
230+
assert len(results) == 3
231+
assert results[0]["id"] == "0"
232+
assert results[1]["id"] == "1"
233+
assert results[2]["id"] == "2"
234+
await aexecute(engine, f'DROP TABLE IF EXISTS "{table_name}"')
235+
212236
async def test_aadd_texts(self, engine, vs):
213237
ids = [str(uuid.uuid4()) for i in range(len(texts))]
214238
await vs.aadd_texts(texts, ids=ids)

0 commit comments

Comments
 (0)