""" title: Azure AI Search author: secondhandrebel funding_url: https://github.com/open-webui version: 0.1.0 license: MIT """ from pydantic import BaseModel, Field from azure.identity import DefaultAzureCredential from azure.search.documents import SearchClient from azure.search.documents.indexes import SearchIndexClient from azure.search.documents.indexes.models import SearchIndex from azure.search.documents.models import VectorizableTextQuery class Tools: class Valves(BaseModel): SEARCH_INDEX: str = Field(default="", description="Azure AI Search Index") SEARCH_SERVICE: str = Field( default="", description="Azure AI Search Service Name" ) RESULTS_NO: str = Field(default="4", description="Number of results to return") # Add your custom tools using pure Python code here, make sure to add type hints # Use Sphinx-style docstrings to document your tools, they will be used for generating tools specifications # Please refer to function_calling_filter_pipeline.py file from pipelines project for an example def __init__(self): print(f"[AI Search] initializing search") self.valves = self.Valves() self.user_valves = None # Add a description of contents of the knowledgebase def azure_ai_search(self, search_query: str) -> str: """ This search tool will search an internal knowledgebase. You must formulate your own search query based on the user's message. :param search_query: A semantic search query used in search engine. :return: The search results in a text block. """ # Do not include :param for __user__ in the docstring as it should not be shown in the tool's specification # The session user object will be passed as a parameter when the function is called print("[AI Search] Executing AI Search") print(f"[AI Search] search_query: {search_query}") azd_credential = DefaultAzureCredential() search_client = SearchClient( endpoint=f"https://{self.valves.SEARCH_SERVICE}.search.windows.net/", index_name=self.valves.SEARCH_INDEX, credential=azd_credential, ) vector_query = VectorizableTextQuery( text=search_query, kind="text", k_nearest_neighbors=self.valves.RESULTS_NO, fields="embeddings", ) search_results = search_client.search( search_text=search_query, top=self.valves.RESULTS_NO, include_total_count=True, query_type="semantic", semantic_configuration_name="default", vector_queries=[vector_query], ) print(f"[AI Search] search complete.") results = "" try: for r in search_results: results = results + f'source:{r["sourcefile"]}\n' results = results + f'url: {r["weburl"]}\n' results = results + f'content: {r["content"]}\n' except Exception as e: print(e) return results