Coverage for paperless_ai/prompts/render.py: 89%

1import dataclasses 

2import enum 

3from typing import ClassVar 

4 

5from jinja2 import Environment 

6from jinja2 import PackageLoader 

7from jinja2 import StrictUndefined 

8 

9 

10class PromptName(enum.Enum): 

11 CLASSIFICATION = "classification" 

12 CLASSIFICATION_RAG_CONTEXT = "classification_rag_context" 

13 LOCALIZATION = "localization" 

14 TAXONOMY_BLOCK = "taxonomy_block" 

15 CHAT_QA = "chat_qa" 

16 CHAT_REFINE = "chat_refine" 

17 

18 

19@dataclasses.dataclass(frozen=True, slots=True) 

20class PromptContext: 

21 template_name: ClassVar[PromptName] 

22 

23 

24# Every render here goes through Environment.get_template() and 

25# .render(**dataclasses.asdict(context)). This is variable substitution, 

26# never a template-source compile. If you're about to call from_string()/Template() 

27# on anything derived from user input, stop: that needs a sandboxed 

28# environment (see documents/templating/environment.py), not this one. 

29_env = Environment( 

30 loader=PackageLoader("paperless_ai", "prompts"), 

31 trim_blocks=True, 

32 lstrip_blocks=True, 

33 keep_trailing_newline=False, 

34 autoescape=False, 

35 undefined=StrictUndefined, 

36) 

37 

38 

39def render_prompt(context: PromptContext) -> str: 

40 template = _env.get_template(f"{context.template_name.value}.j2") 

41 return template.render(**dataclasses.asdict(context)).strip()