Coverage for open_webui/retrieval/web/azure.py: 15%

1import logging 

2from typing import Optional 

3 

4from open_webui.retrieval.web.main import SearchResult, get_filtered_results 

5 

6log = logging.getLogger(__name__) 

7 

8""" 

9Azure AI Search integration for Open WebUI. 

10Documentation: https://learn.microsoft.com/en-us/python/api/overview/azure/search-documents-readme?view=azure-python 

11 

12Required package: azure-search-documents 

13Install: pip install azure-search-documents 

14""" 

15 

16 

17def search_azure( 

18 api_key: str, 

19 endpoint: str, 

20 index_name: str, 

21 query: str, 

22 count: int, 

23 filter_list: Optional[list[str]] = None, 

24) -> list[SearchResult]: 

25 """ 

26 Search using Azure AI Search. 

27 

28 Args: 

29 api_key: Azure Search API key (query key or admin key) 

30 endpoint: Azure Search service endpoint (e.g., https://myservice.search.windows.net) 

31 index_name: Name of the search index to query 

32 query: Search query string 

33 count: Number of results to return 

34 filter_list: Optional list of domains to filter results 

35 

36 Returns: 

37 List of SearchResult objects with link, title, and snippet 

38 """ 

39 try: 

40 from azure.core.credentials import AzureKeyCredential 

41 from azure.search.documents import SearchClient 

42 except ImportError: 

43 log.error( 

44 'azure-search-documents package is not installed. Install it with: pip install azure-search-documents' 

45 ) 

46 raise ImportError( 

47 'azure-search-documents is required for Azure AI Search. ' 

48 'Install it with: pip install azure-search-documents' 

49 ) 

50 

51 try: 

52 # Create search client with API key authentication 

53 credential = AzureKeyCredential(api_key) 

54 search_client = SearchClient(endpoint=endpoint, index_name=index_name, credential=credential) 

55 

56 # Perform the search 

57 results = search_client.search(search_text=query, top=count) 

58 

59 # Convert results to list and extract fields 

60 search_results = [] 

61 for result in results: 

62 # Azure AI Search returns documents with custom schemas 

63 # We need to extract common fields that might represent URL, title, and content 

64 # Common field names to look for: 

65 result_dict = dict(result) 

66 

67 # Try to find URL field (common names) 

68 link = ( 

69 result_dict.get('url') 

70 or result_dict.get('link') 

71 or result_dict.get('uri') 

72 or result_dict.get('metadata_storage_path') 

73 or '' 

74 ) 

75 

76 # Try to find title field (common names) 

77 title = ( 

78 result_dict.get('title') 

79 or result_dict.get('name') 

80 or result_dict.get('metadata_title') 

81 or result_dict.get('metadata_storage_name') 

82 or None 

83 ) 

84 

85 # Try to find content/snippet field (common names) 

86 snippet = ( 

87 result_dict.get('content') 

88 or result_dict.get('snippet') 

89 or result_dict.get('description') 

90 or result_dict.get('summary') 

91 or result_dict.get('text') 

92 or None 

93 ) 

94 

95 # Truncate snippet if too long 

96 if snippet and len(snippet) > 500: 

97 snippet = snippet[:497] + '...' 

98 

99 if link: # Only add if we found a valid link 

100 search_results.append( 

101 { 

102 'link': link, 

103 'title': title, 

104 'snippet': snippet, 

105 } 

106 ) 

107 

108 # Apply domain filtering if specified 

109 if filter_list: 

110 search_results = get_filtered_results(search_results, filter_list) 

111 

112 # Convert to SearchResult objects 

113 return [ 

114 SearchResult( 

115 link=result['link'], 

116 title=result.get('title'), 

117 snippet=result.get('snippet'), 

118 ) 

119 for result in search_results 

120 ] 

121 

122 except Exception as ex: 

123 log.error(f'Azure AI Search error: {ex}') 

124 raise ex