QA_MA.py

from __future__ import annotations
import asyncio, json, logging, math
from typing import Any, Dict, List, Literal, Optional, Union

import pandas as pd
from openai import AsyncOpenAI
from pydantic import BaseModel, Field, create_model

from utils import count_llama_tokens, safe_json_load        
from rag_retrieve import create_documents, hybrid_query      


LLAMA3_70B_MAX_TOKENS = 22000
logger = logging.getLogger(__name__)


class Response(BaseModel):
    reasoning: str = Field(..., description="Step-by-step reasoning leading to your choice.")
    choice:    Literal["A","B","C","D","E"] = Field(..., description="Your choice to the question.")

class LLMAgent:
    def __init__(self,
        system_prompt: str,
        model_name: str = "meta-llama/Llama-3.3-70B-Instruct",
        client          = AsyncOpenAI(base_url="http://localhost:8000/v1", api_key="dummy"),
        max_tokens: int = LLAMA3_70B_MAX_TOKENS,
        summarization_threshold: float = 0.7):

        self.model_name  = model_name
        self.client      = client
        self.messages    = [{"role": "system", "content": system_prompt}]
        self.max_tokens  = max_tokens
        self.token_thres = int(max_tokens * summarization_threshold)

        logger.info(
            "[%s] init – token limit %d, summarize ≥%d",
            self.__class__.__name__, max_tokens, self.token_thres)


    async def _summarize_once(self, text: str, max_chars: Optional[int] = None) -> Optional[str]:
        length_rule = f"Do **not** exceed {max_chars} characters." if max_chars else "Do **not** exceed 1000 words."
        prompt = (
            "Summarize the following message concisely, preserving every key fact and reasoning step. "
            f"{length_rule}\n\n<<<MESSAGE_START>>>\n{text}\n<<<MESSAGE_END>>>"
        )
        try:
            resp = await self.client.chat.completions.create(
                model       = self.model_name,
                messages    = [
                    {"role": "system", "content": "You are a summarization assistant."},
                    {"role": "user",   "content": prompt}
                ],
                temperature = 0.1,
                max_tokens  = 1500
            )
            return resp.choices[0].message.content.strip()
        except Exception as e:
            logger.error("Summarization failed: %s", e)
            return None

    async def _summarize_history(self,
        inplace: bool = True,
        message: str = "",
        max_chars: Optional[int] = None
    ) -> Union[bool, str]:
      
        if inplace:
            if len(self.messages) < 3:
                logger.warning("Not enough messages to summarize.")
                return False
            tried: set[int] = set()
            while count_llama_tokens(self.messages) >= self.token_thres:
                idx_len = sorted(
                    ((i, count_llama_tokens([m]))
                     for i, m in enumerate(self.messages)
                     if m["role"] != "system" and i not in tried),
                    key=lambda t: t[1], reverse=True
                )
                if not idx_len:
                    logger.error("No further messages left to summarize.")
                    return False
                idx = idx_len[0][0]
                summary = await self._summarize_once(self.messages[idx]["content"], max_chars)
                if summary is None or count_llama_tokens([{"role": "assistant", "content": summary}]) >= count_llama_tokens([self.messages[idx]]):
                    tried.add(idx)
                    continue
                self.messages[idx]["content"] = f"[Summary] {summary}"
                logger.info("Summarized message #%d → total %d tokens", idx, count_llama_tokens(self.messages))
            return True

        if not message:
            logger.warning("No message supplied for external summarization.")
            return False
        summary = await self._summarize_once(message, max_chars)
        if summary is None:
            return False
        while count_llama_tokens(summary) >= self.token_thres or (max_chars and len(summary) > max_chars):
            summary = await self._summarize_once(summary, max_chars)
            if summary is None:
                return False
        return summary


    async def llm_call(self, user_prompt: str,
        temperature: float = 0.5,
        guided_: Optional[dict] = None
    ) -> str:
        if count_llama_tokens(self.messages + [{"role": "user", "content": user_prompt}]) > self.token_thres:
            ok = await self._summarize_history()
            if not ok and count_llama_tokens(self.messages) > self.max_tokens:
                raise RuntimeError("Context length still exceeds hard limit after summarization.")
        self.messages.append({"role": "user", "content": user_prompt})

        params: Dict[str, Any] = {
            "model":       self.model_name,
            "messages":    self.messages,
            "temperature": temperature,
        }
        if guided_:
            params["extra_body"] = guided_

        resp = await self.client.chat.completions.create(**params)
        return resp.choices[0].message.content


    def append_message(self, content: Any, role: str = "assistant") -> None:
        self.messages.append({"role": role, "content": str(content)})



class BaselineZS(LLMAgent):
    def __init__(self):
        super().__init__("You are an agent in a medical question-answering system.")
        self.schema = Response.model_json_schema()

    async def analyse_question(self, q_text: str) -> Optional[dict]:
        user_prompt = (
            "Read the following medical question and provide your reasoning, then your choice.\n"
            f"<<<QUESTION_START>>{q_text}<<<QUESTION_END>>>\n\n"
            "Please give your reasoning, then the choice (A, B, C, D, or E)."
        )
        raw = await self.llm_call(user_prompt, guided_={"guided_json": self.schema})
        return safe_json_load(raw)


class GenericAgent(LLMAgent):
    def __init__(self, agent_id: str, state: dict):
        self.agent_id = str(agent_id)
        self.state    = state
        self.state.setdefault("generic_agents", {})[self.agent_id] = {}
        self.round    = 0
        self.schema   = Response.model_json_schema()

        super().__init__("You are a collaborating agent in a medical question-answering multi‑agent system.")

    async def analyse_question(self, q_text: str) -> dict:
        self.round += 1
        user_prompt = (
            "Read the following medical question and provide your reasoning, then your choice.\n"
            f"<<<QUESTION_START>>{q_text}<<<QUESTION_END>>>\n\n"
            "Please give your reasoning, then the choice (A, B, C, D, or E)."
        )
        raw = await self.llm_call(user_prompt, guided_={"guided_json": self.schema})
        parsed = safe_json_load(raw)
        self.append_message(raw)
        self.state["generic_agents"][self.agent_id][f"round_{self.round}"] = parsed
        return parsed

    async def debate(self) -> dict:
        self.round += 1
        peers: Dict[str, Any] = {}
        for a_id, hist in self.state["generic_agents"].items():
            if a_id == self.agent_id:      
                continue
            if f"round_{self.round - 1}" in hist:
                peers[a_id] = hist[f"round_{self.round - 1}"]
        for p in [k for k in self.state if k.startswith("panel_")]:
            for role, info in self.state[p]["Collected Specialists"].items():
                if f"round_{self.round - 1}" in info["answer_history"]:
                    peers[role] = info["answer_history"][f"round_{self.round - 1}"]

        prompt = (
            "Here are your peers’ previous answers:\n"
            f"<<<PEERS_START>>>\n{json.dumps(peers, indent=2)}\n<<<PEERS_END>>>\n"
            "Review their reasoning against yours. Considering their input and your own "
            "analysis, decide whether to revise your answer.\n"
            "Return your updated reasoning and final choice (A, B, C, D, or E)."
        )
        raw    = await self.llm_call(prompt, guided_={"guided_json": self.schema})
        parsed = safe_json_load(raw)
        self.append_message(raw)
        self.state["generic_agents"][self.agent_id][f"round_{self.round}"] = parsed
        return parsed


class DynamicSpecialist(LLMAgent):
    def __init__(self, specialist: str, panel_id: int, state: dict):
        self.specialist   = specialist
        self.panel_id    = panel_id
        self.state       = state
        rec              = state[f"panel_{panel_id}"]["Collected Specialists"][specialist]
        self.expertise   = rec["expertise"]
        self.answer_hist = rec["answer_history"]
        self.round       = 0
        self.schema      = Response.model_json_schema()

        super().__init__(
            f"You are a {specialist} agent in a medical question-answering multi‑agent system.\n"
            f"Your expertise includes:\n{self.expertise}\n"
            f"Solve the question from the perspective of a {specialist}."
        )

    async def analyse_question(self, q_text: str) -> dict:
        self.round += 1
        user_prompt = (
            "Read the following medical question and provide your reasoning, then your choice.\n"
            f"<<<QUESTION_START>>{q_text}<<<QUESTION_END>>>\n\n"
            "Please give your reasoning, then the choice (A, B, C, D, or E)."
        )
        raw = await self.llm_call(user_prompt, guided_={"guided_json": self.schema})
        parsed = safe_json_load(raw)
        self.append_message(raw)
        self.answer_hist[f"round_{self.round}"] = parsed
        return parsed

    async def debate(self) -> dict:
        self.round += 1
        peers: Dict[str, Any] = {}

        # specialists in the same panel
        for role, info in self.state[f"panel_{self.panel_id}"]["Collected Specialists"].items():
            if role == self.specialist:
                continue
            tag = f"round_{self.round - 1}"
            if tag in info["answer_history"]:
                peers[role] = info["answer_history"][tag]

        # generic agents
        for gid, hist in self.state.get("generic_agents", {}).items():
            tag = f"round_{self.round - 1}"
            if tag in hist:
                peers[gid] = hist[tag]

        prompt = (
            "Here are your peers’ previous answers:\n"
            f"<<<PEERS_START>>>\n{json.dumps(peers, indent=2)}\n<<<PEERS_END>>>\n"
            "Review their reasoning against yours. Considering their input and your own "
            "analysis, decide whether to revise your answer.\n"
            "Return your updated reasoning and final choice (A, B, C, D, or E)."
        )
        raw    = await self.llm_call(prompt, guided_={"guided_json": self.schema})
        parsed = safe_json_load(raw)
        self.append_message(raw)
        self.answer_hist[f"round_{self.round}"] = parsed
        return parsed



class Manager(LLMAgent):
    def __init__(
        self,
        q_text: str,
        q_id: str,
        label: str,
        n_specialists: Union[int, Literal["auto"]] = "auto",
        n_generic: int = 0,
        static_specialists: Optional[List[str]] = None,
        consensus_thresh: float = 0.8,
        max_consensus_rounds: int = 3,
        max_assignment_attempts: int = 2,
        per_agent_summary_chars: int = 300,
    ):
        super().__init__(
            "You are the manager of a medical question-answering multi‑agent system. "
            "Your role is to coordinate sub‑agents to reach a final answer."
        )

        self.q_text         = q_text
        self.q_id           = q_id
        self.label          = label

        self.n_specialists  = n_specialists
        self.n_generic      = n_generic
        self.static_specs   = static_specialists or []

        self.cons_thresh    = consensus_thresh
        self.max_cons_round = max_consensus_rounds
        self.max_assign     = max_assignment_attempts
        self.per_agent_chars= per_agent_summary_chars

        self.assign_attempts = 0
        self.state: Dict[str, Any] = {
            "question": q_text,
            "q_id":     q_id,
            "label":    label,
            "generic_agents": {},
            "final":    {}
        }

    # ──────────────────────────────────────────────────────────────────────
    #  Specialist selection
    # ──────────────────────────────────────────────────────────────────────
    async def _assign_specialists(self) -> None:
        self.assign_attempts += 1
        pid = self.assign_attempts
        self.state[f"panel_{pid}"] = {
            "Initially Identified Specialties": [],
            "Collected Specialists": {}
        }

        # (1) decide which specialties we want
        specialties: List[str] = list(self.static_specs)
        additional_needed: Union[int, Literal["auto"]]
        if self.n_specialists == "auto":
            additional_needed = "auto"
        else:
            additional_needed = max(self.n_specialists - len(self.static_specs), 0)

        if additional_needed == 0 and not specialties:
            return                                          # nothing to add

        if additional_needed:
            ask_specialties = (
                f"Below is the **medical question** we need to solve:\n\n{self.q_text}\n\n"
                f"Already chosen specialties: {self.static_specs or '(none)'}.\n"
                "Please list any **additional medical specialties** whose expertise would help answer this question."
            )
            if additional_needed == "auto":
                class SpecialtyList(BaseModel):
                    specialties: List[str]
                raw = await self.llm_call(ask_specialties, guided_={"guided_json": SpecialtyList.model_json_schema()})
                specialties.extend(safe_json_load(raw)["specialties"])
            else:
                fld = {f"specialty_{i+1}": (str, ...) for i in range(additional_needed)}
                Req = create_model("RequestedSpecialties", **fld)
                raw = await self.llm_call(ask_specialties, guided_={"guided_json": Req.model_json_schema()})
                specialties.extend([safe_json_load(raw)[f"specialty_{i+1}"] for i in range(additional_needed)])

        self.state[f"panel_{pid}"]["Initially Identified Specialties"] = specialties

        # (2) flesh them out
        ask_panel = (
            f"We will run a multi-agent debate to answer one question:\n"
            f"{self.q_text}\n\n"
            "The required specialties are:\n"
            f"{specialties}\n"
            "For each listed specialty, return an object with keys:\n"
            "  • `specialist` – full job title\n"
            "  • `expertise`  – 1‑3 short phrases explaining the specialist’s expertise" 
        )
        class Specialist(BaseModel):
            specialist: str
            expertise:  List[str]
        fld  = {f"specialist_{i+1}": (Specialist, ...) for i in range(len(specialties))}
        PanelOut  = create_model("SpecialistPanel", **fld)
        raw  = await self.llm_call(ask_panel, guided_={"guided_json": PanelOut.model_json_schema()})
        panel_json = safe_json_load(raw)
        self.append_message(raw)                          

        for obj in panel_json.values():
            self.state[f"panel_{pid}"]["Collected Specialists"][obj["specialist"]] = {
                "expertise":       obj["expertise"],
                "answer_history":  {}
            }


    async def _panel_summary(self, pid: int) -> str:
        round_ids = []
        for info in self.state[f"panel_{pid}"]["Collected Specialists"].values():
            round_ids.extend(int(tag.split('_')[1]) for tag in info["answer_history"])
        for hist in self.state.get("generic_agents", {}).values():
            round_ids.extend(int(tag.split('_')[1]) for tag in hist)
        if not round_ids:
            return "(no valid answers to summarize)"
        last_round = max(round_ids)

        blobs: List[tuple[str,str,str]] = []
        for role, info in self.state[f"panel_{pid}"]["Collected Specialists"].items():
            tag = f"round_{last_round}"
            if tag in info["answer_history"]:
                ans = info["answer_history"][tag]
                blobs.append((role, ans["choice"], ans["reasoning"]))
        for gid, hist in self.state.get("generic_agents", {}).items():
            tag = f"round_{last_round}"
            if tag in hist:
                ans = hist[tag]
                blobs.append((gid, ans["choice"], ans["reasoning"]))

        async def _compress(text: str) -> str:
            out = await self._summarize_history(inplace=False, message=text, max_chars=self.per_agent_chars)
            return out if isinstance(out, str) else text.strip()

        summaries = await asyncio.gather(*(_compress(r) for _, _, r in blobs))
        return "\n".join(f"- **{name}** → {ch} | {short}" for (name, ch, _), short in zip(blobs, summaries))


    def _check_consensus(self, pid: int, round_id: int, total_agents: int) -> Optional[str]:
        count: Dict[str, int] = {}
        for role, info in self.state[f"panel_{pid}"]["Collected Specialists"].items():
            tag = f"round_{round_id}"
            if tag in info["answer_history"]:
                c = info["answer_history"][tag]["choice"]
                count[c] = count.get(c, 0) + 1

        for gid, hist in self.state.get("generic_agents", {}).items():
            tag = f"round_{round_id}"
            if tag in hist:
                c = hist[tag]["choice"]
                count[c] = count.get(c, 0) + 1

        for choice, n in count.items():
            if n >= math.ceil(total_agents * self.cons_thresh):
                return choice
        return None


    async def run_hybrid(self) -> dict:
        crashed_agents: List[str] = []
        while self.assign_attempts < self.max_assign:
            if self.assign_attempts:              # re‑assignment
                prev = self.assign_attempts
                summary = await self._panel_summary(prev)
                self.append_message(
                    "The previous specialist panel failed to reach a consensus.\n\n"
                    f"**Final positions from the panel:**\n{summary}\n\n"
                    "Please assemble a *fresh* set of medical specialties to resolve the disagreement.",
                    role="system"
                )

            await self._assign_specialists()
            pid = self.assign_attempts

            specialists = [
                DynamicSpecialist(role, pid, self.state)
                for role in self.state[f"panel_{pid}"]["Collected Specialists"]
            ]
            generics = [GenericAgent(f"generic_{i+1}", self.state) for i in range(self.n_generic)]
            active_agents = specialists + generics
            crashed_agents.clear()

            # ─── Round 1
            results = await asyncio.gather(
                *(a.analyse_question(self.q_text) for a in active_agents),
                return_exceptions=True
            )
            new_agents = []
            for ag, res in zip(active_agents, results):
                name = getattr(ag, "agent_id", getattr(ag, "specialist", ag.__class__.__name__))
                if isinstance(res, Exception):
                    logger.error("Agent %s crashed during analysis: %s", name, res)
                    crashed_agents.append(name)
                else:
                    new_agents.append(ag)
            active_agents = new_agents
            total = len(active_agents)
            if total == 0:
                raise RuntimeError(f"All agents crashed for QID {self.q_id}")

            choice = self._check_consensus(pid, 1, total)
            if choice:
                self.state["final"] = {"final_reasoning": "Consensus reached", "final_choice": choice}
                self.state["meta"]  = {"crashed_agents": crashed_agents, "active_agents": total, "round": 1}
                return self.state

            # ─── Debate rounds
            for r in range(2, self.max_cons_round + 1):
                results = await asyncio.gather(*(ag.debate() for ag in active_agents), return_exceptions=True)

                new_agents = []
                for ag, res in zip(active_agents, results):
                    name = getattr(ag, "agent_id", getattr(ag, "specialist", ag.__class__.__name__))
                    if isinstance(res, Exception):
                        logger.error("Agent %s crashed during debate round %d: %s", name, r, res)
                        crashed_agents.append(name)
                    else:
                        new_agents.append(ag)
                active_agents = new_agents
                total = len(active_agents)
                if total == 0:
                    raise RuntimeError(f"No agents survived round {r} for QID {self.q_id}")

                choice = self._check_consensus(pid, r, total)
                if choice:
                    self.state["final"] = {"final_reasoning": "Consensus reached", "final_choice": choice}
                    self.state["meta"]  = {"crashed_agents": crashed_agents, "active_agents": total, "round": r}
                    return self.state

        # Fallback aggregator
        hist = {k: v for k, v in self.state.items() if k not in ("label", "final")}
        parsed = await self._aggregate(hist)
        self.state["meta"] = {"crashed_agents": crashed_agents, "active_agents": len(active_agents), "round": None}
        return self.state

    async def _aggregate(self, chat_history):
        prompt = (
            "The specialists failed to reach consensus. Below is the entire "
            "conversation history:\n\n"
            f"{json.dumps(chat_history, indent=2)}\n\n"
            "Analyze their reasoning and provide:\n"
            " 1) A concise explanation of how you reached the final decision\n"
            " 2) The single best-supported choice (A, B, C, D, or E)"
        )
        class AggregatedResponse(BaseModel):
            final_reasoning: str = Field(..., description="Concise explanation.")
            final_choice:   Literal["A","B","C","D","E"] = Field(..., description="Final choice.")
        raw = await self.llm_call(prompt, guided_={"guided_json": AggregatedResponse.model_json_schema()})
        parsed = safe_json_load(raw)
        self.state["final"] = parsed
        return parsed


async def run_baseline(q_text: str, q_id: str, label: str) -> dict:
    agent = BaselineZS()
    out   = await agent.analyse_question(q_text)
    return {
        "method":   "baseline_zs",
        "q_id":     q_id,
        "label":    label,
        "choice":   (out or {}).get("choice", "ERROR"),
        "reasoning":(out or {}).get("reasoning", "failed"),
        "raw_state": out,
    }


async def run_generic(q_text: str, q_id: str, label: str) -> dict:
    mgr = Manager(q_text, q_id, label, n_specialists=0, n_generic=5)
    st  = await mgr.run_hybrid()
    return {
        "method":    "generic",
        "q_id":      q_id,
        "label":     label,
        "choice":    st["final"]["final_choice"],
        "reasoning": st["final"]["final_reasoning"],
        "raw_state": st,
    }


async def run_dynamic(q_text: str, q_id: str, label: str) -> dict:
    mgr = Manager(q_text, q_id, label, n_specialists=5, n_generic=0)
    st  = await mgr.run_hybrid()
    return {
        "method":    "dynamic",
        "q_id":      q_id,
        "label":     label,
        "choice":    st["final"]["final_choice"],
        "reasoning": st["final"]["final_reasoning"],
        "raw_state": st,
    }


async def run_special_generic(q_text: str, q_id: str, label: str) -> dict:
    mgr = Manager(q_text, q_id, label, n_specialists=3, n_generic=2, static_specialists=[])
    st  = await mgr.run_hybrid()
    return {
        "method":    "hybrid_special_generic",
        "q_id":      q_id,
        "label":     label,
        "choice":    st["final"]["final_choice"],
        "reasoning": st["final"]["final_reasoning"],
        "raw_state": st,
    }


async def run_static_dynamic(q_text: str, q_id: str, label: str) -> dict:
    mgr = Manager(q_text, q_id, label, n_specialists=5, n_generic=0, static_specialists=[])
    st  = await mgr.run_hybrid()
    return {
        "method":    "static_dynamic",
        "q_id":      q_id,
        "label":     label,
        "choice":    st["final"]["final_choice"],
        "reasoning": st["final"]["final_reasoning"],
        "raw_state": st,
    }



async def process_row(row: pd.Series) -> List[dict]:
    q_text = f"{row['question']}\n\n{str(row['choice'])}"
    q_id  = str(row["qn_num"])
    label = str(row["ground_truth"])

    method_runners = [
        run_baseline,
        run_generic,
        run_dynamic,
        run_special_generic,
        # run_static_dynamic,
    ]
    # names   = ["baseline_zs", "generic", "dynamic", "hybrid_special_generic", "static_dynamic"]
    method_names   = ["baseline_zs", "generic", "dynamic", "hybrid_special_generic"]

    tasks   = [m(q_text, q_id, label) for m in method_runners]
    results = await asyncio.gather(*tasks, return_exceptions=True)

    out = []
    for name, res in zip(method_names, results):
        if isinstance(res, Exception):
            logger.error("Method %s failed for QID %s: %s", name, q_id, res)
            out.append({"method": name, "q_id": q_id, "label": label, "error": str(res)})
        else:
            out.append(res)
    return out

async def process_failed_row(row):
    q_text = f"{row['question']}\n\n{str(row['choice'])}"
    q_id  = str(row["qn_num"])
    label = str(row["ground_truth"])

    method_runners = [
        run_baseline,
        run_generic,
        run_dynamic,
        run_special_generic,
        # run_static_dynamic,
    ]
    # names   = ["baseline_zs", "generic", "dynamic", "hybrid_special_generic", "static_dynamic"]
    method_names   = ["baseline_zs", "generic", "dynamic", "hybrid_special_generic"]

    tasks = []
    names = []
    for name, fn in zip(method_names, method_runners):
        if (q_id, name) in ERROR_KEYS:
            tasks.append(fn(q_text, q_id, label))
            names.append(name)

    if not tasks:
        return []

    results = await asyncio.gather(*tasks, return_exceptions=True)
    out = []
    for name, res in zip(names, results):
        if isinstance(res, Exception):
            logger.error("Method %s failed for QID %s: %s", name, q_id, res)
            out.append({"method": name, "q_id": q_id, "label": label, "error": str(res)})
        else:
            out.append(res)
    return out

async def process_dataset(df: pd.DataFrame) -> List[dict]:
    logger.info("Processing %d questions …", len(df))
    results = []
    for _, row in df.iterrows():
        results.extend(await process_row(row))
        # all_results.extend(await process_failed_row(row))
    return results


async def main() -> None:
    # 1) load dataset  (adjust the path as necessary)
    df = pd.read_csv("/home/yl3427/cylab/SOAP_MA/Input/filtered_merged_QA.csv", lineterminator="\n")

    # 2) run the multi‑agent protocols
    results = await process_dataset(df)

    # 3) save
    out_path = "/home/yl3427/cylab/SOAP_MA/Output/MedicalQA/medical_QA_MA_results.json"
    with open(out_path, "w") as f:
        json.dump(results, f, indent=2)

    logger.info("✅ Saved → %s", out_path)


if __name__ == "__main__":
    logging.basicConfig(
        level   = logging.INFO,
        format  = "%(asctime)s — %(levelname)s — %(message)s",
        datefmt = "%Y‑%m‑%d %H:%M:%S",
        handlers=[
            logging.FileHandler("log/0530_medical_QA_MA.log", "w"),
            logging.StreamHandler(),
        ],
    )
    asyncio.run(main())