| """Answer: hits -> abstain, or Claude Haiku answer with verified citations. |
| |
| Try one question: python -m multimodal_rag.answer "your question" |
| """ |
|
|
| import logging |
| import sys |
| import time |
| from collections.abc import Sequence |
| from dataclasses import replace |
|
|
| import anthropic |
| from dotenv import load_dotenv |
|
|
| from multimodal_rag.config import ( |
| ABSTAIN_TEXT, |
| ANSWER_K, |
| ANSWER_MAX_TOKENS, |
| CHECK_MAX_TOKENS, |
| CLAUDE_MODEL, |
| CLAUDE_TIMEOUT_S, |
| PROJECT_DIR, |
| REFUSAL_TEXT, |
| SELF_CHECK, |
| SELF_CHECK_MIN, |
| ) |
| from multimodal_rag.contracts import Answer, Chunk, Citation, Hit |
| from multimodal_rag.retrieve import ( |
| QueryError, |
| Retriever, |
| cost_usd, |
| gap_query, |
| search_query, |
| ) |
|
|
| log = logging.getLogger(__name__) |
|
|
| SYSTEM_PROMPT = ( |
| "You answer questions about FDA documents using only the documents provided. " |
| "Cite the passages that support each claim. Do not use outside knowledge. " |
| "If the question asks about a specific version, edition, year or period and " |
| "the documents only cover a different one, the documents do not contain the " |
| "answer. " |
| "If the documents do not contain the answer, reply with exactly this sentence " |
| f"and nothing else: {ABSTAIN_TEXT}" |
| ) |
| CHECK_PROMPT = ( |
| "You check an answer to a question about FDA documents. Using only the " |
| "documents shown, score from 1 to 10 how completely and correctly the answer " |
| "answers the question: 10 = fully answers it and every claim is supported by " |
| "the documents; 1 = does not answer it or contradicts the documents. " |
| "If the score is below 10, describe in one sentence the information that is " |
| "missing or unsupported; otherwise leave the gap empty. " |
| "Record your verdict with the record_check tool." |
| ) |
| |
| |
| CHECK_TOOL = { |
| "name": "record_check", |
| "description": "Record the 1-10 score and the missing information.", |
| "input_schema": { |
| "type": "object", |
| "properties": { |
| "score": {"type": "integer", "enum": list(range(1, 11))}, |
| "gap": {"type": "string"}, |
| }, |
| "required": ["score", "gap"], |
| }, |
| } |
|
|
|
|
| def doc_title(chunk: Chunk) -> str: |
| """Document title Claude sees: the 1-indexed PDF page a reader would look up |
| (data structures keep the 0-indexed page), marked when the text is a |
| generated description rather than the PDF's own words.""" |
| page = f"{chunk.doc_id} p.{chunk.page + 1}" |
| return f"generated visual description, {page}" if chunk.kind == "vision" else page |
|
|
|
|
| def build_documents(hits: list[Hit]) -> list[dict]: |
| """One citable plain-text document per chunk; list position = document_index.""" |
| return [ |
| { |
| "type": "document", |
| "source": { |
| "type": "text", |
| "media_type": "text/plain", |
| "data": h.chunk.text, |
| }, |
| "title": doc_title(h.chunk), |
| "citations": {"enabled": True}, |
| } |
| for h in hits |
| ] |
|
|
|
|
| def verify_citations(content: list, hits: list[Hit]) -> list[Citation]: |
| """Citations from Claude's text blocks whose quote matches the chunk exactly. |
| |
| A citation is kept only if `cited_text == chunk.text[start:end]` for the chunk |
| at its `document_index`; any other is dropped and logged. `generated` is set |
| for quotes from vision chunks. |
| """ |
| out: list[Citation] = [] |
| for block in content: |
| for c in getattr(block, "citations", None) or []: |
| chunk = hits[c.document_index].chunk |
| span = chunk.text[c.start_char_index : c.end_char_index] |
| if c.type != "char_location" or span != c.cited_text: |
| log.warning("dropped unverified citation of %s", chunk.chunk_id) |
| continue |
| score = hits[c.document_index].score |
| generated = chunk.kind == "vision" |
| out.append( |
| Citation( |
| chunk.chunk_id, chunk.doc_id, chunk.page, span, score, generated |
| ) |
| ) |
| return out |
|
|
|
|
| def answer(question: str, hits: list[Hit], client: anthropic.Anthropic) -> Answer: |
| """Answer from the top ANSWER_K hits with verified citations; abstain without |
| calling Claude only if there are no hits. Claude decides every other abstain |
| (no score threshold: in testing it caught no unanswerable query Claude |
| missed and blocked answerable ones with the gold page in the top 10). |
| |
| If Claude stops with `refusal`, returns REFUSAL_TEXT with `refused=True`: |
| partial text and citations are discarded, the cost is kept. Raises |
| QueryError if the call fails or stops any other way than `end_turn` |
| (e.g. `max_tokens` cut the answer and its citations short). |
| """ |
| t0 = time.monotonic() |
| if not hits: |
| return Answer(ABSTAIN_TEXT, [], True, False, 0.0, time.monotonic() - t0) |
|
|
| hits = hits[:ANSWER_K] |
| try: |
| msg = client.messages.create( |
| model=CLAUDE_MODEL, |
| max_tokens=ANSWER_MAX_TOKENS, |
| system=SYSTEM_PROMPT, |
| messages=[ |
| { |
| "role": "user", |
| "content": build_documents(hits) |
| + [{"type": "text", "text": question}], |
| } |
| ], |
| timeout=CLAUDE_TIMEOUT_S, |
| ) |
| except anthropic.APIError as e: |
| raise QueryError(f"Claude call failed: {type(e).__name__}: {e}") from e |
|
|
| cost = cost_usd(msg.usage) |
| log.info( |
| "claude %s in=%d out=%d cost=$%.6f", |
| msg.stop_reason, |
| msg.usage.input_tokens, |
| msg.usage.output_tokens, |
| cost, |
| ) |
| if msg.stop_reason == "refusal": |
| return Answer(REFUSAL_TEXT, [], False, True, cost, time.monotonic() - t0) |
| if msg.stop_reason != "end_turn": |
| raise QueryError(f"Claude stopped with {msg.stop_reason} (cost ${cost:.6f})") |
|
|
| text = "".join(b.text for b in msg.content if b.type == "text").strip() |
| citations = verify_citations(msg.content, hits) |
| abstained = text == ABSTAIN_TEXT and not citations |
| return Answer(text, citations, abstained, False, cost, time.monotonic() - t0) |
|
|
|
|
| def answered(a: Answer) -> bool: |
| """Claude answered from the documents: neither an abstain nor a refusal.""" |
| return not (a.abstained or a.refused) |
|
|
|
|
| def self_check( |
| question: str, ans: Answer, hits: list[Hit], client: anthropic.Anthropic |
| ) -> tuple[int, str, float]: |
| """Haiku scores `ans` 1-10 against the question and the ANSWER_K chunks it |
| was written from; (score, gap, cost_usd). `gap` = one sentence on what is |
| missing or unsupported ("" at 10). Raises QueryError on an API error or a |
| score outside 1-10.""" |
| docs = "\n".join( |
| f'<document title="{doc_title(h.chunk)}">{h.chunk.text}</document>' |
| for h in hits[:ANSWER_K] |
| ) |
| prompt = ( |
| f"<documents>\n{docs}\n</documents>\n" |
| f"<question>{question}</question>\n<answer>{ans.text}</answer>" |
| ) |
| try: |
| msg = client.messages.create( |
| model=CLAUDE_MODEL, |
| max_tokens=CHECK_MAX_TOKENS, |
| system=CHECK_PROMPT, |
| tools=[CHECK_TOOL], |
| tool_choice={"type": "tool", "name": CHECK_TOOL["name"]}, |
| messages=[{"role": "user", "content": prompt}], |
| timeout=CLAUDE_TIMEOUT_S, |
| ) |
| except anthropic.APIError as e: |
| raise QueryError(f"self-check failed: {type(e).__name__}: {e}") from e |
| cost = cost_usd(msg.usage) |
| if msg.stop_reason != "tool_use": |
| raise QueryError(f"self-check stopped with {msg.stop_reason} (${cost:.6f})") |
| verdict = next(b.input for b in msg.content if b.type == "tool_use") |
| if verdict.get("score") not in range(1, 11): |
| raise QueryError(f"self-check returned an invalid score: {verdict}") |
| return verdict["score"], verdict.get("gap", ""), cost |
|
|
|
|
| def check_and_retry( |
| question: str, |
| first: Answer, |
| hits: list[Hit], |
| retriever: Retriever, |
| client: anthropic.Anthropic, |
| ) -> tuple[Answer, list[Hit]]: |
| """Self-check `first`; below SELF_CHECK_MIN, search once more and answer again. |
| |
| An abstain or a refusal is final and not checked. Otherwise the score becomes |
| `confidence`. Below SELF_CHECK_MIN: the gap becomes a new search query |
| (`gap_query`), searched with the normal pipeline and reranked against |
| `question`, and Claude answers `question` from the new hits. That second |
| answer is final, even an abstain; it is checked once more only to fill |
| `confidence` (`retried=True`). Returns the final answer, with the cost and |
| latency of every step, and the ANSWER_K hits it was written from. |
| """ |
| if not answered(first): |
| return first, hits[:ANSWER_K] |
| t0 = time.monotonic() |
| score, gap, check_cost = self_check(question, first, hits, client) |
| if score >= SELF_CHECK_MIN: |
| return ( |
| replace( |
| first, |
| cost_usd=first.cost_usd + check_cost, |
| latency_s=first.latency_s + time.monotonic() - t0, |
| confidence=score, |
| ), |
| hits[:ANSWER_K], |
| ) |
| query, query_cost = gap_query(question, gap, client) |
| new_hits = retriever.search(query, question=question) |
| second = answer(question, new_hits, client) |
| cost = first.cost_usd + check_cost + query_cost + second.cost_usd |
| confidence = None |
| if answered(second): |
| confidence, _, second_check_cost = self_check( |
| question, second, new_hits, client |
| ) |
| cost += second_check_cost |
| log.info("self-check %d < %d, retried with %r", score, SELF_CHECK_MIN, query) |
| final = replace( |
| second, |
| cost_usd=cost, |
| latency_s=first.latency_s + time.monotonic() - t0, |
| confidence=confidence, |
| retried=True, |
| ) |
| return final, new_hits[:ANSWER_K] |
|
|
|
|
| def ask( |
| question: str, |
| retriever: Retriever, |
| client: anthropic.Anthropic, |
| history: Sequence[str] = (), |
| self_check_on: bool = SELF_CHECK, |
| ) -> tuple[Answer, list[Hit]]: |
| """Transform the question (follow-up rewrite, REWRITE, HYDE), retrieve, answer, |
| and with SELF_CHECK check the answer and retry once (`check_and_retry`). |
| |
| `history` = the conversation's earlier user questions, oldest first. |
| Returns the answer and the ANSWER_K hits it was written from (after a |
| retry, the retry's). `cost_usd` includes the rewrite / HyDE calls; |
| `latency_s` covers all steps. |
| """ |
| t0 = time.monotonic() |
| sq = search_query(question, client, history) |
| hits = retriever.search(sq.text, passage=sq.passage, question=sq.question) |
| result = answer(sq.question, hits, client) |
| hits = hits[:ANSWER_K] |
| if self_check_on: |
| result, hits = check_and_retry(sq.question, result, hits, retriever, client) |
| final = replace( |
| result, |
| cost_usd=result.cost_usd + sq.cost_usd, |
| latency_s=time.monotonic() - t0, |
| ) |
| return final, hits |
|
|
|
|
| def main() -> None: |
| """Answer the question given on the command line and print its citations.""" |
| logging.basicConfig(level=logging.WARNING, format="%(levelname)s %(message)s") |
| log.setLevel(logging.INFO) |
| load_dotenv(PROJECT_DIR / ".env") |
| result, _ = ask(sys.argv[1], Retriever(), anthropic.Anthropic()) |
| print(result.text, "\n") |
| for c in result.citations: |
| where = f"{c.doc_id} p.{c.page + 1} [{c.chunk_id}, {c.score:.3f}]" |
| mark = " (generated description)" if c.generated else "" |
| print(f"- {where}{mark} {c.cited_text!r}") |
| print( |
| f"\nabstained={result.abstained} refused={result.refused} " |
| f"confidence={result.confidence} retried={result.retried} " |
| f"cost=${result.cost_usd:.6f} " |
| f"latency={result.latency_s:.2f}s" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
