File size: 4,171 Bytes
cfb0fa4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import logging
from typing import Optional
from open_webui.retrieval.web.main import SearchResult, get_filtered_results

log = logging.getLogger(__name__)

"""
Azure AI Search integration for Open WebUI.
Documentation: https://learn.microsoft.com/en-us/python/api/overview/azure/search-documents-readme?view=azure-python

Required package: azure-search-documents
Install: pip install azure-search-documents
"""


def search_azure(
    api_key: str,
    endpoint: str,
    index_name: str,
    query: str,
    count: int,
    filter_list: Optional[list[str]] = None,
) -> list[SearchResult]:
    """
    Search using Azure AI Search.

    Args:
        api_key: Azure Search API key (query key or admin key)
        endpoint: Azure Search service endpoint (e.g., https://myservice.search.windows.net)
        index_name: Name of the search index to query
        query: Search query string
        count: Number of results to return
        filter_list: Optional list of domains to filter results

    Returns:
        List of SearchResult objects with link, title, and snippet
    """
    try:
        from azure.core.credentials import AzureKeyCredential
        from azure.search.documents import SearchClient
    except ImportError:
        log.error(
            "azure-search-documents package is not installed. "
            "Install it with: pip install azure-search-documents"
        )
        raise ImportError(
            "azure-search-documents is required for Azure AI Search. "
            "Install it with: pip install azure-search-documents"
        )

    try:
        # Create search client with API key authentication
        credential = AzureKeyCredential(api_key)
        search_client = SearchClient(
            endpoint=endpoint, index_name=index_name, credential=credential
        )

        # Perform the search
        results = search_client.search(search_text=query, top=count)

        # Convert results to list and extract fields
        search_results = []
        for result in results:
            # Azure AI Search returns documents with custom schemas
            # We need to extract common fields that might represent URL, title, and content
            # Common field names to look for:
            result_dict = dict(result)

            # Try to find URL field (common names)
            link = (
                result_dict.get("url")
                or result_dict.get("link")
                or result_dict.get("uri")
                or result_dict.get("metadata_storage_path")
                or ""
            )

            # Try to find title field (common names)
            title = (
                result_dict.get("title")
                or result_dict.get("name")
                or result_dict.get("metadata_title")
                or result_dict.get("metadata_storage_name")
                or None
            )

            # Try to find content/snippet field (common names)
            snippet = (
                result_dict.get("content")
                or result_dict.get("snippet")
                or result_dict.get("description")
                or result_dict.get("summary")
                or result_dict.get("text")
                or None
            )

            # Truncate snippet if too long
            if snippet and len(snippet) > 500:
                snippet = snippet[:497] + "..."

            if link:  # Only add if we found a valid link
                search_results.append(
                    {
                        "link": link,
                        "title": title,
                        "snippet": snippet,
                    }
                )

        # Apply domain filtering if specified
        if filter_list:
            search_results = get_filtered_results(search_results, filter_list)

        # Convert to SearchResult objects
        return [
            SearchResult(
                link=result["link"],
                title=result.get("title"),
                snippet=result.get("snippet"),
            )
            for result in search_results
        ]

    except Exception as ex:
        log.error(f"Azure AI Search error: {ex}")
        raise ex