pipeline.py

import json
import re
from pathlib import Path

import yaml

from .evidence import LocalEvidence, NONE
from .openai_client import OpenAIChat, image_part


def _task_prompt(root, task):
    prompt = root[task]
    return prompt


def _ordered(claims):
    return sorted(
        claims,
        key=lambda item: (
            int(item.get("sentence_id", 10**9)),
            int(str(item.get("claim_id", "claim999999")).replace("claim", "")),
        ),
    )


class MedDetectPipeline:
    """Compact end-to-end implementation of full detection and iterative repair."""

    def __init__(self, config_path):
        self.config_path = Path(config_path).resolve()
        self.config = yaml.safe_load(self.config_path.read_text(encoding="utf-8"))
        self._resolve_paths()
        model_config = self.config["openai"]
        self.chat = OpenAIChat(
            model=model_config["model"],
            temperature=model_config.get("temperature", 0.0),
            max_tokens=model_config.get("max_tokens", 8192),
        )
        self.prompts = {
            name: yaml.safe_load(Path(path).read_text(encoding="utf-8"))
            for name, path in self.config["prompts"].items()
        }
        self.local_evidence = LocalEvidence(self.config, self.chat)

    def _resolve_paths(self):
        base = self.config_path.parent
        for key, value in self.config["prompts"].items():
            path = Path(value)
            self.config["prompts"][key] = str(
                path if path.is_absolute() else (base / path).resolve()
            )

    def run(self, sample):
        task = sample.get("task") or self.config["general"]["task"]
        response = str(sample["response"])
        segments = sample.get("segments") or self.split_response(response, task)
        image_path = sample.get("image_path") or None
        if image_path and not Path(image_path).is_file():
            image_path = None
        query_claims = self.extract_query_claims(sample.get("query"), task)
        initial = self.detect(
            response=response,
            segments=segments,
            query_claims=query_claims,
            image_path=image_path,
            image_modality=sample.get("image_modality"),
            task=task,
            stop_on_first=False,
        )
        detection = {
            "num_claims": len(initial["claims"]),
            "num_hallucinations": len(initial["errors"]),
            "first_error": initial["errors"][0] if initial["errors"] else None,
            "hallucinations": initial["errors"],
        }
        result = {
            "uid": sample.get("uid"),
            "mode": self.config["pipeline"]["mode"],
            "detection": detection,
        }
        if self.config["pipeline"]["mode"] == "iterative_correction":
            result["correction"] = self.correct(
                response=response,
                first_error=detection["first_error"],
                initial=initial,
                query=sample.get("query"),
                query_claims=query_claims,
                image_path=image_path,
                image_modality=sample.get("image_modality"),
                task=task,
            )
        return result

    def detect(
        self,
        response,
        segments,
        query_claims,
        image_path,
        image_modality,
        task,
        *,
        stop_on_first,
        start_sentence_id=1,
        trusted_prefix=None,
        claims=None,
    ):
        claims = claims or self.extract_response_claims(segments, task, image_modality)
        dependencies = self.link_dependencies(
            query_claims, claims, target_start_sentence_id=start_sentence_id
        )
        prefix_decisions = {}
        errors = []
        checked = []
        trusted_prefix = trusted_prefix or []

        for claim in _ordered(claims):
            if int(claim["sentence_id"]) < int(start_sentence_id):
                continue
            local = self.local_evidence.ground(claim, image_path, task)
            packet = self.build_evidence_packet(
                claim=claim,
                query_claims=query_claims,
                dependencies=dependencies,
                prefix_decisions=prefix_decisions,
                trusted_prefix=trusted_prefix,
                local=local,
                image_modality=image_modality,
                task=task,
            )
            decision = self.verify(claim, packet, image_path, task)
            checked.append(claim["claim_id"])
            if decision["verdict"] == "hallucination":
                errors.append(self.error_result(claim, decision))
                if stop_on_first:
                    break
            prefix_decisions[claim["claim_id"]] = {
                "claim": claim["claim"],
                "verdict": decision["verdict"],
                "reason": decision.get("reason", ""),
            }
        return {
            "response": response,
            "claims": claims,
            "dependencies": dependencies,
            "errors": _ordered(errors),
            "checked": checked,
        }

    def correct(
        self,
        response,
        first_error,
        initial,
        query,
        query_claims,
        image_path,
        image_modality,
        task,
    ):
        if not first_error:
            return {
                "corrected_response": response,
                "rounds": 0,
                "remaining_first_error": None,
            }

        current_response = response
        current_error = first_error
        current_claims = initial["claims"]
        current_dependencies = initial["dependencies"]
        rounds = 0
        for _ in range(self.config["pipeline"].get("max_repair_rounds", 4)):
            repaired_response = self.revise(
                response=current_response,
                first_error=current_error,
                claims=current_claims,
                dependencies=current_dependencies,
                query=query,
                image_path=image_path,
                task=task,
            )
            rounds += 1
            if repaired_response == current_response:
                break

            start = int(current_error["sentence_id"])
            if self.config["pipeline"].get("recheck_after_repair_start") == "next_sentence":
                start += 1
            repaired_segments = self.split_response(repaired_response, task)
            repaired_claims = self.extract_response_claims(
                repaired_segments, task, image_modality
            )
            trusted_prefix = [
                claim
                for claim in repaired_claims
                if int(claim["sentence_id"]) < int(start)
            ]
            post = self.detect(
                response=repaired_response,
                segments=repaired_segments,
                query_claims=query_claims,
                image_path=image_path,
                image_modality=image_modality,
                task=task,
                stop_on_first=True,
                start_sentence_id=start,
                trusted_prefix=trusted_prefix,
                claims=repaired_claims,
            )
            current_response = repaired_response
            current_claims = post["claims"]
            current_dependencies = post["dependencies"]
            current_error = post["errors"][0] if post["errors"] else None
            if current_error is None:
                break

        return {
            "corrected_response": current_response,
            "rounds": rounds,
            "remaining_first_error": current_error,
        }

    def extract_response_claims(self, segments, task, image_modality=None):
        text = segments if isinstance(segments, list) else str(segments)
        serialized = json.dumps(text, ensure_ascii=False) if isinstance(text, list) else text
        prompt = self._claim_stage_prompt(task, "extract")
        user = (
            prompt["user"]
            .replace("{text}", serialized)
            .replace("{segment_count}", str(len(text) if isinstance(text, list) else 1))
            .replace("{attached_image_modality}", str(image_modality or "unknown/not provided"))
        )
        raw = self.chat.complete(
            [
                {"role": "system", "content": prompt["system"]},
                {"role": "user", "content": user},
            ]
        )
        response = json.loads(raw)
        claims = []
        for segment_index, segment in enumerate(response):
            original_segment = text[segment_index] if isinstance(text, list) else segment["segment"]
            for item in segment.get("claims", []):
                claim = {
                    "claim_id": f"claim{len(claims) + 1}",
                    "claim": " ".join(str(item["claim"]).split()),
                    "sentence_id": segment_index + 1 if isinstance(text, list) else int(segment["sentence_id"]),
                    "segment": str(original_segment),
                    "section": item.get("section", segment.get("section", "")),
                    "source": "response",
                }
                claims.append(claim)
        if claims:
            type_prompt = self._claim_stage_prompt(task, "type")
            fixed_claims = [
                {"claim_id": claim["claim_id"], "claim": claim["claim"]}
                for claim in claims
            ]
            type_user = (
                type_prompt["user"]
                .replace("{claims}", json.dumps(fixed_claims, ensure_ascii=False))
                .replace(
                    "{attached_image_modality}",
                    str(image_modality or "unknown/not provided"),
                )
            )
            raw_types = self.chat.complete(
                [
                    {"role": "system", "content": type_prompt["system"]},
                    {"role": "user", "content": type_user},
                ]
            )
            types = {
                item["claim_id"]: str(item["claim_type"]).lower()
                for item in json.loads(raw_types)
            }
            for claim in claims:
                claim["claim_type"] = types[claim["claim_id"]]
        return claims

    def _claim_stage_prompt(self, task, stage):
        prompt_name = self.config["claim_generation"]["prompt"]
        prompt_root = self.prompts[prompt_name]
        return prompt_root[task][stage]

    def extract_query_claims(self, query, task):
        if not query:
            return []
        prompt = _task_prompt(self.prompts["query_claim_generate"], task)
        text = json.dumps(query, ensure_ascii=False) if isinstance(query, (dict, list)) else str(query)
        raw = self.chat.complete(
            [
                {"role": "system", "content": prompt["system"]},
                {"role": "user", "content": prompt["user"].replace("{text}", text)},
            ]
        )
        response = json.loads(raw)
        if isinstance(response, dict):
            response = [response]
        claims = []
        for segment_index, segment in enumerate(response, start=1):
            for item in segment.get("claims", []):
                claims.append(
                    {
                        "claim_id": f"q{len(claims) + 1}",
                        "claim": " ".join(str(item["claim"]).split()),
                        "sentence_id": segment_index,
                        "source": "query",
                    }
                )
        return claims

    def link_dependencies(self, query_claims, response_claims, target_start_sentence_id=1):
        records = {}
        prompt_root = self.prompts["dependency_linking"]
        ordered_claims = _ordered(response_claims)
        for index, claim in enumerate(ordered_claims):
            if int(claim["sentence_id"]) < int(target_start_sentence_id):
                continue
            if claim["claim_type"] == "knowledge":
                records[claim["claim_id"]] = self._dependency_record(claim, [], [])
                continue
            query_ids = self._link_route(claim, query_claims, "query", prompt_root) if query_claims else []
            earlier = ordered_claims[:index]
            response_ids = self._link_route(claim, earlier, "response", prompt_root) if earlier else []
            records[claim["claim_id"]] = self._dependency_record(
                claim, query_ids, response_ids
            )
        for claim in ordered_claims:
            if claim["claim_id"] not in records:
                continue
            record = records[claim["claim_id"]]
            record["resolved_dependency_ids"] = self._resolve_dependencies(
                record["direct_dependency_ids"], records
            )
        return records

    def _link_route(self, claim, candidates, route, prompt_root):
        prompt = prompt_root[route]
        candidate_key = "query_claims" if route == "query" else "earlier_response_claims"
        payload = {
            "current_claim": self._compact_claim(claim),
            candidate_key: [self._compact_claim(item) for item in candidates],
        }
        raw = self.chat.complete(
            [
                {"role": "system", "content": prompt["system"]},
                {
                    "role": "user",
                    "content": prompt["user"].replace(
                        "{input_json}", json.dumps(payload, ensure_ascii=False, indent=2)
                    ),
                },
            ],
            max_tokens=self.config["pipeline"].get("dependency_max_tokens", 1024),
        )
        parsed = json.loads(raw)
        allowed = {item["claim_id"] for item in candidates}
        selected = []
        for item in parsed.get("dependencies", []):
            dependency_id = item.get("claim_id") if isinstance(item, dict) else item
            if dependency_id in allowed and dependency_id not in selected:
                selected.append(dependency_id)
        return selected

    @staticmethod
    def _dependency_record(claim, query_ids, response_ids):
        return {
            "claim_id": claim["claim_id"],
            "query_dependency_ids": query_ids,
            "response_dependency_ids": response_ids,
            "direct_dependency_ids": query_ids + response_ids,
        }

    def _resolve_dependencies(self, direct_ids, records):
        resolved = []

        def visit(claim_id):
            for parent_id in records.get(claim_id, {}).get("direct_dependency_ids", []):
                visit(parent_id)
            if claim_id not in resolved:
                resolved.append(claim_id)

        for direct_id in direct_ids:
            visit(direct_id)
        return resolved

    def build_evidence_packet(
        self,
        claim,
        query_claims,
        dependencies,
        prefix_decisions,
        trusted_prefix,
        local,
        image_modality,
        task,
    ):
        claim_type = claim["claim_type"]
        packet = {
            "current_claim": {
                "claim_id": claim["claim_id"],
                "claim": claim["claim"],
            }
        }
        if claim_type == "knowledge":
            if local.get("medical_knowledge"):
                packet["medical_knowledge"] = local["medical_knowledge"]
            return packet

        dependency = dependencies[claim["claim_id"]]
        dependency_ids = (
            dependency["resolved_dependency_ids"]
            if task == "report"
            else dependency["direct_dependency_ids"]
        )
        query_ids = {item for item in dependency_ids if str(item).startswith("q")}
        response_ids = {item for item in dependency_ids if str(item).startswith("claim")}
        case_context = [
            {"claim": item["claim"]}
            for item in query_claims
            if item["claim_id"] in query_ids
        ]
        trusted_by_id = {item["claim_id"]: item for item in trusted_prefix}
        linked_trusted = [trusted_by_id[item] for item in response_ids if item in trusted_by_id]
        prefix_trace = [prefix_decisions[item] for item in response_ids if item in prefix_decisions]

        if task == "report" and linked_trusted:
            prefix_trace = [
                {
                    "claim": item["claim"],
                    "verdict": "non-hallucination",
                    "reason": "Trusted repaired response before the recheck cursor.",
                }
                for item in linked_trusted
            ] + prefix_trace
        elif task == "diagnose":
            case_context.extend({"claim": item["claim"]} for item in linked_trusted)

        if claim_type in {"context", "reasoning"}:
            if task == "diagnose" and case_context:
                packet["case_context"] = case_context
            if prefix_trace:
                packet["prefix_trace"] = prefix_trace
            if claim_type == "reasoning" and local.get("medical_knowledge"):
                packet["medical_knowledge"] = local["medical_knowledge"]
            return packet

        if task == "diagnose" and case_context:
            packet["case_context"] = case_context
        if prefix_trace:
            packet["prefix_trace"] = prefix_trace
        if image_modality:
            packet["attached_image_scope"] = {"modality": image_modality}
        visual = {key: value for key, value in local.items() if value not in (None, [], {}, NONE)}
        if visual:
            packet["local_visual_evidence"] = visual
        return packet

    def verify(self, claim, packet, image_path, task):
        prompt_group = _task_prompt(self.prompts["unified_verify"], task)
        prompt = prompt_group[claim["claim_type"]]
        use_image = bool(prompt.get("use_images")) and bool(image_path)
        content = [{"type": "text", "text": prompt["user"]}]
        example_text = "\n\n".join(
            f'Example {index}:\n{prompt[key]}'
            for index, key in enumerate(("example1", "example2"), start=1)
            if prompt.get(key)
        )
        if example_text:
            content.append({"type": "text", "text": "BEGIN FEW-SHOT EXAMPLE."})
            if use_image:
                for example in prompt.get("example_images", []):
                    example_path = (Path(self.config["prompts"]["unified_verify"]).parent / example).resolve()
                    content.append(image_part(example_path))
            content.append({"type": "text", "text": example_text + "\nEND FEW-SHOT EXAMPLE."})
        content.append(
            {
                "type": "text",
                "text": "BEGIN CURRENT CASE.\nFinal input JSON:\n" + json.dumps(packet, ensure_ascii=False, indent=2),
            }
        )
        if use_image:
            content.append(image_part(image_path))
        content.append({"type": "text", "text": "END CURRENT CASE."})
        raw = self.chat.complete(
            [
                {"role": "system", "content": prompt["system"]},
                {"role": "user", "content": content},
            ]
        )
        return json.loads(raw)

    def revise(self, response, first_error, claims, dependencies, query, image_path, task):
        prompt = _task_prompt(self.prompts["first_error_revision"], task)
        compact_error = {
            key: first_error.get(key)
            for key in ("claim_id", "claim", "sentence_id", "segment", "reason")
        }
        payload = {
            "query": query or "",
            "response": response,
            "first_error": compact_error,
        } if task == "diagnose" else {
            "report": response,
            "first_error": compact_error,
        }
        cascade = self._cascade_candidates(first_error, claims, dependencies)
        if cascade:
            payload["cascade_candidates"] = cascade
        content = [
            {
                "type": "text",
                "text": prompt["user"].replace("{payload}", json.dumps(payload, ensure_ascii=False, indent=2)),
            }
        ]
        if image_path and Path(image_path).is_file():
            content.append(image_part(image_path))
        raw = self.chat.complete(
            [
                {"role": "system", "content": prompt["system"]},
                {"role": "user", "content": content},
            ]
        )
        return str(json.loads(raw)["repaired_response"])

    @staticmethod
    def _cascade_candidates(first_error, claims, dependencies):
        root_id = first_error["claim_id"]
        root_sentence = int(first_error["sentence_id"])
        grouped = {}
        for claim in claims:
            if int(claim["sentence_id"]) <= root_sentence:
                continue
            if root_id not in dependencies[claim["claim_id"]]["response_dependency_ids"]:
                continue
            group = grouped.setdefault(
                int(claim["sentence_id"]),
                {"sentence_id": int(claim["sentence_id"]), "segment": claim["segment"], "linked_claims": []},
            )
            group["linked_claims"].append(
                {"claim_id": claim["claim_id"], "claim": claim["claim"]}
            )
        return [grouped[key] for key in sorted(grouped)]

    @staticmethod
    def error_result(claim, decision):
        return {
            "claim_id": claim["claim_id"],
            "claim": claim["claim"],
            "claim_type": claim["claim_type"],
            "sentence_id": claim["sentence_id"],
            "segment": claim["segment"],
            "verdict": decision["verdict"],
            "error_source": decision.get("error_source"),
            "error_type": decision.get("error_type"),
            "repair_action": decision.get("repair_action"),
            "reason": decision.get("reason", ""),
        }

    @staticmethod
    def _compact_claim(claim):
        compact = {
            key: claim[key]
            for key in ("claim_id", "claim", "claim_type", "sentence_id")
            if key in claim
        }
        compact["source"] = claim.get("source", "response")
        return compact

    @staticmethod
    def split_response(text, task):
        text = str(text or "").replace("\r\n", "\n").strip()
        if not text:
            return []
        if task == "report":
            sections = re.findall(
                r"(?:^|\n)\s*([^:\n]+?)\s*:\s*(.*?)(?=\n\s*[^:\n]+?\s*:|\Z)",
                text,
                re.DOTALL,
            )
            if sections:
                rows = []
                for section, body in sections:
                    for sentence in re.split(r"(?<=[.!?])\s+", body.strip()):
                        if sentence.strip():
                            rows.append(f"{section.strip()}: {sentence.strip()}")
                return rows
        blocks = [line.strip() for line in text.splitlines() if line.strip()]
        return [
            sentence.strip()
            for block in blocks
            for sentence in re.split(r"(?<=[.!?])\s+", block)
            if sentence.strip()
        ]