Switch vectorestore to azure search

app.py

@@ -7,7 +7,7 @@ from azure.storage.fileshare import ShareServiceClient
 # Import custom modules
 from climateqa.engine.embeddings import get_embeddings_function
 from climateqa.engine.llm import get_llm
-from climateqa.engine.vectorstore import get_pinecone_vectorstore
+from climateqa.engine.vectorstore import get_vectorstore
 from climateqa.engine.reranker import get_reranker
 from climateqa.engine.graph import make_graph_agent, make_graph_agent_poc
 from climateqa.engine.chains.retrieve_papers import find_papers
@@ -66,17 +66,11 @@ user_id = create_user_id()
 
 # Create vectorstore and retriever
 embeddings_function = get_embeddings_function()
-vectorstore = get_pinecone_vectorstore(
-    embeddings_function, index_name=os.getenv("PINECONE_API_INDEX")
-)
-vectorstore_graphs = get_pinecone_vectorstore(
-    embeddings_function,
-    index_name=os.getenv("PINECONE_API_INDEX_OWID"),
-    text_key="description",
-)
-vectorstore_region = get_pinecone_vectorstore(
-    embeddings_function, index_name=os.getenv("PINECONE_API_INDEX_LOCAL_V2")
-)
+
+vectorstore = get_vectorstore(provider="azure_search", embeddings=embeddings_function, index_name="climateqa-ipx")
+vectorstore_graphs = get_vectorstore(provider="azure_search", embeddings=embeddings_function, index_name="climateqa-owid", text_key="description")
+vectorstore_region = get_vectorstore(provider="azure_search", embeddings=embeddings_function, index_name="climateqa-v2")
+
 
 llm = get_llm(provider="openai", max_tokens=1024, temperature=0.0)
 if os.environ["GRADIO_ENV"] == "local":

climateqa/engine/chains/retrieve_documents.py

@@ -19,7 +19,7 @@ from ..llm import get_llm
 from .prompts import retrieve_chapter_prompt_template
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
-from ..vectorstore import get_pinecone_vectorstore
+from ..vectorstore import get_vectorstore
 from ..embeddings import get_embeddings_function
 import ast
 
@@ -134,7 +134,7 @@ def get_ToCs(version: str) :
         "version": version
     }
     embeddings_function = get_embeddings_function()
-    vectorstore = get_pinecone_vectorstore(embeddings_function, index_name="climateqa-v2")
+    vectorstore = get_vectorstore(provider="qdrant", embeddings=embeddings_function, index_name="climateqa")
     tocs = vectorstore.similarity_search_with_score(query="",filter = filters_text)
 
     # remove duplicates or almost duplicates
@@ -236,7 +236,7 @@ async def get_POC_documents_by_ToC_relevant_documents(
     filters_text_toc = {
         **filters,
         "chunk_type":"text",
-        "toc_level0": {"$in": toc_filters},
+        "toc_level0": toc_filters,  # Changed from {"$in": toc_filters} to direct list
         "version": version
         # "report_type": {}, # TODO  to be completed to choose the right documents / chapters according to the analysis of the question
     }
@@ -273,6 +273,22 @@ async def get_POC_documents_by_ToC_relevant_documents(
         "docs_images" : docs_images
     }
     
+def filter_for_full_report_documents(filters: dict) -> dict:
+    """
+    Filter for full report documents.
+    Returns a dictionary format compatible with all vectorstore providers.
+    """
+    # Start with the base filters
+    full_filters = filters.copy()
+    
+    # Add chunk_type filter
+    full_filters["chunk_type"] = "text"
+    
+    # Add report_type exclusion using the new _exclude suffix format
+    # This will be converted to appropriate OData filter by Azure Search wrapper
+    full_filters["report_type_exclude"] = ["SPM"]
+    
+    return full_filters
 
 async def get_IPCC_relevant_documents(
     query: str,
@@ -299,9 +315,9 @@ async def get_IPCC_relevant_documents(
     filters = {}
 
     if len(reports) > 0:
-        filters["short_name"] = {"$in":reports}
+        filters["short_name"] = reports  # Changed from {"$in":reports} to direct list
     else:
-        filters["source"] = { "$in": sources}
+        filters["source"] = sources  # Changed from {"$in": sources} to direct list
 
     # INIT 
     docs_summaries = []
@@ -323,18 +339,16 @@ async def get_IPCC_relevant_documents(
         filters_summaries = {
             **filters,
             "chunk_type":"text",
-            "report_type": { "$in":["SPM"]},
+            "report_type": ["SPM"],  # Changed from {"$in":["SPM"]} to direct list
         }
 
         docs_summaries = vectorstore.similarity_search_with_score(query=query,filter = filters_summaries,k = k_summary)
         docs_summaries = [x for x in docs_summaries if x[1] > threshold]
 
         # Search for k_total - k_summary documents in the full reports dataset
-        filters_full = {
-            **filters,
-            "chunk_type":"text",
-            "report_type": { "$nin":["SPM"]},
-        }
+        filters_full = filter_for_full_report_documents(filters)
+        
+        
         docs_full = vectorstore.similarity_search_with_score(query=query,filter = filters_full,k = k_total)
         
         if search_figures:

climateqa/engine/graph_retriever.py

@@ -60,10 +60,9 @@ async def retrieve_graphs(
         assert sources
         assert any([x in ["OWID"] for x in sources])
 
-        # Prepare base search kwargs
-        filters = {}
-
-        filters["source"] = {"$in": sources}
+        # Prepare base search kwargs for Azure AI Search
+        # Azure expects a filter string, e.g. "source eq 'OWID' or source eq 'IEA'"
+        filters = {"source":"OWID"}
 
         docs = vectorstore.similarity_search_with_score(query=query, filter=filters, k=k_total)
         

climateqa/engine/llm/openai.py

@@ -8,7 +8,6 @@ except Exception:
     pass
 
 def get_llm(model="gpt-4o-mini",max_tokens=1024, temperature=0.0, streaming=True,timeout=30, **kwargs):
-    
     llm = ChatOpenAI(
         model=model,
         api_key=os.environ.get("THEO_API_KEY", None),

climateqa/engine/vectorstore.py

@@ -1,11 +1,11 @@
-# Pinecone
-# More info at https://docs.pinecone.io/docs/langchain
-# And https://python.langchain.com/docs/integrations/vectorstores/pinecone
+# Azure AI Search: https://python.langchain.com/docs/integrations/vectorstores/azuresearch
 import os
-from pinecone import Pinecone
-from langchain_community.vectorstores import Pinecone as PineconeVectorstore
 
-# LOAD ENVIRONMENT VARIABLES
+# Azure AI Search imports
+from langchain_community.vectorstores.azuresearch import AzureSearch
+
+
+# Load environment variables
 try:
     from dotenv import load_dotenv
     load_dotenv()
@@ -13,44 +13,136 @@ except:
     pass
 
 
-
-
-def get_pinecone_vectorstore(embeddings,text_key = "content", index_name = os.getenv("PINECONE_API_INDEX")):
-
-    # # initialize pinecone
-    # pinecone.init(
-    #     api_key=os.getenv("PINECONE_API_KEY"),  # find at app.pinecone.io
-    #     environment=os.getenv("PINECONE_API_ENVIRONMENT"),  # next to api key in console
-    # )
-
-    # index_name = os.getenv("PINECONE_API_INDEX")
-    # vectorstore = Pinecone.from_existing_index(index_name, embeddings,text_key = text_key)
-
-    # return vectorstore
-
-    pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
-    index = pc.Index(index_name)
-
-    vectorstore = PineconeVectorstore(
-        index, embeddings, text_key,
+class AzureSearchWrapper:
+    """
+    Wrapper class for Azure AI Search vectorstore to handle filter conversion.
+    
+    This wrapper automatically converts dictionary-style filters to Azure Search OData filter format,
+    ensuring seamless compatibility when switching from other providers.
+    """
+    
+    def __init__(self, azure_search_vectorstore):
+        self.vectorstore = azure_search_vectorstore
+        
+    def __getattr__(self, name):
+        """Delegate all other attributes to the wrapped vectorstore."""
+        return getattr(self.vectorstore, name)
+    
+    def _convert_dict_filter_to_odata(self, filter_dict):
+        """
+        Convert dictionary-style filters to Azure Search OData filter format.
+        
+        Args:
+            filter_dict (dict): Dictionary-style filter
+            
+        Returns:
+            str: OData filter string
+        """
+        if not filter_dict:
+            return None
+            
+        conditions = []
+        
+        for key, value in filter_dict.items():
+            if key.endswith('_exclude'):
+                # Handle exclusion filters (e.g., report_type_exclude)
+                base_key = key.replace('_exclude', '')
+                if isinstance(value, list):
+                    if len(value) == 1:
+                        conditions.append(f"{base_key} ne '{value[0]}'")
+                    else:
+                        exclude_conditions = [f"{base_key} ne '{v}'" for v in value]
+                        conditions.append(f"({' and '.join(exclude_conditions)})")
+                else:
+                    conditions.append(f"{base_key} ne '{value}'")
+            elif isinstance(value, list):
+                # Handle list values (equivalent to $in operator)
+                if len(value) == 1:
+                    conditions.append(f"{key} eq '{value[0]}'")
+                else:
+                    list_conditions = [f"{key} eq '{v}'" for v in value]
+                    conditions.append(f"({' or '.join(list_conditions)})")
+            else:
+                # Handle single values
+                conditions.append(f"{key} eq '{value}'")
+        
+        return " and ".join(conditions) if conditions else None
+    
+    def similarity_search_with_score(self, query, k=4, filter=None, **kwargs):
+        """Override similarity_search_with_score to convert filters."""
+        if filter is not None:
+            filter = self._convert_dict_filter_to_odata(filter)
+
+        return self.vectorstore.hybrid_search_with_score(
+            query=query, k=k, filters=filter, **kwargs
+        )
+
+    
+    def similarity_search(self, query, k=4, filter=None, **kwargs):
+        """Override similarity_search to convert filters."""
+        if filter is not None:
+            filter = self._convert_dict_filter_to_odata(filter)
+        
+        return self.vectorstore.similarity_search(
+            query=query, k=k, filter=filter, **kwargs
+        )
+    
+    def similarity_search_by_vector(self, embedding, k=4, filter=None, **kwargs):
+        """Override similarity_search_by_vector to convert filters."""
+        if filter is not None:
+            filter = self._convert_dict_filter_to_odata(filter)
+        
+        return self.vectorstore.similarity_search_by_vector(
+            embedding=embedding, k=k, filter=filter, **kwargs
+        )
+    
+    def as_retriever(self, search_type="similarity", search_kwargs=None, **kwargs):
+        """Override as_retriever to handle filter conversion in search_kwargs."""
+        if search_kwargs and "filter" in search_kwargs:
+            # Convert the filter in search_kwargs
+            search_kwargs = search_kwargs.copy()  # Don't modify the original
+            if search_kwargs["filter"] is not None:
+                search_kwargs["filter"] = self._convert_dict_filter_to_odata(search_kwargs["filter"])
+        
+        return self.vectorstore.as_retriever(
+            search_type=search_type, search_kwargs=search_kwargs, **kwargs
+        )
+
+
+def get_azure_search_vectorstore(embeddings, text_key="content", index_name=None):
+    """
+    Create an Azure AI Search vectorstore instance.
+    
+    Args:
+        embeddings: The embeddings function to use
+        text_key: The key for text content in the payload (default: "content")
+        index_name: The name of the Azure Search index
+    
+    Returns:
+        AzureSearchWrapper: A wrapped Azure AI Search vectorstore instance with filter compatibility
+    """
+    # Get Azure AI Search configuration from environment variables
+    azure_search_endpoint = os.getenv("AI_SEARCH_INDEX_ENDPOINT")
+    azure_search_key = os.getenv("AI_SEARCH_KEY")
+    
+    if not azure_search_endpoint:
+        raise ValueError("AI_SEARCH_INDEX_ENDPOINT environment variable is required")
+    
+    if not azure_search_key:
+        raise ValueError("AI_SEARCH_KEY environment variable is required")
+    
+    if not index_name:
+        raise ValueError("index_name must be provided for Azure Search")
+    
+    # Create Azure Search vectorstore
+    vectorstore = AzureSearch(
+        azure_search_endpoint=azure_search_endpoint,
+        azure_search_key=azure_search_key,
+        index_name=index_name,
+        embedding_function=embeddings.embed_query,
+        content_key=text_key,
     )
-    return vectorstore
-
-
-
-# def get_pinecone_retriever(vectorstore,k = 10,namespace = "vectors",sources = ["IPBES","IPCC"]):
-
-#     assert isinstance(sources,list)
-
-#     # Check if all elements in the list are either IPCC or IPBES
-#     filter = {
-#         "source": { "$in":sources},
-#     }
-
-#     retriever = vectorstore.as_retriever(search_kwargs={
-#         "k": k,
-#         "namespace":"vectors",
-#         "filter":filter
-#     })
+    
+    # Wrap the vectorstore to handle filter conversion
+    return AzureSearchWrapper(vectorstore)
 
-#     return retriever

climateqa/utils.py

@@ -25,7 +25,7 @@ def remove_duplicates_keep_highest_score(documents):
     unique_docs = {}
     
     for doc in documents:
-        doc_id = doc.metadata.get('doc_id')
+        doc_id = doc.metadata.get('id')
         if doc_id in unique_docs:
             if doc.metadata['reranking_score'] > unique_docs[doc_id].metadata['reranking_score']:
                 unique_docs[doc_id] = doc

requirements.txt

@@ -1,6 +1,9 @@
 gradio==5.0.2
 azure-storage-file-share==12.11.1
 azure-storage-blob==12.23.0
+# Azure AI Search support
+azure-search-documents>=11.4.0
+azure-core>=1.29.0
 python-dotenv==1.0.0
 langchain==0.2.1
 langchain_openai==0.1.7