|
40 | 40 | docs = [ |
41 | 41 | Document(page_content=texts[i], metadata=metadatas[i]) for i in range(len(texts)) |
42 | 42 | ] |
| 43 | +id_column_as_metadata = [{"id": str(i)} for i in range(len(texts))] |
43 | 44 |
|
44 | 45 | embeddings = [embeddings_service.embed_query(texts[i]) for i in range(len(texts))] |
45 | 46 |
|
@@ -209,6 +210,29 @@ async def test_post_init(self, engine): |
209 | 210 | metadata_json_column="mymeta", |
210 | 211 | ) |
211 | 212 |
|
| 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 | + |
212 | 236 | async def test_aadd_texts(self, engine, vs): |
213 | 237 | ids = [str(uuid.uuid4()) for i in range(len(texts))] |
214 | 238 | await vs.aadd_texts(texts, ids=ids) |
|
0 commit comments