Construct upload_file_map on runtime and avoid save as env var

Author: howieleung

src/api/main.py

@@ -9,7 +9,7 @@
 from typing import Dict
 
 from azure.ai.projects.aio import AIProjectClient
-from azure.ai.projects.models import FilePurpose, FileSearchTool, AsyncToolSet
+from azure.ai.projects.models import FilePurpose, FileSearchTool, AsyncToolSet, FileSearchToolResource
 from azure.identity import DefaultAzureCredential
 
 import fastapi
@@ -76,7 +76,28 @@ async def lifespan(app: fastapi.FastAPI):
 
         app.state.ai_client = ai_client
         app.state.agent = agent
-
+        
+        
+        file_map: Dict[str, Dict[str, str]] = {}
+
+        logger.info(f"Creating UPLOADED_FILE_MAP")
+        if agent.tool_resources and agent.tool_resources.file_search:
+            for vector_store_id in agent.tool_resources.file_search.vector_store_ids:
+                file = await ai_client.agents.list_vector_store_files(vector_store_id)
+                for file in file.data:
+                    openAI_file = await ai_client.agents.get_file(file.id)
+                    logger.info(f"Retrieved file, file ID: {openAI_file.filename}")
+                    file_name = openAI_file.filename
+                    file_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..\\files', file_name))
+
+                    if not os.path.exists(file_path):
+                        logger.warning(f"File path does not exist: {file_path}")
+                        continue 
+                                       
+                    # Store both file id and the file path using the file name as key.
+                    file_map[file_name] = {"id": file.id, "path": file_path}
+        app.state.upload_file_map = file_map
+        logger.info(f"Set UPLOADED_FILE_MAP {file_map}")                    
         yield
 
     except Exception as e:

src/api/routes.py

@@ -272,14 +272,7 @@ async def fetch_document(request: Request):
     if not file_name:
         raise HTTPException(status_code=400, detail="file_name is required")
 
-    # Reconstruct the file dictionary from the env variable:
-    files_env = os.environ['UPLOADED_FILE_MAP']
-    try:
-        files = json.loads(files_env)
-        logger.info("Successfully parsed UPLOADED_FILE_MAP from environment variable.")
-    except json.JSONDecodeError:
-        files = {}
-        logger.warning("Failed to parse UPLOADED_FILE_MAP from environment variable.", exc_info=True)
+    files = request.app.state.upload_file_map
 
     logger.info(f"File requested: {file_name}. Current file keys: {list(files.keys())}")
 

src/gunicorn.conf.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 # Licensed under the MIT license.
 # See LICENSE file in the project root for full license information.
-from typing import Dict
+from typing import Dict, List
 
 import asyncio
 import csv
@@ -110,7 +110,7 @@ async def get_available_toolset(
     :return: The tool set, available based on the environment.
     """
     # File name -> {"id": file_id, "path": file_path}
-    files: Dict[str, Dict[str, str]] = {}
+    file_ids: List[str] = []
     # First try to get an index search.
     conn_id = ""
     if os.environ.get('AZURE_AI_SEARCH_INDEX_NAME'):
@@ -129,11 +129,6 @@ async def get_available_toolset(
             index_name=os.environ.get('AZURE_AI_SEARCH_INDEX_NAME'))
 
         toolset.add(ai_search)
-        # Register the files
-        for file_name in FILES_NAMES:
-            file_path = _get_file_path(file_name)
-            files[file_name] = {"id": file_name, "path": file_path}
-        logger.info("agent: initialized index")
     else:
         logger.info(
             "agent: index was not initialized, falling back to file search.")
@@ -144,22 +139,17 @@ async def get_available_toolset(
             file = await ai_client.agents.upload_file_and_poll(
                 file_path=file_path, purpose=FilePurpose.AGENTS)
             # Store both file id and the file path using the file name as key.
-            files[file_name] = {"id": file.id, "path": file_path}
+            file_ids.append(file.id)
 
         # Create the vector store using the file IDs.
         vector_store = await ai_client.agents.create_vector_store_and_poll(
-            file_ids=[info["id"] for info in files.values()],
+            file_ids=file_ids,
             name="sample_store"
         )
         logger.info("agent: file store and vector store success")
 
         file_search_tool = FileSearchTool(vector_store_ids=[vector_store.id])
         toolset.add(file_search_tool)
-    # Serialize and store files information in the environment variable (so
-    # workers see it)
-    os.environ["UPLOADED_FILE_MAP"] = json.dumps(files)
-    logger.info(
-        f"Set env UPLOADED_FILE_MAP = {os.environ['UPLOADED_FILE_MAP']}")
 
     return toolset