fix for citations (#24)

sophia-ramsey · web-flow · commit e11f086cf6a4 · 2025-03-06T08:47:32.000-08:00
* fix for citations

* fix citations with files env variable

* address comments
diff --git a/src/api/main.py b/src/api/main.py
@@ -5,6 +5,7 @@
 import logging
 import os
 import sys
+import json
 from typing import Dict
 
 from azure.ai.projects.aio import AIProjectClient
@@ -56,10 +57,7 @@
 
 @contextlib.asynccontextmanager
 async def lifespan(app: fastapi.FastAPI):
-    files: Dict[str, Dict[str, str]] = {}  # File name -> {"id": file_id, "path": file_path}
-    vector_store = None
     agent = None
-    create_new_agent = True
 
     try:
         if not os.getenv("RUNNING_IN_PRODUCTION"):
@@ -89,33 +87,35 @@ async def lifespan(app: fastapi.FastAPI):
         if os.environ.get("AZURE_AI_AGENT_ID") is not None:
             try: 
                 agent = await ai_client.agents.get_agent(os.environ["AZURE_AI_AGENT_ID"])
-                create_new_agent = False
                 logger.info("Agent already exists, skipping creation")
                 logger.info(f"Fetched agent, agent ID: {agent.id}")
                 logger.info(f"Fetched agent, model name: {agent.model}")
             except Exception as e:
                 logger.error(f"Error fetching agent: {e}", exc_info=True)
-                create_new_agent = True
-        if create_new_agent:
-            # Check if a previous agent created by the template exists
+
+        if not agent:
+            # Fallback to searching by name
+            agent_name = os.environ["AZURE_AI_AGENT_NAME"]
             agent_list = await ai_client.agents.list_agents()
             if agent_list.data:
                 for agent_object in agent_list.data:
-                    if agent_object.name == os.environ["AZURE_AI_AGENT_NAME"]:
+                    if agent_object.name == agent_name:
                         agent = agent_object
-        if agent == None:
-            raise Exception("Agent not found")
+                        logger.info(f"Found agent by name '{agent_name}', ID={agent_object.id}")
+                        break
 
-    except Exception as e:
-        logger.error(f"Error creating agent: {e}", exc_info=True)
-        raise RuntimeError(f"Failed to create the agent: {e}")
+        if not agent:
+            raise RuntimeError("No agent found. Ensure qunicorn.py created one or set AZURE_AI_AGENT_ID.")
 
-    app.state.ai_client = ai_client
-    app.state.agent = agent
-    app.state.files = files
+        app.state.ai_client = ai_client
+        app.state.agent = agent
 
-    try:
         yield
+
+    except Exception as e:
+        logger.error(f"Error during startup: {e}", exc_info=True)
+        raise RuntimeError(f"Error during startup: {e}")
+
     finally:
         try:
             await ai_client.close()
diff --git a/src/api/routes.py b/src/api/routes.py
@@ -197,7 +197,15 @@ async def fetch_document(request: Request):
     if not file_name:
         raise HTTPException(status_code=400, detail="file_name is required")
 
-    files = getattr(request.app.state, "files", {})
+    # 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)
+
     logger.info(f"File requested: {file_name}. Current file keys: {list(files.keys())}")
 
     if file_name not in files:
diff --git a/src/gunicorn.conf.py b/src/gunicorn.conf.py
@@ -1,11 +1,12 @@
 import multiprocessing
 import os
 import sys
+import json
 from typing import Dict
 import asyncio
 import logging
 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, Agent
 from azure.identity import DefaultAzureCredential
 
 from dotenv import load_dotenv
@@ -32,66 +33,122 @@
     file_handler.setFormatter(file_formatter)
     logger.addHandler(file_handler)
 
-async def list_or_create_agent():
-    files: Dict[str, Dict[str, str]] = {}  # File name -> {"id": file_id, "path": file_path}
-    vector_store = None
-    agent = None
+FILES_NAMES = ["product_info_1.md", "product_info_2.md"]
 
+async def create_agent(ai_client: AIProjectClient) -> Agent:
+    files: Dict[str, Dict[str, str]] = {}
+
+    # Create a new agent with the required resources
+    logger.info("Creating new agent with resources")
+
+    # Upload files for file search
+    for file_name in FILES_NAMES:
+        file_path = os.path.abspath(os.path.join(os.path.dirname(__file__), 'files', file_name))
+        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}
+
+    # 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']}")
+
+    # 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()],
+        name="sample_store"
+    )
+    logger.info("agent: file store and vector store success")
+
+    file_search_tool = FileSearchTool(vector_store_ids=[vector_store.id])
+    toolset = AsyncToolSet()
+    toolset.add(file_search_tool)
+
+    agent = await ai_client.agents.create_agent(
+        model=os.environ["AZURE_AI_AGENT_DEPLOYMENT_NAME"],
+        name=os.environ["AZURE_AI_AGENT_NAME"], 
+        instructions="You are helpful assistant",
+        toolset=toolset
+    )
+    return agent
+
+
+async def update_agent(agent: Agent, ai_client: AIProjectClient) -> Agent:
+    logger.info("Updating agent with resources")
+    files: Dict[str, Dict[str, str]] = {}
+
+    # Upload files for file search
+    for file_name in FILES_NAMES:
+        file_path = os.path.abspath(os.path.join(os.path.dirname(__file__), 'files', file_name))
+        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}
+
+    # 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']}")
+
+    # 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()],
+        name="sample_store"
+    )
+    logger.info("agent: file store and vector store success")
+
+    file_search_tool = FileSearchTool(vector_store_ids=[vector_store.id])
+    toolset = AsyncToolSet()
+    toolset.add(file_search_tool)
+
+    agent = await ai_client.agents.update_agent(
+        assistant_id=agent.id,
+        model=os.environ["AZURE_AI_AGENT_DEPLOYMENT_NAME"],
+        name=os.environ["AZURE_AI_AGENT_NAME"],
+        instructions="You are helpful assistant",
+        toolset=toolset
+    )
+    return agent
+
+
+async def initialize_resources():
     try:
         ai_client = AIProjectClient.from_connection_string(
             credential=DefaultAzureCredential(exclude_shared_token_cache_credential=True),
             conn_str=os.environ["AZURE_AIPROJECT_CONNECTION_STRING"],
         )
 
-        if os.environ.get("AZURE_AI_AGENT_ID"):
+        # If the environment already has AZURE_AI_AGENT_ID, try fetching that agent
+        if os.environ.get("AZURE_AI_AGENT_ID") is not None:
             try: 
                 agent = await ai_client.agents.get_agent(os.environ["AZURE_AI_AGENT_ID"])
+                logger.info(f"Found agent by ID: {agent.id}")
+                # Update the agent with the latest resources
+                agent = await update_agent(agent, ai_client)
                 return
             except Exception as e:
-                logger.info("Error with agent ID")
+                logger.warning(f"Could not retrieve agent by AZURE_AI_AGENT_ID = {os.environ['AZURE_AI_AGENT_ID']}, error: {e}")
 
-        # Check if a previous agent created by the template exists
+        # Check if an agent with the same name already exists
         agent_list = await ai_client.agents.list_agents()
         if agent_list.data:
             for agent_object in agent_list.data:
                 if agent_object.name == os.environ["AZURE_AI_AGENT_NAME"]:
-                    return 
-        
-        # Create a new agent with the required resources
-        logger.info("Creating new agent with resources")
-        file_names = ["product_info_1.md", "product_info_2.md"]
-        for file_name in file_names:
-            file_path = os.path.abspath(os.path.join(os.path.dirname(__file__), 'files', file_name))
-            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}
-        
-        # 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()],
-            name="sample_store"
-        )
-        logger.info("agent: file store and vector store success")
+                    logger.info(f"Found existing agent named '{agent_object.name}', ID: {agent_object.id}")
+                    os.environ["AZURE_AI_AGENT_ID"] = agent_object.id
+                    # Update the agent with the latest resources
+                    agent = await update_agent(agent_object, ai_client)
+                    return
 
-        file_search_tool = FileSearchTool(vector_store_ids=[vector_store.id])
-        toolset = AsyncToolSet()
-        toolset.add(file_search_tool)
+        # Create a new agent
+        agent = await create_agent(ai_client)
+        os.environ["AZURE_AI_AGENT_ID"] = agent.id
+        logger.info(f"Created agent, agent ID: {agent.id}")
 
-        agent = await ai_client.agents.create_agent(
-            model=os.environ["AZURE_AI_AGENT_DEPLOYMENT_NAME"],
-            name=os.environ["AZURE_AI_AGENT_NAME"], 
-            instructions="You are helpful assistant",
-            toolset=toolset
-        )
-        logger.info("Created agent, agent ID: {agent.id}")
-    
     except Exception as e:
         logger.info("Error creating agent: {e}", exc_info=True)
         raise RuntimeError(f"Failed to create the agent: {e}")
-    
+
 def on_starting(server):
     """This code runs once before the workers will start."""
-    asyncio.get_event_loop().run_until_complete(list_or_create_agent())
+    asyncio.get_event_loop().run_until_complete(initialize_resources())
 
 max_requests = 1000
 max_requests_jitter = 50
