rag_query_generator.py

import re

from chemical_formatter import to_latex_subscript, to_unicode_subscript
from tracernet.services.pathway_evidence import parse_path, path_components


class TracePathQueryGenerator:
    @staticmethod
    def _to_nougat_latex(formula):
        latex = to_latex_subscript(formula)
        if not latex:
            return latex
        return re.sub(r"_\{(\d+)\}", r"\\({}_\1\\)", latex)

    @staticmethod
    def _ordered_unique(items):
        seen = set()
        out = []
        for item in items:
            query = str(item or "").strip()
            if not query or query in seen:
                continue
            seen.add(query)
            out.append(query)
        return out

    @staticmethod
    def parse_trace_report(trace_report):
        unique_steps = set()
        parsed_steps = []
        for start_material, paths in trace_report.items():
            for path_str in paths:
                for edge in parse_path(path_str):
                    step_id = f"{edge.reactant}|{edge.condition}|{edge.product}"
                    if step_id in unique_steps:
                        continue
                    unique_steps.add(step_id)
                    parsed_steps.append(
                        {
                            "source": edge.reactant,
                            "target": edge.product,
                            "condition": edge.condition,
                            "root_material": start_material,
                        }
                    )
        return parsed_steps

    def generate_queries(self, trace_report):
        queries = []
        for step in self.parse_trace_report(trace_report):
            src = step["source"]
            tgt = step["target"]
            cond = step["condition"]
            src_unicode = to_unicode_subscript(src)
            tgt_unicode = to_unicode_subscript(tgt)
            src_latex = to_latex_subscript(src)
            tgt_latex = to_latex_subscript(tgt)
            src_nougat = self._to_nougat_latex(src)
            tgt_nougat = self._to_nougat_latex(tgt)
            readable_cond = cond.replace("+", " and ")
            queries.append(
                f"Chemical reaction converting {src_unicode} to {tgt_unicode} "
                f"under {readable_cond}"
            )
            if src_latex != src or tgt_latex != tgt:
                queries.append(
                    f"Chemical reaction converting {src_latex} to {tgt_latex} "
                    f"under {readable_cond}"
                )
                queries.append(f"Formation mechanism of {tgt_latex} from {src_latex}")
            if src_nougat != src or tgt_nougat != tgt:
                queries.append(
                    f"Chemical reaction converting {src_nougat} to {tgt_nougat} "
                    f"under {readable_cond}"
                )
                queries.append(f"Formation mechanism of {tgt_nougat} from {src_nougat}")
            queries.append(
                f"Influence of {readable_cond} on the degradation of "
                f"{src_unicode} in murals"
            )
            queries.append(
                f"Evidence for pigment degradation pathway from {src_unicode} "
                f"to {tgt_unicode} under {readable_cond}"
            )

        for root, paths in (trace_report or {}).items():
            root_unicode = to_unicode_subscript(root)
            root_latex = to_latex_subscript(root)
            queries.append(f"Degradation pathways of {root_unicode}")
            if root_latex != root:
                queries.append(f"Degradation pathways of {root_latex}")
            for path_str in paths:
                species, conditions = path_components(path_str)
                if species:
                    start = to_unicode_subscript(species[0])
                    end = to_unicode_subscript(species[-1])
                    condition_text = (
                        " then ".join(
                            condition.replace("+", " and ")
                            for condition in conditions
                        )
                        if conditions
                        else "the reported conditions"
                    )
                    queries.append(
                        f"Literature evidence for full pathway from {start} "
                        f"to {end} under {condition_text}"
                    )
                queries.append(f"Mural degradation pathway: {path_str}")
        return self._ordered_unique(queries)