Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/citations/_models.py: 98%

1"""Citation domain models and enums.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass, field 

6from datetime import UTC, datetime 

7from enum import StrEnum 

8from typing import Any 

9 

10 

11class CitationStyle(StrEnum): 

12 """Citation formatting styles.""" 

13 

14 NUMERIC = "numeric" # [1], [2], etc. 

15 AUTHOR_YEAR = "author_year" # (Smith, 2023) 

16 FOOTNOTE = "footnote" # Superscript numbers 

17 INLINE = "inline" # Inline source references 

18 APA = "apa" # APA style 

19 MLA = "mla" # MLA style 

20 CHICAGO = "chicago" # Chicago style 

21 CUSTOM = "custom" # Custom format 

22 

23 

24class SourceType(StrEnum): 

25 """Types of sources.""" 

26 

27 DOCUMENT = "document" 

28 WEB_PAGE = "web_page" 

29 ARTICLE = "article" 

30 BOOK = "book" 

31 PAPER = "paper" 

32 DATABASE = "database" 

33 API = "api" 

34 INTERNAL = "internal" 

35 UNKNOWN = "unknown" 

36 

37 

38@dataclass 

39class Source: 

40 """A source of information.""" 

41 

42 id: str 

43 content: str 

44 source_type: SourceType = SourceType.UNKNOWN 

45 title: str | None = None 

46 author: str | None = None 

47 url: str | None = None 

48 publication_date: str | None = None 

49 page_number: int | None = None 

50 chapter: str | None = None 

51 metadata: dict[str, Any] = field(default_factory=dict) 

52 timestamp: str = field( 

53 default_factory=lambda: datetime.now(UTC).isoformat(), 

54 ) 

55 

56 def __repr__(self) -> str: 

57 """Return string representation.""" 

58 parts = [f"id={self.id}"] 

59 if self.title: 

60 parts.append(f"title={self.title[:30]}") 

61 if self.author: 

62 parts.append(f"author={self.author}") 

63 return f"Source({', '.join(parts)})" 

64 

65 

66@dataclass 

67class Citation: 

68 """A citation linking text to source.""" 

69 

70 source_id: str 

71 text_span: str 

72 start_char: int | None = None 

73 end_char: int | None = None 

74 confidence: float = 1.0 

75 relevance_score: float = 1.0 

76 citation_number: int | None = None 

77 metadata: dict[str, Any] = field(default_factory=dict) 

78 

79 def __repr__(self) -> str: 

80 """Return string representation.""" 

81 return ( 

82 f"Citation(source={self.source_id}, " 

83 f"span='{self.text_span[:30]}...', " 

84 f"conf={self.confidence:.2f})" 

85 ) 

86 

87 

88@dataclass 

89class CitedResponse: 

90 """A response with citations.""" 

91 

92 text: str 

93 sources: list[Source] 

94 citations: list[Citation] 

95 citation_style: CitationStyle = CitationStyle.NUMERIC 

96 metadata: dict[str, Any] = field(default_factory=dict) 

97 timestamp: str = field( 

98 default_factory=lambda: datetime.now(UTC).isoformat(), 

99 ) 

100 

101 @property 

102 def num_sources(self) -> int: 

103 """Number of unique sources.""" 

104 return len(self.sources) 

105 

106 @property 

107 def num_citations(self) -> int: 

108 """Number of citations.""" 

109 return len(self.citations) 

110 

111 @property 

112 def avg_confidence(self) -> float: 

113 """Average citation confidence.""" 

114 if not self.citations: 

115 return 0.0 

116 return sum(c.confidence for c in self.citations) / len(self.citations) 

117 

118 def get_source(self, source_id: str) -> Source | None: 

119 """Get source by ID.""" 

120 for source in self.sources: 

121 if source.id == source_id: 

122 return source 

123 return None 

124 

125 def get_citations_for_source(self, source_id: str) -> list[Citation]: 

126 """Get all citations for a source.""" 

127 return list(filter(lambda c: c.source_id == source_id, self.citations)) 

128 

129 def __repr__(self) -> str: 

130 """Return string representation.""" 

131 return ( 

132 f"CitedResponse(sources={self.num_sources}, " 

133 f"citations={self.num_citations}, " 

134 f"style={self.citation_style.value})" 

135 )