Coverage for src / quber / playground / answers / document.py: 100%

1"""The document an answer is about, stated to the model. 

2 

3A chunk carries its page heading and its table's caption, and a table's header 

4rows usually carry the period. A chart on a slide carries none of that: the 

5deck's date is on its cover, not on the slide. Handed a chart alone, the 

6answer model refuses to date it, which is the right reading of its own rules. 

7The period is a property of the document, and the documents table already 

8holds it, so the route states it ahead of the context on every request. 

9 

10The header is data the route knows, not prompt text: it is built from the 

11document row and prepended to the user message, and the system prompt does not 

12change. 

13""" 

14 

15from __future__ import annotations 

16 

17from dataclasses import dataclass 

18from typing import Any, Mapping, Optional 

19 

20 

21@dataclass(frozen=True) 

22class DocumentIdentity: 

23 title: Optional[str] = None 

24 issuer: Optional[str] = None 

25 form: Optional[str] = None 

26 period: Optional[str] = None 

27 year: Optional[int] = None 

28 

29 @classmethod 

30 def from_row(cls, row: Mapping[str, Any]) -> "DocumentIdentity": 

31 """From a documents-table row as the routes carry it (`_doc_dict`).""" 

32 return cls( 

33 title=row.get("title"), 

34 issuer=row.get("folder"), 

35 form=row.get("filing_type"), 

36 period=row.get("period"), 

37 year=row.get("year"), 

38 ) 

39 

40 def header(self) -> str: 

41 """The block that opens the user message. Fields the row lacks are left out.""" 

42 period = " ".join(str(p) for p in (self.period, self.year) if p) 

43 fields = [ 

44 ("Title", self.title), 

45 ("Issuer", self.issuer), 

46 ("Form", self.form), 

47 ("Period", period or None), 

48 ] 

49 lines = ["Document"] + [f" {name}: {value}" for name, value in fields if value] 

50 lines.append( 

51 "Every figure in the context below is from this document unless a chunk states another period." 

52 ) 

53 return "\n".join(lines) 

54 

55 

56def with_document(prompt: str, document: Optional[DocumentIdentity]) -> str: 

57 """The prompt with the document stated first, or unchanged when none is known.""" 

58 if document is None: 

59 return prompt 

60 return f"{document.header()}\n\n{prompt}"