Coverage for paperless_ai/embedding.py: 31%

1import re 

2from typing import TYPE_CHECKING 

3 

4from django.conf import settings 

5 

6if TYPE_CHECKING: 6 ↛ 7line 6 didn't jump to line 7 because the condition on line 6 was never true

7 from llama_index.core.base.embeddings.base import BaseEmbedding 

8 

9from documents.models import Document 

10from paperless.config import AIConfig 

11from paperless.models import LLMEmbeddingBackend 

12from paperless.network import PinnedHostAsyncHTTPTransport 

13from paperless.network import PinnedHostHTTPTransport 

14from paperless.network import create_pinned_async_httpx_client 

15from paperless.network import create_pinned_httpx_client 

16from paperless.network import validate_outbound_http_url 

17from paperless_ai.client import PLACEHOLDER_API_KEY 

18 

19OCR_LEADER_REGEX = re.compile(r"[._\-\u00b7]{4,}") 

20HORIZONTAL_WHITESPACE_REGEX = re.compile(r"[ \t\u00a0]+") 

21 

22 

23def get_embedding_model(config: AIConfig) -> "BaseEmbedding": 

24 match config.llm_embedding_backend: 

25 case LLMEmbeddingBackend.OPENAI_LIKE: 

26 from llama_index.embeddings.openai_like import OpenAILikeEmbedding 

27 

28 endpoint = config.llm_embedding_endpoint or config.llm_endpoint or None 

29 http_client = None 

30 async_http_client = None 

31 if endpoint: 

32 http_client = create_pinned_httpx_client( 

33 endpoint, 

34 allow_internal=config.llm_allow_internal_endpoints, 

35 timeout=config.llm_request_timeout, 

36 ) 

37 async_http_client = create_pinned_async_httpx_client( 

38 endpoint, 

39 allow_internal=config.llm_allow_internal_endpoints, 

40 timeout=config.llm_request_timeout, 

41 ) 

42 return OpenAILikeEmbedding( 

43 model_name=config.llm_embedding_model or "text-embedding-3-small", 

44 api_key=config.llm_embedding_api_key 

45 or config.llm_api_key 

46 or PLACEHOLDER_API_KEY, 

47 api_base=endpoint, 

48 timeout=config.llm_request_timeout, 

49 http_client=http_client, 

50 async_http_client=async_http_client, 

51 ) 

52 case LLMEmbeddingBackend.HUGGINGFACE: 

53 from llama_index.embeddings.huggingface import HuggingFaceEmbedding 

54 

55 return HuggingFaceEmbedding( 

56 model_name=config.llm_embedding_model 

57 or "sentence-transformers/all-MiniLM-L6-v2", 

58 cache_folder=str(settings.DATA_DIR / "hf_cache"), 

59 ) 

60 case LLMEmbeddingBackend.OLLAMA: 

61 from llama_index.embeddings.ollama import OllamaEmbedding 

62 from ollama import AsyncClient 

63 from ollama import Client 

64 

65 endpoint = ( 

66 config.llm_embedding_endpoint 

67 or config.llm_endpoint 

68 or "http://localhost:11434" 

69 ) 

70 validate_outbound_http_url( 

71 endpoint, 

72 allow_internal=config.llm_allow_internal_endpoints, 

73 ) 

74 embedding = OllamaEmbedding( 

75 model_name=config.llm_embedding_model or "embeddinggemma", 

76 base_url=endpoint, 

77 ollama_additional_kwargs={"num_ctx": config.llm_context_size}, 

78 ) 

79 embedding._client = Client( 

80 host=endpoint, 

81 timeout=config.llm_request_timeout, 

82 transport=PinnedHostHTTPTransport( 

83 allow_internal=config.llm_allow_internal_endpoints, 

84 ), 

85 ) 

86 embedding._async_client = AsyncClient( 

87 host=endpoint, 

88 timeout=config.llm_request_timeout, 

89 transport=PinnedHostAsyncHTTPTransport( 

90 allow_internal=config.llm_allow_internal_endpoints, 

91 ), 

92 ) 

93 return embedding 

94 case _: 

95 raise ValueError( 

96 f"Unsupported embedding backend: {config.llm_embedding_backend}", 

97 ) 

98 

99 

100_DEFAULT_MODEL_NAMES = { 

101 LLMEmbeddingBackend.OPENAI_LIKE: "text-embedding-3-small", 

102 LLMEmbeddingBackend.HUGGINGFACE: "sentence-transformers/all-MiniLM-L6-v2", 

103 LLMEmbeddingBackend.OLLAMA: "embeddinggemma", 

104} 

105 

106 

107def get_configured_model_name(config: AIConfig) -> str: 

108 """Return the canonical name of the currently configured embedding model.""" 

109 # dict.get(key, default) overload resolution fails for TextChoices keys in some 

110 # type checkers; use `or` fallback to avoid the ambiguity. 

111 default = ( 

112 _DEFAULT_MODEL_NAMES.get( 

113 config.llm_embedding_backend, 

114 ) 

115 or "sentence-transformers/all-MiniLM-L6-v2" 

116 ) 

117 return config.llm_embedding_model or default 

118 

119 

120def _normalize_llm_index_text(text: str) -> str: 

121 text = OCR_LEADER_REGEX.sub(" ", text) 

122 return HORIZONTAL_WHITESPACE_REGEX.sub(" ", text) 

123 

124 

125def build_llm_index_text(doc: Document) -> str: 

126 # Short structured fields (filename, storage path, ASN, title, tags, ...) live 

127 # in node.metadata: excluded from embeddings, shown to the LLM via metadata 

128 # prepend. Notes and Custom Fields stay in the body: Notes can be long free 

129 # text, Custom Fields are dynamic in count and best kept in the embedding. 

130 lines = [ 

131 f"Notes: {','.join([str(c.note) for c in doc.notes.all()])}", 

132 ] 

133 

134 for instance in doc.custom_fields.all(): 

135 lines.append(f"Custom Field - {instance.field.name}: {instance}") 

136 

137 lines.append("\nContent:\n") 

138 lines.append(doc.content or "") 

139 

140 return _normalize_llm_index_text("\n".join(lines))