deeptutor-visualize 1.7.0__py3-none-any.whl

deeptutor/agents/visualize/__init__.py

@@ -0,0 +1,5 @@
+"""Visualize agents and pipeline."""
+
+from .pipeline import VisualizePipeline
+
+__all__ = ["VisualizePipeline"]

deeptutor/agents/visualize/agents/__init__.py

@@ -0,0 +1,11 @@
+"""Agent building blocks for the visualize capability."""
+
+from .analysis_agent import AnalysisAgent
+from .code_generator_agent import CodeGeneratorAgent
+from .review_agent import ReviewAgent
+
+__all__ = [
+    "AnalysisAgent",
+    "CodeGeneratorAgent",
+    "ReviewAgent",
+]

deeptutor/agents/visualize/agents/analysis_agent.py

@@ -0,0 +1,102 @@
+"""Analysis stage: decide SVG vs Chart.js and produce a structured brief."""
+
+from __future__ import annotations
+
+from deeptutor.agents.base_agent import BaseAgent
+from deeptutor.core.context import Attachment
+from deeptutor.core.trace import build_trace_metadata, new_call_id
+
+from ..models import VisualizationAnalysis
+from ..utils import extract_json_object
+
+
+class AnalysisAgent(BaseAgent):
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        api_version: str | None = None,
+        language: str = "zh",
+    ) -> None:
+        super().__init__(
+            module_name="visualize",
+            agent_name="analysis_agent",
+            api_key=api_key,
+            base_url=base_url,
+            api_version=api_version,
+            language=language,
+        )
+
+    async def process(
+        self,
+        *,
+        user_input: str,
+        history_context: str,
+        render_mode: str = "auto",
+        attachments: list[Attachment] | None = None,
+    ) -> VisualizationAnalysis:
+        # Manim modes short-circuit the LLM call entirely — MathAnimatorPipeline
+        # has its own ConceptAnalysisAgent that produces the manim-specific
+        # design brief. We only need a stub VisualizationAnalysis here so the
+        # capability runner can dispatch on render_type.
+        if render_mode in ("manim_video", "manim_image"):
+            return VisualizationAnalysis(
+                render_type=render_mode,  # type: ignore[arg-type]
+                description="",
+                data_description="",
+                chart_type="",
+                visual_elements=[],
+                rationale=f"User selected {render_mode} explicitly.",
+            )
+
+        if render_mode in ("svg", "chartjs", "mermaid", "mindmap", "html"):
+            system_prompt = self.get_prompt("system_fixed")
+            user_template = self.get_prompt("user_template_fixed")
+        elif render_mode == "figure":
+            # Constrained-auto mode: LLM picks one of svg/chartjs/mermaid
+            # (html is excluded). Used by the Book figure block.
+            system_prompt = self.get_prompt("system_figure")
+            user_template = self.get_prompt("user_template_figure")
+        else:
+            system_prompt = self.get_prompt("system")
+            user_template = self.get_prompt("user_template")
+        if not system_prompt or not user_template:
+            raise ValueError("AnalysisAgent prompts are not configured.")
+
+        format_kwargs: dict[str, str] = {
+            "user_input": user_input.strip(),
+            "history_context": history_context.strip() or "(none)",
+        }
+        if render_mode in ("svg", "chartjs", "mermaid", "mindmap", "html"):
+            format_kwargs["render_type"] = render_mode
+
+        user_prompt = user_template.format(**format_kwargs)
+
+        # Blocking rather than streamed: nothing consumes the partial JSON, and
+        # a reasoning model's <think> prelude never reaches the parser this way.
+        raw = await self.call_llm(
+            user_prompt=user_prompt,
+            system_prompt=system_prompt,
+            response_format={"type": "json_object"},
+            stage="analyzing",
+            attachments=attachments,
+            trace_meta=build_trace_metadata(
+                call_id=new_call_id("viz-analysis"),
+                phase="analyzing",
+                label="Visualization analysis",
+                call_kind="viz_analysis",
+                trace_role="analyze",
+                trace_kind="llm_output",
+            ),
+        )
+        result = VisualizationAnalysis.model_validate(extract_json_object(raw))
+        if render_mode in ("svg", "chartjs", "mermaid", "mindmap", "html"):
+            result.render_type = render_mode  # type: ignore[assignment]
+        elif render_mode == "figure" and result.render_type not in (
+            "svg",
+            "chartjs",
+            "mermaid",
+        ):
+            # Defensive: if the LLM ignored the constraint, force a safe default.
+            result.render_type = "svg"  # type: ignore[assignment]
+        return result

deeptutor/agents/visualize/agents/code_generator_agent.py

@@ -0,0 +1,117 @@
+"""Code generation stage: produce SVG or Chart.js code from the analysis."""
+
+from __future__ import annotations
+
+import json
+
+from deeptutor.agents.base_agent import BaseAgent
+from deeptutor.core.trace import build_trace_metadata, new_call_id
+
+from ..models import VisualizationAnalysis
+from ..utils import extract_code_block
+
+
+class CodeGeneratorAgent(BaseAgent):
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        api_version: str | None = None,
+        language: str = "zh",
+    ) -> None:
+        super().__init__(
+            module_name="visualize",
+            agent_name="code_generator_agent",
+            api_key=api_key,
+            base_url=base_url,
+            api_version=api_version,
+            language=language,
+        )
+
+    async def process(
+        self,
+        *,
+        user_input: str,
+        history_context: str,
+        analysis: VisualizationAnalysis,
+    ) -> str:
+        # Structured prompt assembly: every render type loads the shared
+        # base + general rules, plus exactly one format-specific rule block.
+        # This keeps the per-call prompt dense with the rules that matter for
+        # *this* format without ever paying for the union of all four.
+        # render_type here is always one of svg/chartjs/mermaid/html — manim
+        # is dispatched away in the capability layer before code generation.
+        system_parts = (
+            self.get_prompt("system_base"),
+            self.get_prompt("rules_general"),
+            self.get_prompt(f"rules_{analysis.render_type}"),
+        )
+        system_prompt = "\n\n".join(part.strip() for part in system_parts if part and part.strip())
+        user_template = self.get_prompt("user_template")
+        if not system_prompt or not user_template:
+            raise ValueError("CodeGeneratorAgent prompts are not configured.")
+
+        user_prompt = user_template.format(
+            user_input=user_input.strip(),
+            history_context=history_context.strip() or "(none)",
+            render_type=analysis.render_type,
+            analysis_json=json.dumps(analysis.model_dump(), ensure_ascii=False, indent=2),
+        )
+
+        # Blocking rather than streamed: the fenced block is only usable whole.
+        response = await self.call_llm(
+            user_prompt=user_prompt,
+            system_prompt=system_prompt,
+            stage="generating",
+            trace_meta=build_trace_metadata(
+                call_id=new_call_id("viz-codegen"),
+                phase="generating",
+                label="Code generation",
+                call_kind="viz_code_generation",
+                trace_role="generate",
+                trace_kind="llm_output",
+            ),
+        )
+
+        if analysis.render_type == "svg":
+            lang_hint = "svg"
+        elif analysis.render_type in ("mermaid", "mindmap"):
+            lang_hint = "mermaid"
+        elif analysis.render_type == "html":
+            lang_hint = "html"
+        else:
+            lang_hint = "javascript"
+
+        extracted = extract_code_block(response, lang_hint) or extract_code_block(response)
+
+        # For html, the model sometimes returns the full document with no fence.
+        # `extract_code_block` will then return the trimmed raw response — accept
+        # it as long as it looks like an HTML document.
+        if analysis.render_type == "html" and not extracted:
+            stripped = (response or "").strip()
+            lowered = stripped.lower()
+            if lowered.startswith("<!doctype") or lowered.startswith("<html"):
+                return stripped
+
+        # The renderers expect their payload to start cleanly with the right
+        # root tag. Models often wrap output with prose ("Here you go: <svg>…")
+        # or emit a code fence on the same line as the closing tag, which
+        # `extract_code_block`'s regex (requiring a leading \n before ```) does
+        # not strip. Trim to the outermost root tags as a defensive net.
+        if extracted:
+            if analysis.render_type == "svg":
+                lower = extracted.lower()
+                start = lower.find("<svg")
+                end = lower.rfind("</svg>")
+                if start != -1 and end != -1 and end > start:
+                    extracted = extracted[start : end + len("</svg>")]
+            elif analysis.render_type == "html":
+                lower = extracted.lower()
+                start = lower.find("<!doctype")
+                if start == -1:
+                    start = lower.find("<html")
+                end = lower.rfind("</html>")
+                if start != -1 and end != -1 and end > start:
+                    extracted = extracted[start : end + len("</html>")]
+
+        return extracted

deeptutor/agents/visualize/agents/review_agent.py

@@ -0,0 +1,75 @@
+"""Repair stage: targeted fix after deterministic local validation fails."""
+
+from __future__ import annotations
+
+import json
+
+from deeptutor.agents.base_agent import BaseAgent
+from deeptutor.core.trace import build_trace_metadata, new_call_id
+
+from ..models import ReviewResult, VisualizationAnalysis
+from ..utils import extract_json_object
+
+
+class ReviewAgent(BaseAgent):
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        api_version: str | None = None,
+        language: str = "zh",
+    ) -> None:
+        super().__init__(
+            module_name="visualize",
+            agent_name="review_agent",
+            api_key=api_key,
+            base_url=base_url,
+            api_version=api_version,
+            language=language,
+        )
+
+    async def process(
+        self,
+        *,
+        user_input: str,
+        analysis: VisualizationAnalysis,
+        code: str,
+        error: str,
+    ) -> ReviewResult:
+        """Targeted repair pass — this agent's single job.
+
+        Only invoked *after* the deterministic local validation has failed,
+        and handed the concrete error so the model fixes one specific defect
+        instead of re-judging everything (chat and Book share this path).
+        """
+        system_prompt = self.get_prompt("repair_system")
+        user_template = self.get_prompt("repair_user_template")
+        if not system_prompt or not user_template:
+            raise ValueError("ReviewAgent repair prompts are not configured.")
+
+        user_prompt = user_template.format(
+            user_input=user_input.strip(),
+            render_type=analysis.render_type,
+            error=error.strip() or "(unspecified)",
+            analysis_json=json.dumps(analysis.model_dump(), ensure_ascii=False, indent=2),
+            code=code,
+        )
+
+        chunks: list[str] = []
+        async for chunk in self.stream_llm(
+            user_prompt=user_prompt,
+            system_prompt=system_prompt,
+            response_format={"type": "json_object"},
+            stage="reviewing",
+            trace_meta=build_trace_metadata(
+                call_id=new_call_id("viz-repair"),
+                phase="reviewing",
+                label="Code repair",
+                call_kind="viz_code_repair",
+                trace_role="review",
+                trace_kind="llm_output",
+            ),
+        ):
+            chunks.append(chunk)
+        response = "".join(chunks)
+        return ReviewResult.model_validate(extract_json_object(response))