Add storageurl field if missing in index (#1556)

* Add storageurl field

* Add tests

Author: pamelafox

app/backend/prepdocslib/searchmanager.py

@@ -187,6 +187,18 @@ async def create_index(self, vectorizers: Optional[List[VectorSearchVectorizer]]
                 await search_index_client.create_index(index)
             else:
                 logger.info("Search index %s already exists", self.search_info.index_name)
+                index_definition = await search_index_client.get_index(self.search_info.index_name)
+                if not any(field.name == "storageUrl" for field in index_definition.fields):
+                    logger.info("Adding storageUrl field to index %s", self.search_info.index_name)
+                    index_definition.fields.append(
+                        SimpleField(
+                            name="storageUrl",
+                            type="Edm.String",
+                            filterable=True,
+                            facetable=False,
+                        ),
+                    )
+                    await search_index_client.create_or_update_index(index_definition)
 
     async def update_content(
         self, sections: List[Section], image_embeddings: Optional[List[List[float]]] = None, url: Optional[str] = None

tests/test_searchmanager.py

@@ -6,6 +6,11 @@
 from azure.core.credentials import AzureKeyCredential
 from azure.search.documents.aio import SearchClient
 from azure.search.documents.indexes.aio import SearchIndexClient
+from azure.search.documents.indexes.models import (
+    SearchFieldDataType,
+    SearchIndex,
+    SimpleField,
+)
 from openai.types.create_embedding_response import Usage
 
 from prepdocslib.embeddings import AzureOpenAIEmbeddingService
@@ -75,20 +80,72 @@ async def mock_list_index_names(self):
 
 @pytest.mark.asyncio
 async def test_create_index_does_exist(monkeypatch, search_info):
-    indexes = []
+    created_indexes = []
+    updated_indexes = []
 
     async def mock_create_index(self, index):
-        indexes.append(index)
+        created_indexes.append(index)
+
+    async def mock_list_index_names(self):
+        yield "test"
+
+    async def mock_get_index(self, *args, **kwargs):
+        return SearchIndex(
+            name="test",
+            fields=[
+                SimpleField(
+                    name="storageUrl",
+                    type=SearchFieldDataType.String,
+                    filterable=True,
+                )
+            ],
+        )
+
+    async def mock_create_or_update_index(self, index, *args, **kwargs):
+        updated_indexes.append(index)
+
+    monkeypatch.setattr(SearchIndexClient, "create_index", mock_create_index)
+    monkeypatch.setattr(SearchIndexClient, "list_index_names", mock_list_index_names)
+    monkeypatch.setattr(SearchIndexClient, "get_index", mock_get_index)
+    monkeypatch.setattr(SearchIndexClient, "create_or_update_index", mock_create_or_update_index)
+
+    manager = SearchManager(search_info)
+    await manager.create_index()
+    assert len(created_indexes) == 0, "It should not have created a new index"
+    assert len(updated_indexes) == 0, "It should not have updated the existing index"
+
+
+@pytest.mark.asyncio
+async def test_create_index_add_field(monkeypatch, search_info):
+    created_indexes = []
+    updated_indexes = []
+
+    async def mock_create_index(self, index):
+        created_indexes.append(index)
 
     async def mock_list_index_names(self):
         yield "test"
 
+    async def mock_get_index(self, *args, **kwargs):
+        return SearchIndex(
+            name="test",
+            fields=[],
+        )
+
+    async def mock_create_or_update_index(self, index, *args, **kwargs):
+        updated_indexes.append(index)
+
     monkeypatch.setattr(SearchIndexClient, "create_index", mock_create_index)
     monkeypatch.setattr(SearchIndexClient, "list_index_names", mock_list_index_names)
+    monkeypatch.setattr(SearchIndexClient, "get_index", mock_get_index)
+    monkeypatch.setattr(SearchIndexClient, "create_or_update_index", mock_create_or_update_index)
 
     manager = SearchManager(search_info)
     await manager.create_index()
-    assert len(indexes) == 0, "It should not have created a new index"
+    assert len(created_indexes) == 0, "It should not have created a new index"
+    assert len(updated_indexes) == 1, "It should have updated the existing index"
+    assert len(updated_indexes[0].fields) == 1
+    assert updated_indexes[0].fields[0].name == "storageUrl"
 
 
 @pytest.mark.asyncio