Coverage for paperless_ai/prompts/context.py: 100%

1from dataclasses import dataclass 

2from typing import ClassVar 

3 

4from paperless_ai.prompts.render import PromptContext 

5from paperless_ai.prompts.render import PromptName 

6 

7 

8@dataclass(frozen=True, slots=True) 

9class TaxonomyBlockPromptContext(PromptContext): 

10 template_name: ClassVar[PromptName] = PromptName.TAXONOMY_BLOCK 

11 candidate_payload_json: str 

12 

13 

14@dataclass(frozen=True, slots=True) 

15class ClassificationPromptContext(PromptContext): 

16 template_name: ClassVar[PromptName] = PromptName.CLASSIFICATION 

17 filename: str 

18 content: str 

19 taxonomy_block: str 

20 has_candidates: bool 

21 

22 

23@dataclass(frozen=True, slots=True) 

24class RagContextPromptContext(PromptContext): 

25 template_name: ClassVar[PromptName] = PromptName.CLASSIFICATION_RAG_CONTEXT 

26 base_prompt: str 

27 context: str 

28 

29 

30@dataclass(frozen=True, slots=True) 

31class LocalizationPromptContext(PromptContext): 

32 template_name: ClassVar[PromptName] = PromptName.LOCALIZATION 

33 language_name: str 

34 suggestions_json: str 

35 

36 

37@dataclass(frozen=True, slots=True) 

38class ChatQaPromptContext(PromptContext): 

39 template_name: ClassVar[PromptName] = PromptName.CHAT_QA 

40 output_language: str | None 

41 

42 

43@dataclass(frozen=True, slots=True) 

44class ChatRefinePromptContext(PromptContext): 

45 template_name: ClassVar[PromptName] = PromptName.CHAT_REFINE 

46 output_language: str | None