29327663402

Committed 14 Jul 2026 11:05AM UTC coverage: 97.426% (-0.4%) from 97.804%

Build # 29327663402

Build Type

push

github

Committed by

davidmezzetti

Commit Message

Code cleanup

Coverage Stats

757 of 777 relevant lines covered (97.43%)

0.97 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.43

/src/python/paperai/report/common.py

"""
Report module
"""

import regex as re

from txtai.pipeline import Labels, RAG, Similarity, Tokenizer

from ..index import Index
from ..query import Query

from .column import Column


class Report:
    """
    Methods to build reports from a series of queries
    """

    def __init__(self, embeddings, db, options):
        """
        Creates a new report.

        Args:
            embeddings: embeddings index
            db: database connection
            options: report options
        """

        # Store references to embeddings index and open database cursor
        self.embeddings = embeddings
        self.cur = db.cursor()

        # Report options
        self.options = options

        # Column names
        self.names = []

        # Create similarity and labels pipeline, if necessary
        self.similarity = Similarity(options["similarity"]) if "similarity" in options else None
        self.labels = Labels(model=self.similarity) if self.similarity else None

        # Create RAG pipeline arguments without report options
        args = {x: options[x] for x in options if x not in ["topn", "render", "path", "qa", "indir", "threshold"]}

        # Retrieval Augmented Generation (RAG) pipeline for calculated fields
        self.rag = RAG(
            self.similarity if self.similarity else self.embeddings,
            args["llm"] if args.get("llm") else args["qa"] if args.get("qa") else "NeuML/bert-small-cord19qa",
            **args,
        )

    def build(self, queries, options, output):
        """
        Builds a report using a list of input queries

        Args:
            queries: queries to execute
            options: report options
            output: output I/O object
        """

        # Default to 50 documents if not specified
        topn = options.get("topn", 50)

        for name, config in queries:
            query = config["query"]
            columns = config["columns"]

            # Write query string
            self.query(output, name, query)

            # Write separator
            self.separator(output)

            # Query for best matches
            results = Query.search(self.embeddings, self.cur, query, topn, options.get("threshold"))

            # Generate highlights section
            self.section(output, "Highlights")

            # Generate highlights
            self.highlights(output, results, int(topn / 10))

            # Separator between highlights and articles
            self.separator(output)

            # Generate articles section
            self.section(output, "Articles")

            # Generate table headers
            self.headers([column["name"] for column in columns], output)

            # Generate table rows
            self.articles(output, topn, (name, query, columns), results)

            # Write section separator
            self.separator(output)

    def highlights(self, output, results, topn):
        """
        Builds a highlights section.

        Args:
            output: output file
            results: search results
            topn: number of results to return
        """

        # Extract top sections as highlights
        for highlight in Query.highlights(results, topn):
            # Get matching article
            uid = [article for _, _, article, text in results if text == highlight][0]
            self.cur.execute("SELECT Authors, Reference FROM articles WHERE id = ?", [uid])
            article = self.cur.fetchone()

            # Write out highlight row
            self.highlight(output, article, highlight)

    def articles(self, output, topn, metadata, results):
        """
        Builds an articles section.

        Args:
            output: output file
            topn: number of documents to return
            metadata: query metadata
            results: search results
        """

        # Unpack metadata
        _, query, _ = metadata

        # Retrieve list of documents
        documents = Query.all(self.cur) if query == "*" else Query.documents(results, topn)

        # Collect matching rows
        rows = []

        for x, uid in enumerate(documents):
            # Get article metadata
            self.cur.execute(
                "SELECT Published, Title, Reference, Publication, Source, Entry, Id FROM articles WHERE id = ?",
                [uid],
            )
            article = self.cur.fetchone()

            if x and x % 100 == 0:
                print(f"Processed {x} documents", end="\r")

            # Calculate derived fields
            calculated = self.calculate(uid, metadata)

            # Builds a row for article
            rows.append(self.buildRow(article, documents[uid], calculated))

        # Print report by published desc
        for row in sorted(rows, key=lambda x: x["Date"], reverse=True):
            # Convert row dict to list
            row = [row[column] for column in self.names]

            # Write out row
            self.writeRow(output, row)

    def calculate(self, uid, metadata):
        """
        Builds a dict of calculated fields for a given document. This method calculates
        constant field columns and derived query columns. Derived query columns run through
        an embedding search and either run an additional QA query to extract a value or
        use the top n embedding search matches.

        Args:
            uid: article id
            metadata: query metadata

        Returns:
            {name: value} containing derived column values
        """

        # Parse column parameters
        fields, params = self.params(metadata)

        # Different type of calculations
        #  1. Similarity query
        #  2. Extractor query (similarity + question)
        #  3. Question-answering on other field
        queries, extractions, questions = [], [], []

        # Retrieve indexed document text for article
        sections = self.sections(uid)
        texts = [text for _, text in sections]

        for name, query, question, snippet, _, _, matches, _ in params:
            if query.startswith("$"):
                questions.append((name, query.replace("$", ""), question, snippet))
            elif matches:
                queries.append((name, query, matches))
            else:
                extractions.append((name, query, question, snippet))

        # Run all extractor queries against document text
        results = self.rag.query([query for _, query, _ in queries], texts)

        # Only execute embeddings queries for columns with matches set
        for x, (name, query, matches) in enumerate(queries):
            if results[x]:
                # Get topn text matches
                topn = [text for _, text, _ in results[x]][:matches]

                # Join results into String and return
                value = [self.resolve(params, sections, uid, name, value) for value in topn]
                fields[name] = "\n\n".join(value) if value else ""
            else:
                fields[name] = ""

        # Add extraction fields
        if extractions:
            for name, value in self.rag(extractions, texts, **self.options.get("params", {})):
                # Resolves the full value based on column parameters
                fields[name] = self.resolve(params, sections, uid, name, value) if value else ""

        # Add question fields
        if questions:
            for name, value in self.rag(questions, texts, **self.options.get("params", {})):
                # Resolves the full value based on column parameters
                fields[name] = self.resolve(params, sections, uid, name, value) if value else ""

        return fields

    def params(self, metadata):
        """
        Process and prepare parameters using input metadata.

        Args:
            metadata: query metadata

        Returns:
            fields, params - constant field values, query parameters for query columns
        """

        # Derived field values
        fields = {}

        # Query column parameters
        params = []

        # Unpack metadata
        _, _, columns = metadata

        for column in columns:
            # Constant column
            if "constant" in column:
                fields[column["name"]] = column["constant"]
            # Question-answer column
            elif "query" in column:
                # Query variable substitutions
                query = self.variables(column["query"], metadata)
                question = self.variables(column["question"], metadata) if "question" in column else query

                # Additional context parameters
                section = column.get("section", False)
                surround = column.get("surround", 0)
                matches = column.get("matches", 0)
                dtype = column.get("dtype")
                snippet = column.get("snippet", False)
                snippet = True if section or surround else snippet

                params.append(
                    (
                        column["name"],
                        query,
                        question,
                        snippet,
                        section,
                        surround,
                        matches,
                        dtype,
                    )
                )

        return fields, params

    def variables(self, value, metadata):
        """
        Runs variable substitution for value.

        Args:
            value: input value
            metadata: query metadata

        Returns:
            value with variable substitution
        """

        name, query, _ = metadata

        # Cleanup name for queries
        name = name.replace("_", "").lower()
        query = query.lower()

        if value:
            value = value.replace("$NAME", name).replace("$QUERY", query)

        return value

    def sections(self, uid):
        """
        Retrieves all sections as list for article with given uid.

        Args:
            uid: article id

        Returns:
            list of section text elements
        """

        # Retrieve indexed document text for article
        self.cur.execute(Index.SECTION_QUERY + " WHERE article = ? ORDER BY id", [uid])

        # Get list of document text sections
        sections = []
        for sid, name, text in self.cur.fetchall():
            if not self.embeddings.isweighted() or not name or not re.search(Index.SECTION_FILTER, name.lower()) or self.options.get("allsections"):
                # Check that section has at least 1 token
                if Tokenizer.tokenize(text):
                    sections.append((sid, text))

        return sections

    def resolve(self, params, sections, uid, name, value):
        """
        Fully resolves a value from an extractor call.

         - If section=True, this method pull the full section text
         - If surround is specified, this method will pull the surrounding text
         - Otherwise, the original value is returned

        Args:
            params: query parameters
            sections: section text
            uid: article id
            name: column name
            value: initial query value after running through extractor process

        Returns:
            resolved value
        """

        # Get all column parameters
        index = [params.index(x) for x in params if x[0] == name][0]
        _, _, _, _, section, surround, _, dtype = params[index]

        if value:
            # Find matching section
            sid = [sid for sid, text in sections if value in text]

            if sid:
                sid = sid[0]

                if section:
                    # Get full text for matching subsection
                    value = self.subsection(uid, sid)
                elif surround:
                    value = self.surround(uid, sid, surround)

            # Column dtype formatting
            if dtype == "int":
                value = Column.integer(value)
            elif isinstance(dtype, list):
                value = Column.categorical(self.labels, value, dtype)
            elif dtype in ["days", "weeks", "months", "years"]:
                value = Column.duration(value, dtype)

        return value

    def subsection(self, uid, sid):
        """
        Extracts all subsection text for columns with section=True.

        Args:
            uid: article id
            sid: section id

        Returns:
            full text for matching section
        """

        self.cur.execute(
            "SELECT Text FROM sections WHERE article = ? AND name = (SELECT name FROM sections WHERE id = ?)",
            [uid, sid],
        )
        return " ".join([x[0] for x in self.cur.fetchall()])

    def surround(self, uid, sid, size):
        """
        Extracts surrounding text for section with specified id.

        Args:
            uid: article id
            sid: section id
            size: number of surrounding lines to extract from each side

        Returns:
            matching text with surrounding context
        """

        self.cur.execute(
            "SELECT Text FROM sections WHERE article = ? AND id in (SELECT id FROM sections WHERE id >= ? AND id <= ?) AND "
            + "name = (SELECT name FROM sections WHERE id = ?)",
            [uid, sid - size, sid + size, sid],
        )

        return " ".join([x[0] for x in self.cur.fetchall()])

    def cleanup(self, outfile):
        """
        Allow freeing or cleaning up resources.

        Args:
            outfile: output file path
        """

    def query(self, output, task, query):
        """
        Writes query.

        Args:
            output: output file
            task: task name
            query: query string
        """

    def section(self, output, name):
        """
        Writes a section name

        Args:
            output: output file
            name: section name
        """

    def highlight(self, output, article, highlight):
        """
        Writes a highlight row

        Args:
            output: output file
            article: article reference
            highlight: highlight text
        """

    def headers(self, columns, output):
        """
        Writes table headers.

        Args:
            columns: column names
            output: output file
        """

    def buildRow(self, article, sections, calculated):
        """
        Converts a document to a table row.

        Args:
            article: article
            sections: text sections for article
            calculated: calculated fields
        """

    def writeRow(self, output, row):
        """
        Writes a table row.

        Args:
            output: output file
            row: output row
        """

    def separator(self, output):
        """
        Writes a separator between sections
        """

1	"""
2	Report module
3	"""
4
5	import regex as re	1✔
6
7	from txtai.pipeline import Labels, RAG, Similarity, Tokenizer	1✔
8
9	from ..index import Index	1✔
10	from ..query import Query	1✔
11
12	from .column import Column	1✔
13
14
15	class Report:	1✔
16	"""
17	Methods to build reports from a series of queries
18	"""
19
20	def __init__(self, embeddings, db, options):	1✔
21	"""
22	Creates a new report.
23
24	Args:
25	embeddings: embeddings index
26	db: database connection
27	options: report options
28	"""
29
30	# Store references to embeddings index and open database cursor
31	self.embeddings = embeddings	1✔
32	self.cur = db.cursor()	1✔
33
34	# Report options
35	self.options = options	1✔
36
37	# Column names
38	self.names = []	1✔
39
40	# Create similarity and labels pipeline, if necessary
41	self.similarity = Similarity(options["similarity"]) if "similarity" in options else None	1✔
42	self.labels = Labels(model=self.similarity) if self.similarity else None	1✔
43
44	# Create RAG pipeline arguments without report options
45	args = {x: options[x] for x in options if x not in ["topn", "render", "path", "qa", "indir", "threshold"]}	1✔
46
47	# Retrieval Augmented Generation (RAG) pipeline for calculated fields
48	self.rag = RAG(	1✔
49	self.similarity if self.similarity else self.embeddings,
50	args["llm"] if args.get("llm") else args["qa"] if args.get("qa") else "NeuML/bert-small-cord19qa",
51	**args,
52	)
53
54	def build(self, queries, options, output):	1✔
55	"""
56	Builds a report using a list of input queries
57
58	Args:
59	queries: queries to execute
60	options: report options
61	output: output I/O object
62	"""
63
64	# Default to 50 documents if not specified
65	topn = options.get("topn", 50)	1✔
66
67	for name, config in queries:	1✔
68	query = config["query"]	1✔
69	columns = config["columns"]	1✔
70
71	# Write query string
72	self.query(output, name, query)	1✔
73
74	# Write separator
75	self.separator(output)	1✔
76
77	# Query for best matches
78	results = Query.search(self.embeddings, self.cur, query, topn, options.get("threshold"))	1✔
79
80	# Generate highlights section
81	self.section(output, "Highlights")	1✔
82
83	# Generate highlights
84	self.highlights(output, results, int(topn / 10))	1✔
85
86	# Separator between highlights and articles
87	self.separator(output)	1✔
88
89	# Generate articles section
90	self.section(output, "Articles")	1✔
91
92	# Generate table headers
93	self.headers([column["name"] for column in columns], output)	1✔
94
95	# Generate table rows
96	self.articles(output, topn, (name, query, columns), results)	1✔
97
98	# Write section separator
99	self.separator(output)	1✔
100
101	def highlights(self, output, results, topn):	1✔
102	"""
103	Builds a highlights section.
104
105	Args:
106	output: output file
107	results: search results
108	topn: number of results to return
109	"""
110
111	# Extract top sections as highlights
112	for highlight in Query.highlights(results, topn):	1✔
113	# Get matching article
114	uid = [article for _, _, article, text in results if text == highlight][0]	1✔
115	self.cur.execute("SELECT Authors, Reference FROM articles WHERE id = ?", [uid])	1✔
116	article = self.cur.fetchone()	1✔
117
118	# Write out highlight row
119	self.highlight(output, article, highlight)	1✔
120
121	def articles(self, output, topn, metadata, results):	1✔
122	"""
123	Builds an articles section.
124
125	Args:
126	output: output file
127	topn: number of documents to return
128	metadata: query metadata
129	results: search results
130	"""
131
132	# Unpack metadata
133	_, query, _ = metadata	1✔
134
135	# Retrieve list of documents
136	documents = Query.all(self.cur) if query == "*" else Query.documents(results, topn)	1✔
137
138	# Collect matching rows
139	rows = []	1✔
140
141	for x, uid in enumerate(documents):	1✔
142	# Get article metadata
143	self.cur.execute(	1✔
144	"SELECT Published, Title, Reference, Publication, Source, Entry, Id FROM articles WHERE id = ?",
145	[uid],
146	)
147	article = self.cur.fetchone()	1✔
148
149	if x and x % 100 == 0:	1✔
150	print(f"Processed {x} documents", end="\r")	1✔
151
152	# Calculate derived fields
153	calculated = self.calculate(uid, metadata)	1✔
154
155	# Builds a row for article
156	rows.append(self.buildRow(article, documents[uid], calculated))	1✔
157
158	# Print report by published desc
159	for row in sorted(rows, key=lambda x: x["Date"], reverse=True):	1✔
160	# Convert row dict to list
161	row = [row[column] for column in self.names]	1✔
162
163	# Write out row
164	self.writeRow(output, row)	1✔
165
166	def calculate(self, uid, metadata):	1✔
167	"""
168	Builds a dict of calculated fields for a given document. This method calculates
169	constant field columns and derived query columns. Derived query columns run through
170	an embedding search and either run an additional QA query to extract a value or
171	use the top n embedding search matches.
172
173	Args:
174	uid: article id
175	metadata: query metadata
176
177	Returns:
178	{name: value} containing derived column values
179	"""
180
181	# Parse column parameters
182	fields, params = self.params(metadata)	1✔
183
184	# Different type of calculations
185	# 1. Similarity query
186	# 2. Extractor query (similarity + question)
187	# 3. Question-answering on other field
188	queries, extractions, questions = [], [], []	1✔
189
190	# Retrieve indexed document text for article
191	sections = self.sections(uid)	1✔
192	texts = [text for _, text in sections]	1✔
193
194	for name, query, question, snippet, _, _, matches, _ in params:	1✔
195	if query.startswith("$"):	1✔
196	questions.append((name, query.replace("$", ""), question, snippet))	×
197	elif matches:	1✔
198	queries.append((name, query, matches))	1✔
199	else:
200	extractions.append((name, query, question, snippet))	1✔
201
202	# Run all extractor queries against document text
203	results = self.rag.query([query for _, query, _ in queries], texts)	1✔
204
205	# Only execute embeddings queries for columns with matches set
206	for x, (name, query, matches) in enumerate(queries):	1✔
207	if results[x]:	1✔
208	# Get topn text matches
209	topn = [text for _, text, _ in results[x]][:matches]	1✔
210
211	# Join results into String and return
212	value = [self.resolve(params, sections, uid, name, value) for value in topn]	1✔
213	fields[name] = "\n\n".join(value) if value else ""	1✔
214	else:
215	fields[name] = ""	×
216
217	# Add extraction fields
218	if extractions:	1✔
219	for name, value in self.rag(extractions, texts, **self.options.get("params", {})):	1✔
220	# Resolves the full value based on column parameters
221	fields[name] = self.resolve(params, sections, uid, name, value) if value else ""	1✔
222
223	# Add question fields
224	if questions:	1✔
225	for name, value in self.rag(questions, texts, **self.options.get("params", {})):	×
226	# Resolves the full value based on column parameters
227	fields[name] = self.resolve(params, sections, uid, name, value) if value else ""	×
228
229	return fields	1✔
230
231	def params(self, metadata):	1✔
232	"""
233	Process and prepare parameters using input metadata.
234
235	Args:
236	metadata: query metadata
237
238	Returns:
239	fields, params - constant field values, query parameters for query columns
240	"""
241
242	# Derived field values
243	fields = {}	1✔
244
245	# Query column parameters
246	params = []	1✔
247
248	# Unpack metadata
249	_, _, columns = metadata	1✔
250
251	for column in columns:	1✔
252	# Constant column
253	if "constant" in column:	1✔
254	fields[column["name"]] = column["constant"]	×
255	# Question-answer column
256	elif "query" in column:	1✔
257	# Query variable substitutions
258	query = self.variables(column["query"], metadata)	1✔
259	question = self.variables(column["question"], metadata) if "question" in column else query	1✔
260
261	# Additional context parameters
262	section = column.get("section", False)	1✔
263	surround = column.get("surround", 0)	1✔
264	matches = column.get("matches", 0)	1✔
265	dtype = column.get("dtype")	1✔
266	snippet = column.get("snippet", False)	1✔
267	snippet = True if section or surround else snippet	1✔
268
269	params.append(	1✔
270	(
271	column["name"],
272	query,
273	question,
274	snippet,
275	section,
276	surround,
277	matches,
278	dtype,
279	)
280	)
281
282	return fields, params	1✔
283
284	def variables(self, value, metadata):	1✔
285	"""
286	Runs variable substitution for value.
287
288	Args:
289	value: input value
290	metadata: query metadata
291
292	Returns:
293	value with variable substitution
294	"""
295
296	name, query, _ = metadata	1✔
297
298	# Cleanup name for queries
299	name = name.replace("_", "").lower()	1✔
300	query = query.lower()	1✔
301
302	if value:	1✔
303	value = value.replace("$NAME", name).replace("$QUERY", query)	1✔
304
305	return value	1✔
306
307	def sections(self, uid):	1✔
308	"""
309	Retrieves all sections as list for article with given uid.
310
311	Args:
312	uid: article id
313
314	Returns:
315	list of section text elements
316	"""
317
318	# Retrieve indexed document text for article
319	self.cur.execute(Index.SECTION_QUERY + " WHERE article = ? ORDER BY id", [uid])	1✔
320
321	# Get list of document text sections
322	sections = []	1✔
323	for sid, name, text in self.cur.fetchall():	1✔
324	if not self.embeddings.isweighted() or not name or not re.search(Index.SECTION_FILTER, name.lower()) or self.options.get("allsections"):	1✔
325	# Check that section has at least 1 token
326	if Tokenizer.tokenize(text):	1✔
327	sections.append((sid, text))	1✔
328
329	return sections	1✔
330
331	def resolve(self, params, sections, uid, name, value):	1✔
332	"""
333	Fully resolves a value from an extractor call.
334
335	- If section=True, this method pull the full section text
336	- If surround is specified, this method will pull the surrounding text
337	- Otherwise, the original value is returned
338
339	Args:
340	params: query parameters
341	sections: section text
342	uid: article id
343	name: column name
344	value: initial query value after running through extractor process
345
346	Returns:
347	resolved value
348	"""
349
350	# Get all column parameters
351	index = [params.index(x) for x in params if x[0] == name][0]	1✔
352	_, _, _, _, section, surround, _, dtype = params[index]	1✔
353
354	if value:	1✔
355	# Find matching section
356	sid = [sid for sid, text in sections if value in text]	1✔
357
358	if sid:	1✔
359	sid = sid[0]	1✔
360
361	if section:	1✔
362	# Get full text for matching subsection
363	value = self.subsection(uid, sid)	1✔
364	elif surround:	1✔
365	value = self.surround(uid, sid, surround)	1✔
366
367	# Column dtype formatting
368	if dtype == "int":	1✔
369	value = Column.integer(value)	1✔
370	elif isinstance(dtype, list):	1✔
371	value = Column.categorical(self.labels, value, dtype)	1✔
372	elif dtype in ["days", "weeks", "months", "years"]:	1✔
373	value = Column.duration(value, dtype)	1✔
374
375	return value	1✔
376
377	def subsection(self, uid, sid):	1✔
378	"""
379	Extracts all subsection text for columns with section=True.
380
381	Args:
382	uid: article id
383	sid: section id
384
385	Returns:
386	full text for matching section
387	"""
388
389	self.cur.execute(	1✔
390	"SELECT Text FROM sections WHERE article = ? AND name = (SELECT name FROM sections WHERE id = ?)",
391	[uid, sid],
392	)
393	return " ".join([x[0] for x in self.cur.fetchall()])	1✔
394
395	def surround(self, uid, sid, size):	1✔
396	"""
397	Extracts surrounding text for section with specified id.
398
399	Args:
400	uid: article id
401	sid: section id
402	size: number of surrounding lines to extract from each side
403
404	Returns:
405	matching text with surrounding context
406	"""
407
408	self.cur.execute(	1✔
409	"SELECT Text FROM sections WHERE article = ? AND id in (SELECT id FROM sections WHERE id >= ? AND id <= ?) AND "
410	+ "name = (SELECT name FROM sections WHERE id = ?)",
411	[uid, sid - size, sid + size, sid],
412	)
413
414	return " ".join([x[0] for x in self.cur.fetchall()])	1✔
415
416	def cleanup(self, outfile):	1✔
417	"""
418	Allow freeing or cleaning up resources.
419
420	Args:
421	outfile: output file path
422	"""
423
424	def query(self, output, task, query):	1✔
425	"""
426	Writes query.
427
428	Args:
429	output: output file
430	task: task name
431	query: query string
432	"""
433
434	def section(self, output, name):	1✔
435	"""
436	Writes a section name
437
438	Args:
439	output: output file
440	name: section name
441	"""
442
443	def highlight(self, output, article, highlight):	1✔
444	"""
445	Writes a highlight row
446
447	Args:
448	output: output file
449	article: article reference
450	highlight: highlight text
451	"""
452
453	def headers(self, columns, output):	1✔
454	"""
455	Writes table headers.
456
457	Args:
458	columns: column names
459	output: output file
460	"""
461
462	def buildRow(self, article, sections, calculated):	1✔
463	"""
464	Converts a document to a table row.
465
466	Args:
467	article: article
468	sections: text sections for article
469	calculated: calculated fields
470	"""
471
472	def writeRow(self, output, row):	1✔
473	"""
474	Writes a table row.
475
476	Args:
477	output: output file
478	row: output row
479	"""
480
481	def separator(self, output):	1✔
482	"""
483	Writes a separator between sections
484	"""

neuml / paperai / 29327663402

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous