Fix

Author: nick863

src/api/search_index_manager.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
 import csv
 import glob
@@ -7,7 +7,7 @@
 import time
 
 from azure.core.credentials_async import AsyncTokenCredential
-from azure.search.documents.aio import SearchClient
+from azure.search.documents.aio import AsyncSearchItemPaged, SearchClient 
 from azure.search.documents.indexes.aio import SearchIndexClient
 from azure.core.exceptions import HttpResponseError
 from azure.search.documents.indexes.models import (
@@ -50,6 +50,11 @@ class SearchIndexManager:
     MIN_DIFF_CHARACTERS_IN_LINE = 5
     MIN_LINE_LENGTH = 5
 
+    _SEMANTIC_CONFIG = "semantic_search"
+    _EMBEDDING_CONFIG = "embedding_config"
+    _VECTORIZER = "search_vectorizer"
+
+
     def __init__(
             self,
             endpoint: str,
@@ -141,6 +146,32 @@ def _check_dimensions(self, vector_index_dimensions: Optional[int] = None) -> in
             raise ValueError("vector_index_dimensions is different from dimensions provided to constructor.")
         return vector_index_dimensions
 
+    async def _format_search_results(self, response: AsyncSearchItemPaged[Dict]) -> str:
+        """
+        Format the output of search.
+
+        :param response: The search results.
+        :return: The formatted response string.
+        """
+        results = [f"{result['token']}, source: {result['document_reference']}" async for result in response]
+        return "\n------\n".join(results)
+
+    async def semantic_search(self, message: str) -> str:
+        """
+        Perform the semantic search on the search resource.
+
+        :param message: The customer question.
+        :return: The context for the question.
+        """
+        self._raise_if_no_index()
+        response = await self._get_client().search(
+            search_text=message,
+            query_type="semantic",
+            semantic_configuration_name=SearchIndexManager._SEMANTIC_CONFIG,
+        )
+        return await self._format_search_results(response)
+        
+
     async def search(self, message: str) -> str:
         """
         Search the message in the vector store.
@@ -160,8 +191,7 @@ async def search(self, message: str) -> str:
         )
         # This lag is necessary, despite it is not described in documentation.
         time.sleep(1)
-        results = [f"{result['token']}, source: {result['document_reference']}" async for result in response]
-        return "\n------\n".join(results)
+        return await self._format_search_results(response)
 
     async def create_index(
         self,
@@ -185,7 +215,7 @@ async def create_index(
         """
         vector_index_dimensions = self._check_dimensions(vector_index_dimensions)
         try:
-            self._index = await self._index_create()
+            self._index = await self._index_create(vector_index_dimensions)
             return True
         except HttpResponseError:
             if raise_on_error:
@@ -194,33 +224,44 @@ async def create_index(
                 self._index = await ix_client.get_index(self._index_name)
             return False
 
-    async def _index_create(self) -> SearchIndex:
-        """Create the index."""
+    async def _index_create(self, vector_index_dimensions: int) -> SearchIndex:
+        """
+        Create the index.
+
+        :param vector_index_dimensions: The number of dimensions in the vector index. This parameter is
+               needed if the embedding parameter cannot be set for the given model. It can be
+               figured out by loading the embeddings file, generated by build_embeddings_file,
+               loading the contents of the first row and 'embedding' column as a JSON and calculating
+               the length of the list obtained.
+               Also please see the embedding model documentation
+               https://platform.openai.com/docs/models#embeddings
+        :return: The newly created search index.
+        """
         async with SearchIndexClient(endpoint=self._endpoint, credential=self._credential) as ix_client:
             fields = [
                 SimpleField(name="embedId", type=SearchFieldDataType.String, key=True),
                 SearchField(
                     name="embedding",
                     type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
-                    vector_search_dimensions=self._dimensions,
+                    vector_search_dimensions=vector_index_dimensions,
                     searchable=True,
-                    vector_search_profile_name="embedding_config"
+                    vector_search_profile_name=SearchIndexManager._EMBEDDING_CONFIG
                 ),
                 SearchField(name="token", searchable=True, type=SearchFieldDataType.String, hidden=False),
                 SearchField(name="document_reference", type=SearchFieldDataType.String, hidden=False),
             ]
             vector_search = VectorSearch(
                 profiles=[
                     VectorSearchProfile(
-                        name="embedding_config",
+                        name=SearchIndexManager._EMBEDDING_CONFIG,
                         algorithm_configuration_name="embed-algorithms-config",
-                        vectorizer_name="search_vectorizer"
+                        vectorizer_name=SearchIndexManager._VECTORIZER
                     )
                 ],
                 algorithms=[HnswAlgorithmConfiguration(name="embed-algorithms-config")],
                 vectorizers=[
                     AzureOpenAIVectorizer(
-                        vectorizer_name="search_vectorizer",
+                        vectorizer_name=SearchIndexManager._VECTORIZER,
                         parameters=AzureOpenAIVectorizerParameters(
                             resource_url=self._embeddings_endpoint,
                             deployment_name=self._embedding_deployment,
@@ -231,10 +272,10 @@ async def _index_create(self) -> SearchIndex:
                 ]
             )
             semantic_search = SemanticSearch(
-                default_configuration_name="index_search",
+                default_configuration_name=SearchIndexManager._SEMANTIC_CONFIG,
                 configurations=[
                     SemanticConfiguration(
-                        name="index_search",
+                        name=SearchIndexManager._SEMANTIC_CONFIG,
                         prioritized_fields=SemanticPrioritizedFields(
                             title_field=SemanticField(field_name="embedId"),
                             content_fields=[SemanticField(field_name="token")]

src/gunicorn.conf.py

@@ -4,6 +4,7 @@
 from typing import Dict
 
 import asyncio
+import csv
 import json
 import logging
 import multiprocessing
@@ -101,6 +102,18 @@ async def create_index_maybe(
             await search_mgr.close()
 
 
+def _get_file_path(file_name: str) -> str:
+    """
+    Get absolute file path.
+
+    :param file_name: The file name.
+    """
+    return os.path.abspath(
+        os.path.join(os.path.dirname(__file__),
+                     'files',
+                     file_name))
+
+
 async def get_available_toolset(
         ai_client: AIProjectClient,
         creds: AsyncTokenCredential) -> AsyncToolSet:
@@ -111,10 +124,12 @@ async def get_available_toolset(
     :param creds: The credentials, used for the index.
     :return: The tool set, available based on the environment.
     """
+    # File name -> {"id": file_id, "path": file_path}
+    files: Dict[str, Dict[str, str]] = {}
     # First try to get an index search.
     conn_id = ""
     if os.environ.get('AZURE_AI_SEARCH_INDEX_NAME'):
-        conn_list = ai_client.connections.list()
+        conn_list = await ai_client.connections.list()
         for conn in conn_list:
             if conn.connection_type == ConnectionType.AZURE_AI_SEARCH:
                 conn_id = conn.id
@@ -130,26 +145,28 @@ async def get_available_toolset(
 
         toolset.add(ai_search)
         logger.info("agent: initialized index")
+        # Populate file links.
+        embeddings_path = os.path.join(
+            os.path.dirname(__file__), 'data', 'embeddings.csv')
+        with open(embeddings_path, newline='') as fp:
+            reader = csv.DictReader(fp)
+            for row in reader:
+                if row['document_reference'] in FILES_NAMES:
+                    files[row['document_reference']] = {
+                        "id": row['document_reference'],
+                        "path": _get_file_path(row['document_reference'])
+                    }
     else:
         logger.info(
             "agent: index was not initialized, falling back to file search.")
         # Upload files for file search
-        # File name -> {"id": file_id, "path": file_path}
-        files: Dict[str, Dict[str, str]] = {}
         for file_name in FILES_NAMES:
-            file_path = os.path.abspath(
-                os.path.join(
-                    os.path.dirname(__file__),
-                    'files',
-                    file_name))
+            file_path = _get_file_path(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']}")
 
@@ -162,6 +179,9 @@ async def get_available_toolset(
 
         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)
     return toolset
 
 
@@ -173,7 +193,7 @@ async def create_agent(ai_client: AIProjectClient,
         model=os.environ["AZURE_AI_AGENT_DEPLOYMENT_NAME"],
         name=os.environ["AZURE_AI_AGENT_NAME"],
         instructions="You are helpful assistant",
-        toolset=await get_available_toolset()
+        toolset=await get_available_toolset(ai_client, creds)
     )
     return agent
 
@@ -209,7 +229,7 @@ async def initialize_resources():
                             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)
+                        agent = await update_agent(agent, ai_client, creds)
                         return
                     except Exception as e:
                         logger.warning(
@@ -221,18 +241,19 @@ async def initialize_resources():
                 if agent_list.data:
                     for agent_object in agent_list.data:
                         if agent_object.name == os.environ[
-                          "AZURE_AI_AGENT_NAME"]:
+                                "AZURE_AI_AGENT_NAME"]:
                             logger.info(
                                 "Found existing agent named "
                                 f"'{agent_object.name}'"
                                 f", 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)
+                            agent = await update_agent(
+                                agent_object, ai_client, creds)
                             return
 
                 # Create a new agent
-                agent = await create_agent(ai_client)
+                agent = await create_agent(ai_client, creds)
                 os.environ["AZURE_AI_AGENT_ID"] = agent.id
                 logger.info(f"Created agent, agent ID: {agent.id}")