chore: import upstream snapshot with attribution

examples/chartqa/chartqa_agent.py

@@ -0,0 +1,410 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""ChartQA agent demonstrating LangGraph-based visual reasoning with refinement.
+
+This module defines `ChartQAAgent` plus the supporting prompt utilities used by
+`debug_chartqa_agent.py` and `train_chartqa_agent.py`.
+
+1. `analyze_chart` observes and summarizes the chart.
+2. `extract_data` calls a text-only LLM to extract the requested values.
+3. `calculate_answer` runs calculations grounded in prior steps.
+4. `check_answer` verifies reasoning quality.
+5. `refine_answer` conditionally patches mistakes before responding.
+
+Example usage can be found in `debug_chartqa_agent.py` and `train_chartqa_agent.py`.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import re
+from typing import Any, Dict, Literal, cast
+
+import env_var as chartqa_env_var
+import termcolor
+from langchain.chat_models import BaseChatModel, init_chat_model
+from langchain_core.messages import AnyMessage, BaseMessage, HumanMessage
+from langgraph.graph import END, START, MessagesState, StateGraph
+from langgraph.graph.state import CompiledStateGraph
+from multimodal_utils import encode_image_to_base64
+from prompts import (
+    ANALYZE_CHART_PROMPT,
+    CALCULATE_ANSWER_PROMPT,
+    CHECK_ANSWER_PROMPT,
+    EXTRACT_DATA_PROMPT,
+    REFINE_ANSWER_PROMPT,
+)
+
+import agentlightning as agl
+
+logger = logging.getLogger("chartqa_agent")
+
+
+class ChartState(MessagesState):
+    question: str
+    image_path: str
+    observation: str
+    extracted_data: str
+    calculation: str
+    answer: str
+    feedback: str
+    num_turns: int
+    messages: list[AnyMessage]
+
+
+class ChartQAAgent(agl.LitAgent[Dict[str, Any]]):
+    """LangGraph-powered ChartQA agent with multi-step reasoning and refinement.
+
+    The implementation shares the same [`agl.LitAgent`][agentlightning.LitAgent] interface as
+    the Calc-X sample agent but augments it with image handling and LangGraph state tracking.
+    """
+
+    def __init__(
+        self,
+        model_name: str | None = None,
+        max_turns: int = 3,
+        debug: bool = False,
+        endpoint: str | None = None,
+        temperature: float = 0.0,
+        use_base64_images: bool = False,
+    ):
+        self.debug = debug
+        self.max_turns = max_turns
+        self.use_base64_images = use_base64_images
+        self.model_name = model_name
+        self.endpoint = endpoint
+        self.temperature = temperature
+
+        self._llm: BaseChatModel | None = None
+        self._graph: CompiledStateGraph[ChartState] | None = None
+
+    def _create_llm(self) -> BaseChatModel:
+        if self.model_name is None:
+            raise ValueError("model_name is required for creating LLM")
+        return init_chat_model(
+            self.model_name,
+            model_provider="openai",
+            openai_api_base=self.endpoint,
+            openai_api_key=chartqa_env_var.OPENAI_API_KEY,
+            temperature=self.temperature,
+            max_retries=2,
+            max_tokens=1024,
+            timeout=300,
+        )
+
+    def update_llm_config(self, model_name: str, endpoint: str | None, temperature: float | None) -> None:
+        """Update the LLM configuration. Re-create the LLM if the configuration is changed."""
+        updated: bool = False
+        if model_name != self.model_name:
+            self.model_name = model_name
+            updated = True
+        if endpoint != self.endpoint:
+            self.endpoint = endpoint
+            updated = True
+        if temperature != self.temperature:
+            self.temperature = temperature
+            updated = True
+        if updated:
+            self._llm = self._create_llm()
+
+    def _ensure_llm(self) -> BaseChatModel:
+        """Ensure the LLM is created and cached."""
+        if self._llm is None:
+            self._llm = self._create_llm()
+        return self._llm
+
+    def invoke_prompt(self, prompt: Any) -> AnyMessage:
+        """Invoke LLM with prompt."""
+        if self.debug:
+            for message in prompt.messages:
+                termcolor.cprint(message.pretty_repr(), "blue")
+
+        try:
+            result = self._ensure_llm().invoke(prompt)
+        except Exception as e:
+            logger.error(f"Failed to invoke prompt: {e}")
+            result = self._ensure_llm().invoke([HumanMessage(content="Please provide a reasonable answer.")])
+
+        if self.debug:
+            termcolor.cprint(result.pretty_repr(), "green")
+
+        return result  # type: ignore
+
+    def invoke_prompt_with_image(self, prompt_text: str, image_path: str) -> str:
+        """Invoke vision-language model with image.
+
+        Handles both local vLLM (file:// URLs) and cloud APIs (base64 encoding).
+        Cloud APIs (OpenAI, Anthropic, Google, Azure, etc.) require base64 encoding.
+        """
+        # Determine image URL format based on endpoint
+        if self.use_base64_images:
+            # Cloud APIs require base64 encoding for local files
+            image_url = encode_image_to_base64(image_path)
+        else:
+            # Local vLLM supports file:// URLs
+            if not image_path.startswith("file://"):
+                image_path = f"file://{os.path.realpath(image_path)}"
+            image_url = image_path
+
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt_text},
+                    {"type": "image_url", "image_url": {"url": image_url}},
+                ],
+            }
+        ]
+
+        if self.debug:
+            termcolor.cprint(f"[VLM Call] {prompt_text[:100]}...", "blue")
+
+        try:
+            result = self._ensure_llm().invoke(messages)
+            response = result.content if hasattr(result, "content") else str(result)  # type: ignore
+        except Exception as e:
+            logger.error(f"Failed to invoke VLM: {e}")
+            response = "<observe>Unable to analyze chart</observe>"
+
+        if self.debug:
+            termcolor.cprint(f"[VLM Response] {response[:200]}...", "green")
+
+        return response  # type: ignore
+
+    def extract_content(self, text: str, tag: str) -> str:
+        """Extract content between XML-style tags."""
+        match = re.search(rf"<{tag}>(.*?)</{tag}>", text, re.DOTALL)
+        return match.group(1).strip() if match else ""
+
+    def analyze_chart(self, state: ChartState) -> ChartState:
+        """Step 1: Observe and describe the chart."""
+        prompt: Any = ANALYZE_CHART_PROMPT.invoke({"question": state["question"]})  # type: ignore
+        prompt_text = prompt.messages[1].content
+
+        result_text = self.invoke_prompt_with_image(prompt_text, state["image_path"])
+
+        observation = self.extract_content(result_text, "observe")
+        if not observation:
+            observation = result_text
+
+        return {  # type: ignore
+            **state,
+            "observation": observation,
+            "num_turns": 1,
+            "messages": [HumanMessage(content=result_text)],
+        }
+
+    def extract_data(self, state: ChartState) -> ChartState:
+        """Step 2: Extract specific data values."""
+        prompt: Any = EXTRACT_DATA_PROMPT.invoke(  # type: ignore
+            {
+                "observation": state["observation"],
+                "question": state["question"],
+            }
+        )
+        result = self.invoke_prompt(prompt)
+
+        extracted_data = self.extract_content(result.content, "extract")  # type: ignore
+        if not extracted_data:
+            extracted_data = result.content  # type: ignore
+
+        return {  # type: ignore
+            **state,
+            "extracted_data": extracted_data,  # type: ignore
+            "messages": [*state.get("messages", []), result],
+        }
+
+    def calculate_answer(self, state: ChartState) -> ChartState:
+        """Step 3: Calculate and provide answer."""
+        prompt: Any = CALCULATE_ANSWER_PROMPT.invoke(  # type: ignore
+            {
+                "extracted_data": state["extracted_data"],
+                "question": state["question"],
+            }
+        )
+        result = self.invoke_prompt(prompt)
+
+        calculation = self.extract_content(result.content, "calculate")  # type: ignore
+        answer = self.extract_content(result.content, "answer")  # type: ignore
+        if not answer:
+            answer = cast(str, result.content)  # type: ignore
+
+        return {  # type: ignore
+            **state,
+            "calculation": calculation,
+            "answer": answer,
+            "messages": [*state.get("messages", []), result],
+        }
+
+    def check_answer(self, state: ChartState) -> ChartState:
+        """Step 4: Verify answer quality."""
+        prompt: Any = CHECK_ANSWER_PROMPT.invoke(  # type: ignore
+            {
+                "observation": state["observation"],
+                "extracted_data": state["extracted_data"],
+                "question": state["question"],
+                "answer": state["answer"],
+                "calculation": state.get("calculation", "No calculation shown"),
+            }
+        )
+        result = self.invoke_prompt(prompt)
+
+        if self.debug:
+            termcolor.cprint(f"[Check] {result.content}", "yellow")  # type: ignore
+
+        return {  # type: ignore
+            **state,
+            "feedback": result.content,  # type: ignore
+            "messages": [*state.get("messages", []), *prompt.messages, result],
+        }
+
+    def refine_answer(self, state: ChartState) -> ChartState:
+        """Step 5: Refine answer based on feedback."""
+        prompt: Any = REFINE_ANSWER_PROMPT.invoke(  # type: ignore
+            {
+                "observation": state["observation"],
+                "extracted_data": state["extracted_data"],
+                "question": state["question"],
+                "answer": state["answer"],
+                "calculation": state.get("calculation", ""),
+                "feedback": state["feedback"],
+            }
+        )
+        result = self.invoke_prompt(prompt)
+        content: str = result.content  # type: ignore
+
+        new_extracted = self.extract_content(content, "extract")
+        extracted_data = new_extracted if new_extracted else state["extracted_data"]
+
+        new_calculation = self.extract_content(content, "calculate")
+
+        new_answer = self.extract_content(content, "answer")
+        if not new_answer:
+            new_answer = content
+
+        return {  # type: ignore
+            **state,
+            "extracted_data": extracted_data,
+            "calculation": new_calculation,
+            "answer": new_answer,
+            "num_turns": state.get("num_turns", 0) + 1,
+            "messages": [*prompt.messages, result],
+        }
+
+    def should_continue(self, state: ChartState) -> Literal[END, "refine_answer"]:  # type: ignore
+        """Determine if refinement is needed."""
+        if state["messages"] and isinstance(
+            state["messages"][-1], BaseMessage
+        ):  # pyright: ignore[reportUnnecessaryIsInstance]
+            last_message = state["messages"][-1]
+            if "THE ANSWER IS CORRECT" in last_message.content:  # type: ignore
+                if "THE ANSWER IS INCORRECT" in last_message.content:  # type: ignore
+                    correct_index = last_message.content.rfind("THE ANSWER IS CORRECT")  # type: ignore
+                    incorrect_index = last_message.content.rfind("THE ANSWER IS INCORRECT")  # type: ignore
+                    if correct_index > incorrect_index:
+                        return END
+                else:
+                    return END
+
+        if state.get("num_turns", 0) >= self.max_turns:
+            return END
+
+        return "refine_answer"
+
+    def graph(self) -> CompiledStateGraph[ChartState]:
+        """Build the workflow graph with refinement loop."""
+        # Check if the graph is already built
+        if self._graph is not None:
+            return self._graph
+
+        builder = StateGraph(ChartState)
+        builder.add_node(self.analyze_chart)  # type: ignore
+        builder.add_node(self.extract_data)  # type: ignore
+        builder.add_node(self.calculate_answer)  # type: ignore
+        builder.add_node(self.check_answer)  # type: ignore
+        builder.add_node(self.refine_answer)  # type: ignore
+
+        builder.add_edge(START, "analyze_chart")
+        builder.add_edge("analyze_chart", "extract_data")
+        builder.add_edge("extract_data", "calculate_answer")
+        builder.add_edge("calculate_answer", "check_answer")
+        builder.add_conditional_edges(
+            "check_answer",
+            self.should_continue,  # type: ignore
+        )
+        builder.add_edge("refine_answer", "extract_data")
+
+        self._graph = builder.compile()  # type: ignore
+        return self._graph
+
+    def rollout(self, task: Dict[str, Any], resources: agl.NamedResources, rollout: agl.Rollout) -> float | None:
+        """AgentLightning wrapper for ChartQA agent."""
+
+        question = task["question"]
+
+        rollout = cast(agl.AttemptedRollout, rollout)
+        llm = cast(agl.LLM, resources["main_llm"])
+
+        image_path = os.path.join(chartqa_env_var.CHARTQA_DATA_DIR, task["image_path"])
+        ground_truth = task["answer"]
+
+        if not os.path.exists(image_path):
+            logger.error(f"Image {image_path} does not exist. Skipping.")
+            return None
+
+        # The new rollout could have a different endpoint or temperature.
+        # Update the LLM if necessary.
+        self.update_llm_config(
+            model_name=llm.model,
+            endpoint=llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
+            temperature=llm.sampling_parameters.get("temperature", 0.0),
+        )
+
+        try:
+            handler = self.tracer.get_langchain_handler()
+            result = self.graph().invoke(  # type: ignore
+                {"question": question, "image_path": image_path},  # type: ignore
+                {"callbacks": [handler] if handler else [], "recursion_limit": 100},
+            )
+        except Exception as e:
+            error_msg = f"[Rollout {rollout.rollout_id}] Error during agent invocation: {e}"
+            logger.error(error_msg, exc_info=True)
+            # Return 0.0 as reward to indicate failure
+            return 0.0
+
+        predicted_answer = result["answer"]
+        reward = evaluate_answer(predicted_answer, ground_truth, raise_on_error=False)
+
+        return reward
+
+
+def evaluate_answer(predicted: str, ground_truth: str, raise_on_error: bool = False) -> float:
+    """Evaluate answer accuracy."""
+    try:
+        pred = predicted.lower().strip()
+        gt = ground_truth.lower().strip()
+
+        # Exact match
+        if pred == gt:
+            return 1.0
+
+        # Try numeric comparison
+        try:
+            pred_num = float(pred.replace(",", ""))
+            gt_num = float(gt.replace(",", ""))
+            if abs(pred_num - gt_num) / max(abs(gt_num), 1e-9) < 0.02:
+                return 1.0
+        except (ValueError, AttributeError):
+            pass
+
+        # Partial credit for substring match
+        if pred in gt or gt in pred:
+            return 0.5
+
+        return 0.0
+    except Exception as e:
+        if raise_on_error:
+            raise
+        logger.exception(f"Error evaluating answer: {e}")
+        return 0.0