From 11fb1baed99a341b455732fd6955709f8469831b Mon Sep 17 00:00:00 2001
From: Disha Prakash <dishaprakash@google.com>
Date: Wed, 27 Nov 2024 10:33:54 +0000
Subject: [PATCH] Added warning if struct_id is not mentioned in
 aget_index_struct.

---
 src/llama_index_alloydb_pg/async_index_store.py |  2 ++
 tests/test_async_index_store.py                 | 14 ++++++++++++++
 2 files changed, 16 insertions(+)

diff --git a/src/llama_index_alloydb_pg/async_index_store.py b/src/llama_index_alloydb_pg/async_index_store.py
index 097e645..d7f043a 100644
--- a/src/llama_index_alloydb_pg/async_index_store.py
+++ b/src/llama_index_alloydb_pg/async_index_store.py
@@ -16,6 +16,7 @@
 
 import json
 from typing import List, Optional
+import warnings
 
 from llama_index.core.constants import DATA_KEY
 from llama_index.core.data_structs.data_structs import IndexStruct
@@ -174,6 +175,7 @@ async def aget_index_struct(
             structs = await self.aindex_structs()
             if len(structs) == 1:
                 return structs[0]
+            warnings.warn("No struct_id specified and more than one struct exists.")
             return None
         else:
             query = f"""SELECT * from "{self._schema_name}"."{self._table_name}" WHERE index_id = '{struct_id}';"""
diff --git a/tests/test_async_index_store.py b/tests/test_async_index_store.py
index 8dd10b7..a940803 100644
--- a/tests/test_async_index_store.py
+++ b/tests/test_async_index_store.py
@@ -15,6 +15,7 @@
 import os
 import uuid
 from typing import Sequence
+import warnings
 
 import pytest
 import pytest_asyncio
@@ -153,3 +154,16 @@ async def test_aindex_structs(self, index_store):
         assert index_dict_struct in indexes
         assert index_list_struct in indexes
         assert index_graph_struct in indexes
+
+    async def test_warning(self, index_store):
+        index_dict_struct = IndexDict()
+        index_list_struct = IndexList()
+
+        await index_store.aadd_index_struct(index_dict_struct)
+        await index_store.aadd_index_struct(index_list_struct)
+
+        with warnings.catch_warnings(record=True) as w:
+          index_struct = await index_store.aget_index_struct()
+
+          assert len(w) == 1
+          assert "No struct_id specified and more than one struct exists." in str(w[-1].message)