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

1"""Citation tracker and extraction helpers.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any 

6 

7from lexigram.ai.rag.citations._formatters import ( 

8 AbstractCitationFormatter, 

9 APACitationFormatter, 

10 AuthorYearCitationFormatter, 

11 FootnoteCitationFormatter, 

12 InlineCitationFormatter, 

13 NumericCitationFormatter, 

14) 

15from lexigram.ai.rag.citations._models import ( 

16 Citation, 

17 CitationStyle, 

18 CitedResponse, 

19 Source, 

20 SourceType, 

21) 

22 

23 

24class CitationTracker: 

25 """Tracks citations throughout RAG pipeline.""" 

26 

27 def __init__(self, citation_style: CitationStyle = CitationStyle.NUMERIC): 

28 """Initialize citation tracker. 

29 

30 Args: 

31 citation_style: Default citation style 

32 """ 

33 self.citation_style = citation_style 

34 self.sources: dict[str, Source] = {} 

35 self.citations: list[Citation] = [] 

36 self._citation_counter = 0 

37 

38 def add_source(self, source: Source) -> None: 

39 """Add a source to tracking. 

40 

41 Args: 

42 source: Source to add 

43 """ 

44 self.sources[source.id] = source 

45 

46 def add_citation( 

47 self, 

48 source_id: str, 

49 text_span: str, 

50 start_char: int | None = None, 

51 end_char: int | None = None, 

52 confidence: float = 1.0, 

53 relevance_score: float = 1.0, 

54 ) -> Citation: 

55 """Add a citation. 

56 

57 Args: 

58 source_id: ID of source being cited 

59 text_span: Text being cited 

60 start_char: Start character position 

61 end_char: End character position 

62 confidence: Citation confidence 

63 relevance_score: Relevance score 

64 

65 Returns: 

66 Created citation 

67 

68 Raises: 

69 ValueError: If source not found 

70 """ 

71 if source_id not in self.sources: 

72 msg = f"Source {source_id} not found. Add source first." 

73 raise ValueError(msg) 

74 

75 self._citation_counter += 1 

76 

77 citation = Citation( 

78 source_id=source_id, 

79 text_span=text_span, 

80 start_char=start_char, 

81 end_char=end_char, 

82 confidence=confidence, 

83 relevance_score=relevance_score, 

84 citation_number=self._citation_counter, 

85 ) 

86 

87 self.citations.append(citation) 

88 return citation 

89 

90 def create_cited_response(self, text: str, **metadata: Any) -> CitedResponse: 

91 """Create cited response from tracked data. 

92 

93 Args: 

94 text: Response text 

95 **metadata: Additional metadata 

96 

97 Returns: 

98 CitedResponse with all tracked sources and citations 

99 """ 

100 return CitedResponse( 

101 text=text, 

102 sources=list(self.sources.values()), 

103 citations=self.citations, 

104 citation_style=self.citation_style, 

105 metadata=metadata, 

106 ) 

107 

108 def format_response(self, cited_response: CitedResponse) -> str: 

109 """Format response with citations. 

110 

111 Args: 

112 cited_response: Response to format 

113 

114 Returns: 

115 Formatted response with inline citations and bibliography 

116 """ 

117 formatter = self._get_formatter(cited_response.citation_style) 

118 return formatter.format_response(cited_response) 

119 

120 def _get_formatter(self, style: CitationStyle) -> AbstractCitationFormatter: 

121 """Get formatter for citation style. 

122 

123 Args: 

124 style: Citation style 

125 

126 Returns: 

127 Appropriate formatter 

128 

129 Raises: 

130 ValueError: If style not supported 

131 """ 

132 formatters = { 

133 CitationStyle.NUMERIC: NumericCitationFormatter(), 

134 CitationStyle.AUTHOR_YEAR: AuthorYearCitationFormatter(), 

135 CitationStyle.FOOTNOTE: FootnoteCitationFormatter(), 

136 CitationStyle.INLINE: InlineCitationFormatter(), 

137 CitationStyle.APA: APACitationFormatter(), 

138 } 

139 

140 formatter = formatters.get(style) 

141 if not formatter: 

142 msg = f"Citation style {style} not yet supported" 

143 raise ValueError(msg) 

144 

145 return formatter 

146 

147 def reset(self) -> None: 

148 """Reset tracker state.""" 

149 self.sources.clear() 

150 self.citations.clear() 

151 self._citation_counter = 0 

152 

153 

154def extract_citations_from_chunks( 

155 response_text: str, 

156 retrieved_chunks: list[dict[str, Any]], 

157 citation_style: CitationStyle = CitationStyle.NUMERIC, 

158) -> CitedResponse: 

159 """Extract citations from retrieved chunks. 

160 

161 Args: 

162 response_text: Generated response text 

163 retrieved_chunks: List of retrieved chunks with metadata 

164 citation_style: Citation style to use 

165 

166 Returns: 

167 CitedResponse with citations from chunks 

168 """ 

169 tracker = CitationTracker(citation_style=citation_style) 

170 

171 # Add sources from chunks 

172 for i, chunk in enumerate(retrieved_chunks): 

173 source_id = chunk.get("id", f"source_{i}") 

174 content = chunk.get("content", "") 

175 metadata = chunk.get("metadata", {}) 

176 

177 source = Source( 

178 id=source_id, 

179 content=content, 

180 source_type=SourceType(metadata.get("type", "document")), 

181 title=metadata.get("title"), 

182 author=metadata.get("author"), 

183 url=metadata.get("url"), 

184 publication_date=metadata.get("date"), 

185 page_number=metadata.get("page"), 

186 metadata=metadata, 

187 ) 

188 

189 tracker.add_source(source) 

190 

191 # Simple citation extraction: assume each chunk contributed to response 

192 # In production, would use more sophisticated attribution 

193 relevance_score = chunk.get("score", 1.0) 

194 

195 tracker.add_citation( 

196 source_id=source_id, 

197 text_span=content[:100], # First 100 chars as representative span 

198 confidence=0.8, # Default confidence 

199 relevance_score=relevance_score, 

200 ) 

201 

202 return tracker.create_cited_response(response_text)