1"""Citation formatter implementations."""
2
3from __future__ import annotations
4
5from abc import ABC, abstractmethod
6
7from lexigram.ai.rag.citations._models import (
8 Citation,
9 CitedResponse,
10 Source,
11 SourceType,
12)
13
14
15class AbstractCitationFormatter(ABC):
16 """Base class for citation formatters."""
17
18 @abstractmethod
19 def format_citation(
20 self,
21 citation: Citation,
22 source: Source,
23 citation_number: int | None = None,
24 ) -> str:
25 """Format a single citation.
26
27 Args:
28 citation: Citation to format
29 source: Source being cited
30 citation_number: Optional citation number
31
32 Returns:
33 Formatted citation string
34 """
35
36 @abstractmethod
37 def format_bibliography_entry(
38 self,
39 source: Source,
40 number: int | None = None,
41 ) -> str:
42 """Format a bibliography entry.
43
44 Args:
45 source: Source to format
46 number: Optional entry number
47
48 Returns:
49 Formatted bibliography entry
50 """
51
52 def format_response(self, cited_response: CitedResponse) -> str:
53 """Format complete response with citations.
54
55 Args:
56 cited_response: Response with citations
57
58 Returns:
59 Formatted response string with inline citations and bibliography
60 """
61 # Build text with inline citations
62 text = cited_response.text
63
64 # Sort citations by position (if available)
65 sorted_citations = sorted(
66 cited_response.citations,
67 key=lambda c: c.start_char if c.start_char is not None else 0,
68 reverse=True, # Start from end to preserve positions
69 )
70
71 # Insert citations
72 for citation in sorted_citations:
73 source = cited_response.get_source(citation.source_id)
74 if (
75 source
76 and citation.start_char is not None
77 and citation.end_char is not None
78 ):
79 citation_marker = self.format_citation(
80 citation,
81 source,
82 citation.citation_number,
83 )
84 # Insert after the cited text
85 text = (
86 text[: citation.end_char]
87 + citation_marker
88 + text[citation.end_char :]
89 )
90
91 # Add bibliography
92 bibliography = self.format_bibliography(cited_response.sources)
93
94 return f"{text}\n\n{bibliography}"
95
96 def format_bibliography(self, sources: list[Source]) -> str:
97 """Format bibliography from sources.
98
99 Args:
100 sources: List of sources
101
102 Returns:
103 Formatted bibliography
104 """
105 if not sources:
106 return ""
107
108 entries = []
109 for i, source in enumerate(sources, 1):
110 entry = self.format_bibliography_entry(source, i)
111 entries.append(entry)
112
113 return "References:\n" + "\n".join(entries)
114
115
116class NumericCitationFormatter(AbstractCitationFormatter):
117 """Numeric citation style [1], [2], etc."""
118
119 def format_citation(
120 self,
121 citation: Citation,
122 source: Source,
123 citation_number: int | None = None,
124 ) -> str:
125 """Format as [1], [2], etc."""
126 num = citation_number or citation.citation_number or 1
127 return f"[{num}]"
128
129 def format_bibliography_entry(
130 self,
131 source: Source,
132 number: int | None = None,
133 ) -> str:
134 """Format bibliography entry."""
135 num = number or 1
136 parts = [f"[{num}]"]
137
138 if source.author:
139 parts.append(source.author)
140 if source.title:
141 parts.append(f'"{source.title}"')
142 if source.publication_date:
143 parts.append(f"({source.publication_date})")
144 if source.url:
145 parts.append(source.url)
146
147 return " ".join(parts)
148
149
150class AuthorYearCitationFormatter(AbstractCitationFormatter):
151 """Author-year citation style (Smith, 2023)."""
152
153 def format_citation(
154 self,
155 citation: Citation,
156 source: Source,
157 citation_number: int | None = None,
158 ) -> str:
159 """Format as (Author, Year)."""
160 author = source.author or "Unknown"
161 year = source.publication_date or "n.d."
162
163 # Extract just year if full date provided
164 if len(year) > 4:
165 year = year[:4]
166
167 return f" ({author}, {year})"
168
169 def format_bibliography_entry(
170 self,
171 source: Source,
172 number: int | None = None,
173 ) -> str:
174 """Format bibliography entry."""
175 author = source.author or "Unknown"
176 year = source.publication_date or "n.d."
177 if len(year) > 4:
178 year = year[:4]
179
180 parts = [f"{author} ({year})."]
181
182 if source.title:
183 parts.append(f"{source.title}.")
184 if source.url:
185 parts.append(f"Retrieved from {source.url}")
186
187 return " ".join(parts)
188
189
190class FootnoteCitationFormatter(AbstractCitationFormatter):
191 """Footnote style with superscript numbers."""
192
193 def format_citation(
194 self,
195 citation: Citation,
196 source: Source,
197 citation_number: int | None = None,
198 ) -> str:
199 """Format as superscript number."""
200 num = citation_number or citation.citation_number or 1
201 # Using Unicode superscript numbers
202 superscripts = "⁰¹²³⁴⁵⁶⁷⁸⁹"
203 if num < 10:
204 return superscripts[num]
205 # For larger numbers, just use regular format
206 return f"^{num}"
207
208 def format_bibliography_entry(
209 self,
210 source: Source,
211 number: int | None = None,
212 ) -> str:
213 """Format as footnote."""
214 num = number or 1
215 parts = [f"{num}."]
216
217 if source.author:
218 parts.append(source.author + ",")
219 if source.title:
220 parts.append(f'"{source.title},"')
221 if source.publication_date:
222 parts.append(source.publication_date + ",")
223 if source.url:
224 parts.append(source.url)
225
226 return " ".join(parts)
227
228
229class InlineCitationFormatter(AbstractCitationFormatter):
230 """Inline source references."""
231
232 def format_citation(
233 self,
234 citation: Citation,
235 source: Source,
236 citation_number: int | None = None,
237 ) -> str:
238 """Format as inline reference."""
239 parts = []
240 if source.title:
241 parts.append(source.title)
242 if source.author:
243 parts.append(f"by {source.author}")
244
245 if parts:
246 return f" (Source: {', '.join(parts)})"
247 return f" (Source: {source.id})"
248
249 def format_bibliography_entry(
250 self,
251 source: Source,
252 number: int | None = None,
253 ) -> str:
254 """Format bibliography entry."""
255 parts = []
256
257 if source.author:
258 parts.append(source.author)
259 if source.title:
260 parts.append(f'"{source.title}"')
261 if source.publication_date:
262 parts.append(f"({source.publication_date})")
263 if source.url:
264 parts.append(source.url)
265
266 return " - ".join(parts) if parts else source.id
267
268
269class APACitationFormatter(AbstractCitationFormatter):
270 """APA citation style."""
271
272 def format_citation(
273 self,
274 citation: Citation,
275 source: Source,
276 citation_number: int | None = None,
277 ) -> str:
278 """Format as APA in-text citation."""
279 author = source.author or "Unknown"
280 year = source.publication_date or "n.d."
281 if len(year) > 4:
282 year = year[:4]
283
284 # Extract last name if full name provided
285 if "," in author:
286 author = author.split(",")[0]
287 elif " " in author:
288 author = author.split()[-1]
289
290 return f" ({author}, {year})"
291
292 def format_bibliography_entry(
293 self,
294 source: Source,
295 number: int | None = None,
296 ) -> str:
297 """Format as APA reference entry."""
298 author = source.author or "Unknown"
299 year = source.publication_date or "n.d."
300 if len(year) > 4:
301 year = year[:4]
302
303 parts = [f"{author} ({year})."]
304
305 if source.title:
306 parts.append(f"{source.title}.")
307
308 if source.source_type == SourceType.WEB_PAGE and source.url:
309 parts.append(f"Retrieved from {source.url}")
310 elif source.url:
311 parts.append(f"DOI: {source.url}")
312
313 return " ".join(parts)