Coverage for open_webui/retrieval/loaders/external_document.py: 16%

1import logging 

2import os 

3from typing import Iterator, List, Union 

4from urllib.parse import quote 

5 

6import requests 

7from langchain_core.document_loaders import BaseLoader 

8from langchain_core.documents import Document 

9from open_webui.utils.headers import include_user_info_headers, parse_custom_headers 

10 

11log = logging.getLogger(__name__) 

12 

13 

14class ExternalDocumentLoader(BaseLoader): 

15 def __init__( 

16 self, 

17 file_path, 

18 url: str, 

19 api_key: str, 

20 mime_type=None, 

21 user=None, 

22 user_groups=None, 

23 headers=None, 

24 metadata=None, 

25 **kwargs, 

26 ) -> None: 

27 self.url = url 

28 self.api_key = api_key 

29 

30 self.file_path = file_path 

31 self.mime_type = mime_type 

32 

33 self.user = user 

34 self.user_groups = user_groups 

35 self.headers = headers 

36 self.metadata = metadata 

37 

38 def load(self) -> List[Document]: 

39 with open(self.file_path, 'rb') as f: 

40 data = f.read() 

41 

42 headers = {} 

43 if self.mime_type is not None: 

44 headers['Content-Type'] = self.mime_type 

45 

46 if self.api_key is not None: 

47 headers['Authorization'] = f'Bearer {self.api_key}' 

48 

49 try: 

50 headers['X-Filename'] = quote(os.path.basename(self.file_path)) 

51 except Exception: 

52 pass 

53 

54 headers.update(parse_custom_headers(self.headers, self.user, self.metadata, user_groups=self.user_groups)) 

55 

56 if self.user is not None: 

57 headers = include_user_info_headers(headers, self.user) 

58 

59 url = self.url 

60 if url.endswith('/'): 

61 url = url[:-1] 

62 

63 try: 

64 response = requests.put(f'{url}/process', data=data, headers=headers) 

65 except Exception as e: 

66 log.error(f'Error connecting to endpoint: {e}') 

67 raise Exception(f'Error connecting to endpoint: {e}') 

68 

69 if response.ok: 

70 response_data = response.json() 

71 if response_data: 

72 if isinstance(response_data, dict): 

73 return [ 

74 Document( 

75 page_content=response_data.get('page_content'), 

76 metadata=response_data.get('metadata'), 

77 ) 

78 ] 

79 elif isinstance(response_data, list): 

80 documents = [] 

81 for document in response_data: 

82 documents.append( 

83 Document( 

84 page_content=document.get('page_content'), 

85 metadata=document.get('metadata'), 

86 ) 

87 ) 

88 return documents 

89 else: 

90 raise Exception('Error loading document: Unable to parse content') 

91 

92 else: 

93 raise Exception('Error loading document: No content returned') 

94 else: 

95 raise Exception(f'Error loading document: {response.status_code} {response.text}')