|
| 1 | +"""Analysis stage: decide SVG vs Chart.js and produce a structured brief.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +from deeptutor.agents.base_agent import BaseAgent |
| 6 | +from deeptutor.core.trace import build_trace_metadata, new_call_id |
| 7 | + |
| 8 | +from ..models import VisualizationAnalysis |
| 9 | +from ..utils import extract_json_object |
| 10 | + |
| 11 | + |
| 12 | +class AnalysisAgent(BaseAgent): |
| 13 | + def __init__( |
| 14 | + self, |
| 15 | + api_key: str | None = None, |
| 16 | + base_url: str | None = None, |
| 17 | + api_version: str | None = None, |
| 18 | + language: str = "zh", |
| 19 | + ) -> None: |
| 20 | + super().__init__( |
| 21 | + module_name="visualize", |
| 22 | + agent_name="analysis_agent", |
| 23 | + api_key=api_key, |
| 24 | + base_url=base_url, |
| 25 | + api_version=api_version, |
| 26 | + language=language, |
| 27 | + ) |
| 28 | + |
| 29 | + async def process( |
| 30 | + self, |
| 31 | + *, |
| 32 | + user_input: str, |
| 33 | + history_context: str, |
| 34 | + render_mode: str = "auto", |
| 35 | + ) -> VisualizationAnalysis: |
| 36 | + if render_mode in ("svg", "chartjs"): |
| 37 | + system_prompt = self.get_prompt("system_fixed") |
| 38 | + user_template = self.get_prompt("user_template_fixed") |
| 39 | + else: |
| 40 | + system_prompt = self.get_prompt("system") |
| 41 | + user_template = self.get_prompt("user_template") |
| 42 | + if not system_prompt or not user_template: |
| 43 | + raise ValueError("AnalysisAgent prompts are not configured.") |
| 44 | + |
| 45 | + format_kwargs: dict[str, str] = { |
| 46 | + "user_input": user_input.strip(), |
| 47 | + "history_context": history_context.strip() or "(none)", |
| 48 | + } |
| 49 | + if render_mode in ("svg", "chartjs"): |
| 50 | + format_kwargs["render_type"] = render_mode |
| 51 | + |
| 52 | + user_prompt = user_template.format(**format_kwargs) |
| 53 | + |
| 54 | + chunks: list[str] = [] |
| 55 | + async for chunk in self.stream_llm( |
| 56 | + user_prompt=user_prompt, |
| 57 | + system_prompt=system_prompt, |
| 58 | + response_format={"type": "json_object"}, |
| 59 | + stage="analyzing", |
| 60 | + trace_meta=build_trace_metadata( |
| 61 | + call_id=new_call_id("viz-analysis"), |
| 62 | + phase="analyzing", |
| 63 | + label="Visualization analysis", |
| 64 | + call_kind="viz_analysis", |
| 65 | + trace_role="analyze", |
| 66 | + trace_kind="llm_output", |
| 67 | + ), |
| 68 | + ): |
| 69 | + chunks.append(chunk) |
| 70 | + response = "".join(chunks) |
| 71 | + result = VisualizationAnalysis.model_validate(extract_json_object(response)) |
| 72 | + if render_mode in ("svg", "chartjs"): |
| 73 | + result.render_type = render_mode # type: ignore[assignment] |
| 74 | + return result |
0 commit comments