From 108e87ee87abc98179e8b35eaa9c0bd9d613a04d Mon Sep 17 00:00:00 2001 From: markson14zhang Date: Mon, 13 Apr 2026 16:23:25 +0000 Subject: [PATCH] feat: enhance ARAG framework with performance tracking and visualization - Add QueryPerfTracker for detailed performance monitoring - Add TrajectoryVisualizer and PerfDashboard for result analysis - Refactor LLM client with enhanced error handling - Update agent with improved token budget management --- README.md | 2 +- docs/visualization.md | 344 +++++++++++++++++ src/arag/__init__.py | 18 +- src/arag/agent/base.py | 205 +++++++--- src/arag/core/__init__.py | 4 +- src/arag/core/llm.py | 117 ++++-- src/arag/core/perf_tracker.py | 704 ++++++++++++++++++++++++++++++++++ src/arag/visualizer.py | 296 ++++++++++++++ 8 files changed, 1597 insertions(+), 93 deletions(-) create mode 100644 docs/visualization.md create mode 100644 src/arag/core/perf_tracker.py create mode 100644 src/arag/visualizer.py diff --git a/README.md b/README.md index dec893b..d2bd448 100644 --- a/README.md +++ b/README.md @@ -342,7 +342,7 @@ print(f"Loops: {result['loops']}") - [ ] **Ablation Interfaces**: Complete interfaces for ablation studies (w/o keyword search, w/o semantic search, w/o chunk read) - [ ] **Multi-Provider Support**: Native API support for Anthropic Claude and Google Gemini (currently only OpenAI-compatible APIs) - [ ] **Additional Benchmarks**: Scripts for HotpotQA, 2WikiMQA, and GraphRAG-Bench evaluation -- [ ] **Visualization Tools**: Trajectory visualization and analysis tools +- [x] **Visualization Tools**: Trajectory visualization and analysis tools (see [docs/visualization.md](docs/visualization.md)) Contributions and feedback are welcome! diff --git a/docs/visualization.md b/docs/visualization.md new file mode 100644 index 0000000..7c866a1 --- /dev/null +++ b/docs/visualization.md @@ -0,0 +1,344 @@ +# Visualization Tools + +A-RAG provides comprehensive visualization and analysis tools for agent execution trajectories and performance metrics. + +## Features + +- **Trajectory Visualization**: ASCII art rendering of agent execution steps +- **Performance Dashboard**: Detailed timing, token, and cost breakdowns +- **Batch Comparison**: Side-by-side comparison of multiple query results +- **Flexible Output**: Both text-based and structured data formats + +## Quick Start + +```python +from arag import LLMClient, BaseAgent, ToolRegistry +from arag.tools.keyword_search import KeywordSearchTool +from arag.tools.read_chunk import ReadChunkTool +from arag.visualizer import visualize_result, compare_results + +# Setup agent (with perf_log=True to print performance report) +agent = BaseAgent( + llm_client=client, + tools=tools, + max_loops=15, + max_token_budget=128000, + perf_log=True, # Enable performance tracking +) + +# Run query - result includes 'perf' field +result = agent.run("Your question here?") + +# Visualize single result +print(visualize_result(result)) +``` + +## Output Format + +Each `agent.run()` result now includes a `perf` field containing: + +```python +{ + "answer": "...", + "trajectory": [...], + "total_cost": 0.005, + "loops": 3, + "perf": { + "summary": { + "total_duration": 2.45, + "llm_duration": 1.82, + "tool_duration": 0.63, + "llm_call_count": 4, + "tool_call_count": 3, + "total_input_tokens": 12500, + "total_output_tokens": 3200, + "total_retrieved_tokens": 4850, + "total_cost": 0.005, + "llm_duration_distribution": {"mean": 0.45, "p50": 0.42, "p90": 0.55, ...}, + "tool_duration_distribution": {"mean": 0.21, "p50": 0.18, "p90": 0.32, ...}, + "context_token_records": [ + {"loop": 1, "message_tokens": 3200, "system_prompt_tokens": 800, "history_tokens": 0, "latest_tool_tokens": 0}, + {"loop": 2, "message_tokens": 5800, "system_prompt_tokens": 800, "history_tokens": 2400, "latest_tool_tokens": 1800}, + ... + ] + }, + "llm_calls": [...], + "tool_calls": [...] + } +} +``` + +## Visualization Examples + +### Single Result Visualization + +Given a result with performance data, `visualize_result()` produces: + +``` +====================================================================== + AGENT EXECUTION TRAJECTORY +====================================================================== +─── Loop 1 ─── + [ 1] keyword_search OK | keywords=['institute', 'Collegian'] + Tokens: 1,213 | Result: Chunk ID: 1383, Match Score: 11.840... + [ 2] read_chunk OK | chunk_ids=['58'] + Tokens: 402 | Result: [Chunk 58] Evidence Excerpt: ... + +─── Loop 2 ─── + [ 3] keyword_search OK | keywords=['Collegian'] + Tokens: 1,850 | Result: Chunk ID: 58, Match Score: 10.838... + + TOOL USAGE STATISTICS +-------------------------------------------------- + Tool Calls Tokens Time(s) +-------------------------------------------------- + keyword_search 2 3,063 0.45 + read_chunk 1 402 0.18 + +====================================================================== + + PERFORMANCE SUMMARY +====================================================================== + Duration: 2.45s | LLM: 1.82s | Tools: 0.63s + LLM Calls: 4 | Tool Calls: 3 | Cost: $0.005000 + Tokens - Input: 12,500 | Output: 3,200 | Retrieved: 4,850 + + LLM CALL DURATIONS +---------------------------------------- + Mean: 0.455s | P50: 0.420s | P90: 0.550s + + TOOL CALL DURATIONS +---------------------------------------- + Mean: 0.210s | P50: 0.180s | P90: 0.320s + + CONTEXT TOKEN DISTRIBUTION +------------------------------------------------------------ + Loop Total System History Latest +------------------------------------------------------------ + 1 3,200 800 0 0 + 2 5,800 800 2,400 1,800 + 3 9,200 800 5,400 2,600 +``` + +### Batch Comparison + +Compare multiple results side-by-side: + +```python +from arag.visualizer import compare_results + +print(compare_results(results)) +``` + +Output: +``` +====================================================================== + COMPARISON VIEW +====================================================================== + Metric Result 1 Result 2 Result 3 +---------------------------------------------------------------------- + Duration (s) 2.45 3.21 1.98 + LLM Calls 4 5 3 + Tool Calls 3 4 2 + Cost ($) 0.005000 0.006500 0.003800 + Input Tokens 12,500 15,300 9,800 + Output Tokens 3,200 4,100 2,600 + Retrieved Tokens 4,850 6,200 3,400 +====================================================================== +``` + +## Data Structures + +### TrajectoryVisualizer + +Renders agent execution as ASCII art. + +```python +from arag.visualizer import TrajectoryVisualizer + +viz = TrajectoryVisualizer(trajectory=result["trajectory"], perf_data=result["perf"]) + +viz.render_text() # ASCII trajectory +viz.render_timeline() # Timeline view with duration bars +viz.render_stats() # Per-tool statistics +``` + +### PerfDashboard + +Renders performance metrics. + +```python +from arag.visualizer import PerfDashboard + +dashboard = PerfDashboard(perf_data=result["perf"]) + +dashboard.render_summary() # Overall summary +dashboard.render_timing_breakdown() # LLM/tool timing distributions +dashboard.render_context_tokens() # Context token per loop +``` + +### QueryPerfTracker + +Collects per-query performance data (used internally by `BaseAgent`). + +```python +from arag.core.perf_tracker import QueryPerfTracker + +tracker = QueryPerfTracker(query=question, model="gpt-4o-mini") +tracker.record_llm_call(loop=1, duration=0.5, input_tokens=1000, ...) +tracker.record_tool_call(loop=1, tool_name="keyword_search", duration=0.2, ...) +tracker.finish() + +perf_data = tracker.to_dict() +tracker.log_report() +``` + +### BatchPerfAggregator + +Aggregates performance across multiple queries. + +```python +from arag.core.perf_tracker import BatchPerfAggregator + +aggregator = BatchPerfAggregator() +for result in results: + aggregator.add_query_report(result["perf"]) + +batch_stats = aggregator.to_dict() +print(aggregator.format_report()) +``` + +## Loading Results + +Load results from JSONL files for analysis: + +```python +from arag.visualizer import load_results + +results = load_results("path/to/results.jsonl") +for r in results: + print(visualize_result(r)) +``` + +## Integration with BaseAgent + +The `BaseAgent` automatically tracks performance when `perf_log=False` (default): + +```python +agent = BaseAgent( + llm_client=client, + tools=tools, + perf_log=False, # Don't print report, just collect data +) + +result = agent.run(query) +# result["perf"] contains structured performance data +``` + +Set `perf_log=True` to also print a detailed performance report after each query: + +```python +agent = BaseAgent( + llm_client=client, + tools=tools, + perf_log=True, # Print report after each query +) +``` + +## Performance Report Example + +When `perf_log=True`, the agent prints after each query: + +``` +==================================================================== + ARAG Performance Report +==================================================================== + Query : When was the institute that owned The Collegian founded? + Model : gpt-4o-mini + Total : 2.45s | LLM calls: 4 | Tool calls: 3 | Cost: $0.005000 + +------------------------------------------------------------------------ + Timing Breakdown +-------------------------------------------------------------------- + LLM : 1.82s ( 74.3%) | 4 calls | avg 0.45s + Tools : 0.63s ( 25.7%) | 3 calls | avg 0.21s + Other : 0.00s ( 0.0%) | overhead + +-------------------------------------------------------------------- + LLM Call Details +-------------------------------------------------------------------- + # Loop Duration In Tok Out Tok Cost Tools? Forced? + 1 1 0.42s 2500 820 $0.000820 Yes No + 2 1 0.48s 3200 890 $0.000980 Yes No + 3 2 0.45s 4100 750 $0.000900 Yes No + 4 3 0.47s 5700 740 $0.000900 No FORCED + +-------------------------------------------------------------------- + Tool Call Details +-------------------------------------------------------------------- + # Loop Tool Duration Ret Tok Status + 1 1 keyword_search 0.25s 1213 OK + 2 1 read_chunk 0.18s 402 OK + 3 2 keyword_search 0.20s 1850 OK + +-------------------------------------------------------------------- + Tool Distribution +-------------------------------------------------------------------- + Tool Calls Total Time Avg Time Tokens Errors + keyword_search 2 0.45s 0.22s 3063 0 + read_chunk 1 0.18s 0.18s 402 0 + +-------------------------------------------------------------------- + Token Summary +-------------------------------------------------------------------- + LLM Input : 12,500 tokens + LLM Output : 3,200 tokens + LLM Total : 15,700 tokens + Retrieved : 4,850 tokens + Total Cost : $0.005000 + +==================================================================== +``` + +## Use Cases + +### 1. Debug Agent Behavior + +Analyze why certain queries fail or take too long: + +```python +results = load_results("results.jsonl") +for r in results: + if r["loops"] > 10: + print(visualize_result(r)) +``` + +### 2. Compare Strategies + +Compare different retrieval strategies: + +```python +# Compare results with different max_loops +agent_fast = BaseAgent(..., max_loops=5) +agent_deep = BaseAgent(..., max_loops=20) + +# ... run queries ... + +print(compare_results([result_fast, result_deep], labels=["fast", "deep"])) +``` + +### 3. Cost Analysis + +Track API costs across queries: + +```python +from arag.core.perf_tracker import BatchPerfAggregator + +aggregator = BatchPerfAggregator() +for r in load_results("results.jsonl"): + aggregator.add_query_report(r["perf"]) + +stats = aggregator.to_dict() +print(f"Average cost per query: ${stats['summary']['avg_cost_per_query']:.6f}") +print(f"Total cost: ${stats['summary']['total_cost']:.6f}") +``` \ No newline at end of file diff --git a/src/arag/__init__.py b/src/arag/__init__.py index 0e19547..b795c61 100644 --- a/src/arag/__init__.py +++ b/src/arag/__init__.py @@ -5,16 +5,32 @@ from arag.core.config import Config from arag.core.context import AgentContext from arag.core.llm import LLMClient +from arag.core.perf_tracker import QueryPerfTracker +from arag.core.perf_tracker import BatchPerfAggregator from arag.agent.base import BaseAgent from arag.tools.base import BaseTool from arag.tools.registry import ToolRegistry +from arag.visualizer import ( + TrajectoryVisualizer, + PerfDashboard, + visualize_result, + compare_results, + load_results, +) __all__ = [ "Config", - "AgentContext", + "AgentContext", "LLMClient", + "QueryPerfTracker", + "BatchPerfAggregator", "BaseAgent", "BaseTool", "ToolRegistry", + "TrajectoryVisualizer", + "PerfDashboard", + "visualize_result", + "compare_results", + "load_results", "__version__", ] diff --git a/src/arag/agent/base.py b/src/arag/agent/base.py index 47d0e4a..76d7ad7 100644 --- a/src/arag/agent/base.py +++ b/src/arag/agent/base.py @@ -1,18 +1,23 @@ """Base agent implementation for ARAG.""" import json -from typing import Any, Dict, List, Optional +import time +import logging +from typing import Any, Dict, List import tiktoken from arag.core.context import AgentContext from arag.core.llm import LLMClient +from arag.core.perf_tracker import QueryPerfTracker from arag.tools.registry import ToolRegistry +logger = logging.getLogger("arag.agent") + class BaseAgent: """Base agent with tool calling capabilities.""" - + def __init__( self, llm_client: LLMClient, @@ -21,6 +26,7 @@ def __init__( max_loops: int = 10, max_token_budget: int = 128000, verbose: bool = False, + perf_log: bool = False, ): self.llm = llm_client self.tools = tools @@ -28,8 +34,16 @@ def __init__( self.max_loops = max_loops self.max_token_budget = max_token_budget self.verbose = verbose + self.perf_log = perf_log self.tokenizer = tiktoken.encoding_for_model("gpt-4o") - + + def _finalize_perf(self, tracker: QueryPerfTracker) -> Dict[str, Any]: + """Finalize perf tracking and return structured stats.""" + tracker.finish() + if self.perf_log: + tracker.log_report() + return tracker.to_dict() + def _calculate_message_tokens(self, messages: List[Dict[str, Any]]) -> int: total = len(self.tokenizer.encode(self.system_prompt)) for msg in messages: @@ -37,23 +51,44 @@ def _calculate_message_tokens(self, messages: List[Dict[str, Any]]) -> int: if content: total += len(self.tokenizer.encode(str(content))) return total - - def _force_final_answer(self, messages: List[Dict[str, Any]], context: AgentContext, - total_cost: float, reason: str) -> tuple: + + def _force_final_answer( + self, + messages: List[Dict[str, Any]], + context: AgentContext, + total_cost: float, + reason: str, + tracker: QueryPerfTracker = None, + loop_count: int = 0, + ) -> tuple: """Force the model to give a final answer when limits are reached.""" force_prompt = ( "You have reached the limit. " "You MUST now provide a final answer based on the information you have gathered so far. " "Do NOT call any more tools. Synthesize the available information and respond directly." ) - + messages.append({"role": "user", "content": force_prompt}) - + try: + forced_t0 = time.time() response = self.llm.chat(messages=messages, tools=None, temperature=0.0) + forced_llm_duration = time.time() - forced_t0 total_cost += response["cost"] final_answer = response["message"].get("content", "") - + + if tracker is not None: + tracker.record_llm_call( + loop=loop_count, + duration=forced_llm_duration, + input_tokens=response.get("input_tokens", 0), + output_tokens=response.get("output_tokens", 0), + cost=response.get("cost", 0.0), + has_tool_calls=False, + forced_final=True, + phase_durations=response.get("phase_durations"), + ) + if self.verbose: print(f"Forced answer: {final_answer[:200]}...") print(f"Total cost: ${total_cost:.6f}") @@ -61,130 +96,196 @@ def _force_final_answer(self, messages: List[Dict[str, Any]], context: AgentCont if self.verbose: print(f"Error getting forced answer: {e}") final_answer = f"Error: {reason} and failed to generate final answer." - + return final_answer, total_cost - + def run(self, query: str) -> Dict[str, Any]: context = AgentContext() + tracker = QueryPerfTracker(query=query, model=self.llm.model) messages = [ {"role": "system", "content": self.system_prompt}, {"role": "user", "content": query}, ] - + trajectory = [] total_cost = 0.0 loop_count = 0 tool_schemas = self.tools.get_all_schemas() - + if self.verbose: - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Question: {query}") - print(f"{'='*60}\n") - + print(f"{'=' * 60}\n") + for loop_idx in range(self.max_loops): loop_count = loop_idx + 1 - + current_tokens = self._calculate_message_tokens(messages) if current_tokens > self.max_token_budget: if self.verbose: - print(f"Token budget exceeded ({current_tokens} > {self.max_token_budget}), forcing answer...") - + print( + f"Token budget exceeded ({current_tokens} > {self.max_token_budget}), forcing answer..." + ) + final_answer, total_cost = self._force_final_answer( - messages, context, total_cost, "Token budget exceeded" + messages, + context, + total_cost, + "Token budget exceeded", + tracker=tracker, + loop_count=loop_count, ) - + + perf_data = self._finalize_perf(tracker) + return { "answer": final_answer, "trajectory": trajectory, "total_cost": total_cost, "loops": loop_count, "token_budget_exceeded": True, - **context.get_summary() + "perf": perf_data, + **context.get_summary(), } - + if self.verbose: - print(f"Loop {loop_count}/{self.max_loops} (Tokens: {current_tokens}/{self.max_token_budget})") - + print( + f"Loop {loop_count}/{self.max_loops} (Tokens: {current_tokens}/{self.max_token_budget})" + ) + + system_tokens = len(self.tokenizer.encode(self.system_prompt)) + latest_tool_tokens = 0 + for msg in reversed(messages): + if msg.get("role") in ("tool", "user") and msg.get("content"): + latest_tool_tokens = len(self.tokenizer.encode(str(msg["content"]))) + break + history_tokens = max(current_tokens - system_tokens - latest_tool_tokens, 0) + tracker.record_context_tokens( + loop=loop_count, + message_tokens=current_tokens, + system_prompt_tokens=system_tokens, + history_tokens=history_tokens, + latest_tool_tokens=latest_tool_tokens, + ) + + llm_t0 = time.time() try: response = self.llm.chat(messages=messages, tools=tool_schemas) except Exception as e: if self.verbose: print(f"LLM error: {e}") + logger.exception("LLM call failed in loop %s", loop_count) break - + llm_duration = time.time() - llm_t0 + total_cost += response["cost"] message = response["message"] messages.append(message) - + + tracker.record_llm_call( + loop=loop_count, + duration=llm_duration, + input_tokens=response.get("input_tokens", 0), + output_tokens=response.get("output_tokens", 0), + cost=response.get("cost", 0.0), + has_tool_calls=bool(message.get("tool_calls")), + phase_durations=response.get("phase_durations"), + ) + if self.verbose and message.get("content"): print(f"Assistant: {message['content'][:200]}...") - + tool_calls = message.get("tool_calls") if not tool_calls: - # No tool calls - agent is done final_answer = message.get("content", "") + perf_data = self._finalize_perf(tracker) return { "answer": final_answer, "trajectory": trajectory, "total_cost": total_cost, "loops": loop_count, - **context.get_summary() + "perf": perf_data, + **context.get_summary(), } - - # Execute tool calls + for tc in tool_calls: func_name = tc["function"]["name"] try: func_args = json.loads(tc["function"]["arguments"]) except json.JSONDecodeError: func_args = {} - + if self.verbose: print(f"Tool: {func_name}") print(f" Args: {func_args}") - + + tool_t0 = time.time() try: tool_result, tool_log = self.tools.execute(func_name, context, **func_args) + tool_success = not bool(tool_log.get("error")) except Exception as e: tool_result = f"Error executing tool: {str(e)}" tool_log = {"retrieved_tokens": 0, "error": str(e)} - + tool_success = False + tool_duration = time.time() - tool_t0 + + tracker.record_tool_call( + loop=loop_count, + tool_name=func_name, + duration=tool_duration, + retrieved_tokens=tool_log.get("retrieved_tokens", 0), + success=tool_success, + arguments=func_args, + ) + if self.verbose: - output_preview = tool_result[:300] + "..." if len(tool_result) > 300 else tool_result + output_preview = ( + tool_result[:300] + "..." if len(tool_result) > 300 else tool_result + ) print(f" Result: {output_preview}") if tool_log.get("retrieved_tokens", 0) > 0: - print(f" Tokens: {tool_log['retrieved_tokens']}") + print( + f" Tokens: {tool_log['retrieved_tokens']} | Time: {tool_duration:.2f}s" + ) print() - - messages.append({ - "role": "tool", - "tool_call_id": tc["id"], - "content": tool_result, - }) - - # Record trajectory (same format as original A-RAG) + + messages.append( + { + "role": "tool", + "tool_call_id": tc["id"], + "content": tool_result, + } + ) + traj_entry = { "loop": loop_count, "tool_name": func_name, "arguments": func_args, "tool_result": tool_result, - **tool_log # includes retrieved_tokens, chunks_found, etc. + **tool_log, } trajectory.append(traj_entry) - - # Max loops reached - force final answer + if self.verbose: print(f"Max loops reached ({self.max_loops}), forcing answer...") - + final_answer, total_cost = self._force_final_answer( - messages, context, total_cost, "Maximum loops exceeded" + messages, + context, + total_cost, + "Maximum loops exceeded", + tracker=tracker, + loop_count=loop_count, ) - + + perf_data = self._finalize_perf(tracker) + return { "answer": final_answer, "trajectory": trajectory, "total_cost": total_cost, "loops": loop_count, "max_loops_exceeded": True, - **context.get_summary() + "perf": perf_data, + **context.get_summary(), } diff --git a/src/arag/core/__init__.py b/src/arag/core/__init__.py index be43391..ace56bb 100644 --- a/src/arag/core/__init__.py +++ b/src/arag/core/__init__.py @@ -3,5 +3,7 @@ from arag.core.config import Config from arag.core.context import AgentContext from arag.core.llm import LLMClient +from arag.core.perf_tracker import QueryPerfTracker +from arag.core.perf_tracker import BatchPerfAggregator -__all__ = ["Config", "AgentContext", "LLMClient"] +__all__ = ["Config", "AgentContext", "LLMClient", "QueryPerfTracker", "BatchPerfAggregator"] diff --git a/src/arag/core/llm.py b/src/arag/core/llm.py index b10e87a..27e1f78 100644 --- a/src/arag/core/llm.py +++ b/src/arag/core/llm.py @@ -1,15 +1,19 @@ """LLM client for ARAG - unified interface for OpenAI-compatible APIs.""" import os +import time +import logging from typing import Any, Dict, List, Optional import requests import tiktoken +logger = logging.getLogger(__name__) + class LLMClient: """Unified LLM client for OpenAI-compatible APIs.""" - + # Official pricing (USD per 1M tokens): (input, cached_input, output) PRICING = { # OpenAI GPT-5 series @@ -53,7 +57,7 @@ class LLMClient: # Default fallback "default": (1.0, 0.1, 5.0), } - + def __init__( self, model: str = None, @@ -65,22 +69,26 @@ def __init__( ): self.model = model or os.getenv("ARAG_MODEL", "gpt-4o-mini") self.api_key = api_key or os.getenv("ARAG_API_KEY") - self.base_url = (base_url or os.getenv("ARAG_BASE_URL", "https://api.openai.com/v1")).rstrip('/') + self.base_url = ( + base_url or os.getenv("ARAG_BASE_URL", "https://api.openai.com/v1") + ).rstrip("/") self.temperature = temperature self.max_tokens = max_tokens self.reasoning_effort = reasoning_effort - + if not self.api_key: - raise ValueError("API key required. Set ARAG_API_KEY environment variable or pass api_key parameter.") - + raise ValueError( + "API key required. Set ARAG_API_KEY environment variable or pass api_key parameter." + ) + try: self.tokenizer = tiktoken.encoding_for_model("gpt-4o") except Exception: self.tokenizer = tiktoken.get_encoding("cl100k_base") - + def count_tokens(self, text: str) -> int: return len(self.tokenizer.encode(text)) - + def count_message_tokens(self, messages: List[Dict[str, Any]]) -> int: total = 0 for msg in messages: @@ -96,17 +104,17 @@ def count_message_tokens(self, messages: List[Dict[str, Any]]) -> int: for tc in msg["tool_calls"]: total += self.count_tokens(str(tc.get("function", {}))) return total - + def calculate_cost(self, usage: dict) -> float: """Calculate API cost from usage info.""" model_lower = self.model.lower() - - prompt_tokens = usage.get('prompt_tokens', 0) - completion_tokens = usage.get('completion_tokens', 0) - prompt_details = usage.get('prompt_tokens_details', {}) or {} - cached_tokens = prompt_details.get('cached_tokens', 0) + + prompt_tokens = usage.get("prompt_tokens", 0) + completion_tokens = usage.get("completion_tokens", 0) + prompt_details = usage.get("prompt_tokens_details", {}) or {} + cached_tokens = prompt_details.get("cached_tokens", 0) input_tokens = max(prompt_tokens - cached_tokens, 0) - + # Find matching pricing for key in self.PRICING: if key in model_lower: @@ -114,21 +122,22 @@ def calculate_cost(self, usage: dict) -> float: break else: input_price, cached_price, output_price = self.PRICING["default"] - + usd_cost = ( - (input_tokens / 1_000_000) * input_price + - (cached_tokens / 1_000_000) * cached_price + - (completion_tokens / 1_000_000) * output_price + (input_tokens / 1_000_000) * input_price + + (cached_tokens / 1_000_000) * cached_price + + (completion_tokens / 1_000_000) * output_price ) - + return round(usd_cost, 6) - + def chat( self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]] = None, temperature: float = None, max_tokens: int = None, + max_retries: int = 6, ) -> Dict[str, Any]: url = f"{self.base_url}/chat/completions" headers = { @@ -146,33 +155,65 @@ def chat( payload["tool_choice"] = "auto" if self.reasoning_effort: payload["reasoning_effort"] = self.reasoning_effort - - response = requests.post(url, headers=headers, json=payload, timeout=300) - response.raise_for_status() - result = response.json() - - usage = result.get("usage", {}) - - return { - "message": result["choices"][0]["message"], - "input_tokens": usage.get("prompt_tokens", 0), - "output_tokens": usage.get("completion_tokens", 0), - "cost": self.calculate_cost(usage), - "raw_response": result, - } - + + last_error = None + total_retry_wait: float = 0.0 + for attempt in range(max_retries): + try: + request_build_t0 = time.time() + request_build_duration = time.time() - request_build_t0 + + http_t0 = time.time() + response = requests.post(url, headers=headers, json=payload, timeout=300) + http_wait_duration = time.time() - http_t0 + + response.raise_for_status() + + parse_t0 = time.time() + result = response.json() + + if "choices" not in result or not result["choices"]: + raise KeyError(f"API returned no choices: {str(result)[:200]}") + + usage = result.get("usage", {}) + response_parse_duration = time.time() - parse_t0 + + return { + "message": result["choices"][0]["message"], + "input_tokens": usage.get("prompt_tokens", 0), + "output_tokens": usage.get("completion_tokens", 0), + "cost": self.calculate_cost(usage), + "raw_response": result, + "phase_durations": { + "request_build_duration": round(request_build_duration, 4), + "http_wait_duration": round(http_wait_duration, 4), + "response_parse_duration": round(response_parse_duration, 4), + "retry_wait_duration": round(total_retry_wait, 4), + }, + } + except (requests.exceptions.RequestException, KeyError) as e: + last_error = e + wait = min(2**attempt + 1, 30) + logger.warning( + f"LLM chat attempt {attempt + 1}/{max_retries} failed: {e}, retrying in {wait}s" + ) + time.sleep(wait) + total_retry_wait += wait + + raise last_error + def generate( self, messages: List[Dict[str, Any]], system: str = None, tools: List[Dict[str, Any]] = None, temperature: float = None, - **kwargs + **kwargs, ) -> tuple: """Generate response (compatibility method for eval script).""" if system: messages = [{"role": "system", "content": system}] + messages - + result = self.chat(messages=messages, tools=tools, temperature=temperature) content = result["message"].get("content", "") cost = result["cost"] diff --git a/src/arag/core/perf_tracker.py b/src/arag/core/perf_tracker.py new file mode 100644 index 0000000..37ce532 --- /dev/null +++ b/src/arag/core/perf_tracker.py @@ -0,0 +1,704 @@ +"""Performance tracker for ARAG query and tool call profiling. + +记录每次 query 的耗时、token 消耗和工具调用情况， +并在结束时输出结构化性能报告，为效率优化提供情报支持。 +""" + +import time +import logging +from typing import Any, Dict, List, Optional +from dataclasses import dataclass, field + +logger = logging.getLogger("arag.perf") + + +@dataclass +class ContextTokenRecord: + """单轮上下文 token 分布记录。""" + + loop: int + message_tokens: int + system_prompt_tokens: int + history_tokens: int + latest_tool_tokens: int + + +@dataclass +class LLMCallRecord: + """单次 LLM 调用的性能记录。""" + + loop: int + duration: float + input_tokens: int + output_tokens: int + cost: float + has_tool_calls: bool + forced_final: bool = False + phase_durations: Optional[Dict[str, float]] = None + + +@dataclass +class ToolCallRecord: + """单次工具调用的性能记录。""" + + loop: int + tool_name: str + duration: float + retrieved_tokens: int + success: bool + arguments: Dict[str, Any] = field(default_factory=dict) + + +class QueryPerfTracker: + """单次 query 的完整性能追踪器。 + + 在 BaseAgent.run() 中创建，收集所有 LLM 调用和工具调用的 + 耗时与 token 数据，最终输出格式化报告。 + """ + + def __init__(self, query: str, model: str = ""): + self.query = query + self.model = model + self.start_time = time.time() + self.end_time: Optional[float] = None + self.llm_calls: List[LLMCallRecord] = [] + self.tool_calls: List[ToolCallRecord] = [] + self.context_token_records: List[ContextTokenRecord] = [] + + def record_llm_call( + self, + loop: int, + duration: float, + input_tokens: int, + output_tokens: int, + cost: float, + has_tool_calls: bool, + forced_final: bool = False, + phase_durations: Optional[Dict[str, float]] = None, + ): + self.llm_calls.append( + LLMCallRecord( + loop=loop, + duration=duration, + input_tokens=input_tokens, + output_tokens=output_tokens, + cost=cost, + has_tool_calls=has_tool_calls, + forced_final=forced_final, + phase_durations=phase_durations, + ) + ) + + def record_tool_call( + self, + loop: int, + tool_name: str, + duration: float, + retrieved_tokens: int, + success: bool, + arguments: Optional[Dict[str, Any]] = None, + ): + self.tool_calls.append( + ToolCallRecord( + loop=loop, + tool_name=tool_name, + duration=duration, + retrieved_tokens=retrieved_tokens, + success=success, + arguments=arguments or {}, + ) + ) + + def record_context_tokens( + self, + loop: int, + message_tokens: int, + system_prompt_tokens: int, + history_tokens: int, + latest_tool_tokens: int, + ): + self.context_token_records.append( + ContextTokenRecord( + loop=loop, + message_tokens=message_tokens, + system_prompt_tokens=system_prompt_tokens, + history_tokens=history_tokens, + latest_tool_tokens=latest_tool_tokens, + ) + ) + + def finish(self): + self.end_time = time.time() + + @property + def total_duration(self) -> float: + end = self.end_time or time.time() + return end - self.start_time + + @property + def total_llm_duration(self) -> float: + return sum(r.duration for r in self.llm_calls) + + @property + def total_tool_duration(self) -> float: + return sum(r.duration for r in self.tool_calls) + + @property + def total_input_tokens(self) -> int: + return sum(r.input_tokens for r in self.llm_calls) + + @property + def total_output_tokens(self) -> int: + return sum(r.output_tokens for r in self.llm_calls) + + @property + def total_cost(self) -> float: + return sum(r.cost for r in self.llm_calls) + + @property + def total_retrieved_tokens(self) -> int: + return sum(r.retrieved_tokens for r in self.tool_calls) + + def _tool_distribution(self) -> Dict[str, Dict[str, Any]]: + dist: Dict[str, Dict[str, Any]] = {} + for r in self.tool_calls: + if r.tool_name not in dist: + dist[r.tool_name] = { + "count": 0, + "total_duration": 0.0, + "total_tokens": 0, + "errors": 0, + } + d = dist[r.tool_name] + d["count"] += 1 + d["total_duration"] += r.duration + d["total_tokens"] += r.retrieved_tokens + if not r.success: + d["errors"] += 1 + return dist + + @staticmethod + def _distribution(values: List[float], digits: int = 3) -> Dict[str, Any]: + if not values: + return { + "count": 0, + "sum": 0, + "mean": 0, + "min": 0, + "p50": 0, + "p90": 0, + "max": 0, + } + + sorted_values = sorted(values) + + def pick(p: float) -> float: + if len(sorted_values) == 1: + return sorted_values[0] + idx = (len(sorted_values) - 1) * p + lo = int(idx) + hi = min(lo + 1, len(sorted_values) - 1) + w = idx - lo + return sorted_values[lo] * (1 - w) + sorted_values[hi] * w + + total = sum(values) + return { + "count": len(values), + "sum": round(total, digits), + "mean": round(total / len(values), digits), + "min": round(min(values), digits), + "p50": round(pick(0.5), digits), + "p90": round(pick(0.9), digits), + "max": round(max(values), digits), + } + + def format_report(self) -> str: + sep = "=" * 68 + thin = "-" * 68 + total_dur = self.total_duration + llm_dur = self.total_llm_duration + tool_dur = self.total_tool_duration + other_dur = max(total_dur - llm_dur - tool_dur, 0) + + lines: List[str] = [] + lines.append(f"\n{sep}") + lines.append(" ARAG Performance Report") + lines.append(sep) + + query_preview = self.query[:80] + "..." if len(self.query) > 80 else self.query + lines.append(f" Query : {query_preview}") + if self.model: + lines.append(f" Model : {self.model}") + lines.append( + f" Total : {total_dur:.2f}s | " + f"LLM calls: {len(self.llm_calls)} | " + f"Tool calls: {len(self.tool_calls)} | " + f"Cost: ${self.total_cost:.6f}" + ) + + lines.append(f"\n{thin}") + lines.append(" Timing Breakdown") + lines.append(thin) + pct = lambda v: f"{v / total_dur * 100:.1f}%" if total_dur > 0 else "0.0%" + avg_llm = llm_dur / len(self.llm_calls) if self.llm_calls else 0 + avg_tool = tool_dur / len(self.tool_calls) if self.tool_calls else 0 + lines.append( + f" LLM : {llm_dur:7.2f}s ({pct(llm_dur):>6}) | " + f"{len(self.llm_calls)} calls | avg {avg_llm:.2f}s" + ) + lines.append( + f" Tools : {tool_dur:7.2f}s ({pct(tool_dur):>6}) | " + f"{len(self.tool_calls)} calls | avg {avg_tool:.2f}s" + ) + lines.append(f" Other : {other_dur:7.2f}s ({pct(other_dur):>6}) | overhead") + + if self.llm_calls: + lines.append(f"\n{thin}") + lines.append(" LLM Call Details") + lines.append(thin) + lines.append( + f" {'#':>3} {'Loop':>4} {'Duration':>8} " + f"{'In Tok':>8} {'Out Tok':>8} {'Cost':>10} {'Tools?'} {'Forced?'}" + ) + for i, r in enumerate(self.llm_calls, 1): + forced_mark = "FORCED" if r.forced_final else "No" + lines.append( + f" {i:3d} {r.loop:4d} {r.duration:7.2f}s " + f"{r.input_tokens:8d} {r.output_tokens:8d} " + f"${r.cost:9.6f} {'Yes' if r.has_tool_calls else 'No':>6} {forced_mark}" + ) + + if self.tool_calls: + lines.append(f"\n{thin}") + lines.append(" Tool Call Details") + lines.append(thin) + lines.append( + f" {'#':>3} {'Loop':>4} {'Tool':<20} " + f"{'Duration':>8} {'Ret Tok':>8} {'Status'}" + ) + for i, r in enumerate(self.tool_calls, 1): + status = "OK" if r.success else "ERR" + lines.append( + f" {i:3d} {r.loop:4d} {r.tool_name:<20} " + f"{r.duration:7.2f}s {r.retrieved_tokens:8d} {status}" + ) + + dist = self._tool_distribution() + if dist: + lines.append(f"\n{thin}") + lines.append(" Tool Distribution") + lines.append(thin) + lines.append( + f" {'Tool':<20} {'Calls':>5} {'Total Time':>10} " + f"{'Avg Time':>8} {'Tokens':>8} {'Errors':>6}" + ) + for name, d in sorted(dist.items(), key=lambda x: -x[1]["total_duration"]): + avg_t = d["total_duration"] / d["count"] if d["count"] else 0 + lines.append( + f" {name:<20} {d['count']:5d} " + f"{d['total_duration']:9.2f}s {avg_t:7.2f}s " + f"{d['total_tokens']:8d} {d['errors']:6d}" + ) + + lines.append(f"\n{thin}") + lines.append(" Token Summary") + lines.append(thin) + lines.append(f" LLM Input : {self.total_input_tokens:>10,d} tokens") + lines.append(f" LLM Output : {self.total_output_tokens:>10,d} tokens") + lines.append( + f" LLM Total : {self.total_input_tokens + self.total_output_tokens:>10,d} tokens" + ) + lines.append(f" Retrieved : {self.total_retrieved_tokens:>10,d} tokens") + lines.append(f" Total Cost : ${self.total_cost:>9.6f}") + + lines.append(sep) + return "\n".join(lines) + + def to_dict(self) -> Dict[str, Any]: + llm_call_durations = [r.duration for r in self.llm_calls] + tool_call_durations = [r.duration for r in self.tool_calls] + tool_call_tokens = [float(r.retrieved_tokens) for r in self.tool_calls] + total_tool_errors = sum(0 if r.success else 1 for r in self.tool_calls) + + return { + "query_duration": round(self.total_duration, 3), + "model": self.model, + "llm_calls": [ + { + "call_index": i + 1, + "loop": r.loop, + "duration": round(r.duration, 3), + "input_tokens": r.input_tokens, + "output_tokens": r.output_tokens, + "cost": r.cost, + "has_tool_calls": r.has_tool_calls, + "forced_final": r.forced_final, + "phase_durations": r.phase_durations, + } + for i, r in enumerate(self.llm_calls) + ], + "tool_calls": [ + { + "call_index": i + 1, + "loop": r.loop, + "tool_name": r.tool_name, + "duration": round(r.duration, 3), + "retrieved_tokens": r.retrieved_tokens, + "success": r.success, + "arguments": r.arguments, + } + for i, r in enumerate(self.tool_calls) + ], + "summary": { + "total_duration": round(self.total_duration, 3), + "llm_duration": round(self.total_llm_duration, 3), + "tool_duration": round(self.total_tool_duration, 3), + "llm_call_count": len(self.llm_calls), + "tool_call_count": len(self.tool_calls), + "total_tool_errors": total_tool_errors, + "total_input_tokens": self.total_input_tokens, + "total_output_tokens": self.total_output_tokens, + "total_retrieved_tokens": self.total_retrieved_tokens, + "total_cost": self.total_cost, + "llm_duration_distribution": self._distribution(llm_call_durations, digits=3), + "tool_duration_distribution": self._distribution(tool_call_durations, digits=3), + "tool_token_distribution": self._distribution(tool_call_tokens, digits=1), + "tool_distribution": { + name: { + "count": d["count"], + "total_duration": round(d["total_duration"], 3), + "total_tokens": d["total_tokens"], + "errors": d["errors"], + } + for name, d in self._tool_distribution().items() + }, + "context_token_records": [ + { + "loop": r.loop, + "message_tokens": r.message_tokens, + "system_prompt_tokens": r.system_prompt_tokens, + "history_tokens": r.history_tokens, + "latest_tool_tokens": r.latest_tool_tokens, + } + for r in self.context_token_records + ], + }, + } + + def log_report(self): + report = self.format_report() + if logger.isEnabledFor(logging.INFO): + logger.info(report) + else: + print(report) + + +class BatchPerfAggregator: + """批量 query 级别性能聚合器。""" + + def __init__(self): + self.query_reports: List[Dict[str, Any]] = [] + + def add_query_report(self, report: Dict[str, Any]): + if not report or not isinstance(report, dict): + return + if "summary" not in report: + return + self.query_reports.append(report) + + @staticmethod + def _percentile(values: List[float], p: float) -> float: + if not values: + return 0.0 + sorted_values = sorted(values) + if len(sorted_values) == 1: + return sorted_values[0] + idx = (len(sorted_values) - 1) * p + lower = int(idx) + upper = min(lower + 1, len(sorted_values) - 1) + weight = idx - lower + return sorted_values[lower] * (1 - weight) + sorted_values[upper] * weight + + def _tool_distribution(self) -> Dict[str, Dict[str, Any]]: + dist: Dict[str, Dict[str, Any]] = {} + for report in self.query_reports: + summary = report.get("summary", {}) + tool_dist = summary.get("tool_distribution", {}) + for name, item in tool_dist.items(): + if name not in dist: + dist[name] = { + "count": 0, + "total_duration": 0.0, + "total_tokens": 0, + "errors": 0, + } + dist[name]["count"] += item.get("count", 0) + dist[name]["total_duration"] += item.get("total_duration", 0.0) + dist[name]["total_tokens"] += item.get("total_tokens", 0) + dist[name]["errors"] += item.get("errors", 0) + return dist + + @staticmethod + def _safe_mean(values: List[float]) -> float: + if not values: + return 0.0 + return sum(values) / len(values) + + def _query_level_analysis(self) -> List[Dict[str, Any]]: + rows: List[Dict[str, Any]] = [] + for i, report in enumerate(self.query_reports, 1): + summary = report.get("summary", {}) + llm_tokens = summary.get("total_input_tokens", 0) + summary.get( + "total_output_tokens", 0 + ) + rows.append( + { + "query_index": i, + "query_duration": report.get("query_duration", 0.0), + "llm_duration": summary.get("llm_duration", 0.0), + "tool_duration": summary.get("tool_duration", 0.0), + "llm_call_count": summary.get("llm_call_count", 0), + "tool_call_count": summary.get("tool_call_count", 0), + "total_tool_errors": summary.get("total_tool_errors", 0), + "llm_input_tokens": summary.get("total_input_tokens", 0), + "llm_output_tokens": summary.get("total_output_tokens", 0), + "llm_total_tokens": llm_tokens, + "retrieved_tokens": summary.get("total_retrieved_tokens", 0), + "total_cost": summary.get("total_cost", 0.0), + } + ) + return sorted(rows, key=lambda x: x["query_duration"], reverse=True) + + def to_dict(self) -> Dict[str, Any]: + total_queries = len(self.query_reports) + if total_queries == 0: + return { + "total_queries": 0, + "summary": {}, + } + + query_durations = [r.get("query_duration", 0.0) for r in self.query_reports] + llm_durations = [r.get("summary", {}).get("llm_duration", 0.0) for r in self.query_reports] + tool_durations = [ + r.get("summary", {}).get("tool_duration", 0.0) for r in self.query_reports + ] + llm_total_tokens = [ + r.get("summary", {}).get("total_input_tokens", 0) + + r.get("summary", {}).get("total_output_tokens", 0) + for r in self.query_reports + ] + retrieved_tokens = [ + r.get("summary", {}).get("total_retrieved_tokens", 0) for r in self.query_reports + ] + + total_input_tokens = sum( + r.get("summary", {}).get("total_input_tokens", 0) for r in self.query_reports + ) + total_output_tokens = sum( + r.get("summary", {}).get("total_output_tokens", 0) for r in self.query_reports + ) + total_retrieved_tokens = sum(retrieved_tokens) + total_cost = sum(r.get("summary", {}).get("total_cost", 0.0) for r in self.query_reports) + total_llm_calls = sum( + r.get("summary", {}).get("llm_call_count", 0) for r in self.query_reports + ) + total_tool_calls = sum( + r.get("summary", {}).get("tool_call_count", 0) for r in self.query_reports + ) + total_tool_errors = sum( + r.get("summary", {}).get("total_tool_errors", 0) for r in self.query_reports + ) + + query_level = self._query_level_analysis() + tool_dist = self._tool_distribution() + + tool_level_rows: List[Dict[str, Any]] = [] + for name, item in sorted(tool_dist.items(), key=lambda x: -x[1]["total_duration"]): + calls = item.get("count", 0) + avg_duration = item.get("total_duration", 0.0) / calls if calls else 0.0 + avg_tokens = item.get("total_tokens", 0) / calls if calls else 0.0 + error_rate = item.get("errors", 0) / calls if calls else 0.0 + tool_level_rows.append( + { + "tool_name": name, + "call_count": calls, + "total_duration": round(item.get("total_duration", 0.0), 3), + "avg_duration": round(avg_duration, 3), + "total_tokens": item.get("total_tokens", 0), + "avg_tokens": round(avg_tokens, 2), + "errors": item.get("errors", 0), + "error_rate": round(error_rate, 4), + "duration_share": round( + item.get("total_duration", 0.0) / max(sum(tool_durations), 1e-9), 4 + ), + } + ) + + return { + "total_queries": total_queries, + "summary": { + "total_duration": round(sum(query_durations), 3), + "total_llm_duration": round(sum(llm_durations), 3), + "total_tool_duration": round(sum(tool_durations), 3), + "total_llm_calls": total_llm_calls, + "total_tool_calls": total_tool_calls, + "total_tool_errors": total_tool_errors, + "total_input_tokens": total_input_tokens, + "total_output_tokens": total_output_tokens, + "total_llm_tokens": total_input_tokens + total_output_tokens, + "total_retrieved_tokens": total_retrieved_tokens, + "total_cost": round(total_cost, 6), + "avg_query_duration": round(sum(query_durations) / total_queries, 3), + "avg_llm_call_per_query": round(total_llm_calls / total_queries, 3), + "avg_tool_call_per_query": round(total_tool_calls / total_queries, 3), + "avg_llm_tokens": round( + (total_input_tokens + total_output_tokens) / total_queries, 1 + ), + "avg_retrieved_tokens": round(total_retrieved_tokens / total_queries, 1), + "avg_cost_per_query": round(total_cost / total_queries, 6), + "query_duration_distribution": { + "min": round(min(query_durations), 3), + "p50": round(self._percentile(query_durations, 0.5), 3), + "p90": round(self._percentile(query_durations, 0.9), 3), + "max": round(max(query_durations), 3), + }, + "llm_duration_distribution": { + "min": round(min(llm_durations), 3), + "p50": round(self._percentile(llm_durations, 0.5), 3), + "p90": round(self._percentile(llm_durations, 0.9), 3), + "max": round(max(llm_durations), 3), + }, + "tool_duration_distribution": { + "min": round(min(tool_durations), 3), + "p50": round(self._percentile(tool_durations, 0.5), 3), + "p90": round(self._percentile(tool_durations, 0.9), 3), + "max": round(max(tool_durations), 3), + }, + "llm_token_distribution": { + "min": min(llm_total_tokens) if llm_total_tokens else 0, + "p50": round(self._percentile([float(x) for x in llm_total_tokens], 0.5), 1), + "p90": round(self._percentile([float(x) for x in llm_total_tokens], 0.9), 1), + "max": max(llm_total_tokens) if llm_total_tokens else 0, + }, + "retrieved_token_distribution": { + "min": min(retrieved_tokens) if retrieved_tokens else 0, + "p50": round(self._percentile([float(x) for x in retrieved_tokens], 0.5), 1), + "p90": round(self._percentile([float(x) for x in retrieved_tokens], 0.9), 1), + "max": max(retrieved_tokens) if retrieved_tokens else 0, + }, + "tool_distribution": tool_dist, + "cost_distribution": { + "min": round( + min( + r.get("summary", {}).get("total_cost", 0.0) for r in self.query_reports + ), + 6, + ), + "p50": round( + self._percentile( + [ + r.get("summary", {}).get("total_cost", 0.0) + for r in self.query_reports + ], + 0.5, + ), + 6, + ), + "p90": round( + self._percentile( + [ + r.get("summary", {}).get("total_cost", 0.0) + for r in self.query_reports + ], + 0.9, + ), + 6, + ), + "max": round( + max( + r.get("summary", {}).get("total_cost", 0.0) for r in self.query_reports + ), + 6, + ), + }, + }, + "query_level_analysis": query_level, + "tool_level_analysis": tool_level_rows, + } + + def format_report(self) -> str: + data = self.to_dict() + total_queries = data.get("total_queries", 0) + if total_queries == 0: + return "\n[ARAG Batch Performance] no query performance data." + + summary = data["summary"] + dist = summary.get("query_duration_distribution", {}) + llm_dist = summary.get("llm_token_distribution", {}) + ret_dist = summary.get("retrieved_token_distribution", {}) + tool_dist = summary.get("tool_distribution", {}) + + lines: List[str] = [] + lines.append("\n" + "=" * 72) + lines.append("ARAG Batch Performance Summary") + lines.append("=" * 72) + lines.append( + f"Queries: {total_queries} | Total time(sum): {summary.get('total_duration', 0):.2f}s | " + f"Total cost: ${summary.get('total_cost', 0):.6f}" + ) + lines.append( + f"LLM calls: {summary.get('total_llm_calls', 0)} | " + f"Tool calls: {summary.get('total_tool_calls', 0)}" + ) + + lines.append("-" * 72) + lines.append("Timing Distribution") + lines.append("-" * 72) + lines.append( + f"avg={summary.get('avg_query_duration', 0):.2f}s, " + f"p50={dist.get('p50', 0):.2f}s, p90={dist.get('p90', 0):.2f}s, " + f"min={dist.get('min', 0):.2f}s, max={dist.get('max', 0):.2f}s" + ) + lines.append( + f"LLM time(sum)={summary.get('total_llm_duration', 0):.2f}s, " + f"Tool time(sum)={summary.get('total_tool_duration', 0):.2f}s" + ) + + lines.append("-" * 72) + lines.append("Token Distribution") + lines.append("-" * 72) + lines.append( + f"LLM input={summary.get('total_input_tokens', 0):,}, " + f"output={summary.get('total_output_tokens', 0):,}, " + f"total={summary.get('total_llm_tokens', 0):,}" + ) + lines.append( + f"Retrieved total={summary.get('total_retrieved_tokens', 0):,}, " + f"avg/query={summary.get('avg_retrieved_tokens', 0):.1f}" + ) + lines.append( + f"LLM token/query: p50={llm_dist.get('p50', 0):.1f}, " + f"p90={llm_dist.get('p90', 0):.1f}, min={llm_dist.get('min', 0)}, max={llm_dist.get('max', 0)}" + ) + lines.append( + f"Retrieved token/query: p50={ret_dist.get('p50', 0):.1f}, " + f"p90={ret_dist.get('p90', 0):.1f}, min={ret_dist.get('min', 0)}, max={ret_dist.get('max', 0)}" + ) + + if tool_dist: + lines.append("-" * 72) + lines.append("Tool Distribution") + lines.append("-" * 72) + lines.append( + f"{'Tool':<20} {'Calls':>6} {'Time(s)':>10} {'Tokens':>10} {'Errors':>6}" + ) + for name, item in sorted( + tool_dist.items(), key=lambda x: -x[1].get("total_duration", 0.0) + ): + lines.append( + f"{name:<20} {item.get('count', 0):6d} {item.get('total_duration', 0.0):10.2f} " + f"{item.get('total_tokens', 0):10d} {item.get('errors', 0):6d}" + ) + + lines.append("=" * 72) + return "\n".join(lines) diff --git a/src/arag/visualizer.py b/src/arag/visualizer.py new file mode 100644 index 0000000..7037224 --- /dev/null +++ b/src/arag/visualizer.py @@ -0,0 +1,296 @@ +"""Visualization tools for ARAG trajectory and performance analysis.""" + +import json +from typing import Any, Dict, List, Optional +from dataclasses import dataclass + + +@dataclass +class TrajectoryNode: + """Single node in the agent trajectory.""" + + loop: int + step: int + tool_name: str + arguments: Dict[str, Any] + result_preview: str + tokens: int + duration: float + success: bool + + +class TrajectoryVisualizer: + """Visualize agent execution trajectory.""" + + def __init__(self, trajectory: List[Dict[str, Any]], perf_data: Dict[str, Any] = None): + self.trajectory = trajectory + self.perf_data = perf_data or {} + + def render_text(self) -> str: + """Render trajectory as ASCII art.""" + if not self.trajectory: + return "[Empty trajectory]" + + lines = [] + lines.append("=" * 70) + lines.append(" AGENT EXECUTION TRAJECTORY") + lines.append("=" * 70) + + prev_loop = 0 + for i, entry in enumerate(self.trajectory, 1): + loop = entry.get("loop", 0) + if loop != prev_loop: + if prev_loop != 0: + lines.append("") + lines.append(f"─── Loop {loop} ───") + prev_loop = loop + + tool_name = entry.get("tool_name", "unknown") + args = entry.get("arguments", {}) + result = entry.get("tool_result", "") + tokens = entry.get("retrieved_tokens", 0) + success = entry.get("success", True) if "success" in entry else True + + status = "OK" if success else "ERR" + result_preview = result[:100] + "..." if len(result) > 100 else result + result_preview = result_preview.replace("\n", " ") + + args_str = ", ".join(f"{k}={v}" for k, v in args.items()) if args else "" + + lines.append(f" [{i:2d}] {tool_name:<20} {status:<4} | {args_str[:40]}") + if tokens > 0: + lines.append(f" Tokens: {tokens:,} | Result: {result_preview[:50]}") + else: + lines.append(f" Result: {result_preview[:60]}") + + lines.append("=" * 70) + return "\n".join(lines) + + def render_timeline(self) -> str: + """Render timeline view of execution.""" + if not self.trajectory: + return "[Empty trajectory]" + + lines = [] + lines.append("\n TIMELINE VIEW") + lines.append("-" * 50) + + loop_groups: Dict[int, List[Dict]] = {} + for entry in self.trajectory: + loop = entry.get("loop", 0) + if loop not in loop_groups: + loop_groups[loop] = [] + loop_groups[loop].append(entry) + + for loop, entries in sorted(loop_groups.items()): + lines.append(f"\n Loop {loop}:") + for entry in entries: + tool_name = entry.get("tool_name", "?") + duration = entry.get("duration", 0) + tokens = entry.get("retrieved_tokens", 0) + bars = "█" * min(int(duration * 10), 40) if duration > 0 else "" + lines.append( + f" {tool_name:<18} {duration:5.2f}s {bars} " + f"{f'({tokens:,} tok)' if tokens > 0 else ''}" + ) + + return "\n".join(lines) + + def get_tool_stats(self) -> Dict[str, Dict[str, Any]]: + """Get per-tool statistics.""" + stats: Dict[str, Dict[str, Any]] = {} + for entry in self.trajectory: + tool_name = entry.get("tool_name", "unknown") + if tool_name not in stats: + stats[tool_name] = {"count": 0, "total_tokens": 0, "total_duration": 0.0} + stats[tool_name]["count"] += 1 + stats[tool_name]["total_tokens"] += entry.get("retrieved_tokens", 0) + stats[tool_name]["total_duration"] += entry.get("duration", 0) + return stats + + def render_stats(self) -> str: + """Render trajectory statistics.""" + stats = self.get_tool_stats() + if not stats: + return "[No statistics available]" + + lines = [] + lines.append("\n TOOL USAGE STATISTICS") + lines.append("-" * 50) + lines.append(f" {'Tool':<20} {'Calls':>6} {'Tokens':>10} {'Time(s)':>10}") + lines.append("-" * 50) + + for tool_name, data in sorted(stats.items(), key=lambda x: -x[1]["total_tokens"]): + lines.append( + f" {tool_name:<20} {data['count']:6d} " + f"{data['total_tokens']:10,d} {data['total_duration']:10.2f}" + ) + + return "\n".join(lines) + + +class PerfDashboard: + """Dashboard for performance data analysis.""" + + def __init__(self, perf_data: Dict[str, Any]): + self.perf_data = perf_data + + def render_summary(self) -> str: + """Render performance summary.""" + summary = self.perf_data.get("summary", {}) + if not summary: + return "[No performance data]" + + lines = [] + sep = "=" * 50 + + lines.append(f"\n{sep}") + lines.append(" PERFORMANCE SUMMARY") + lines.append(sep) + + lines.append( + f" Duration: {summary.get('total_duration', 0):.2f}s | " + f"LLM: {summary.get('llm_duration', 0):.2f}s | " + f"Tools: {summary.get('tool_duration', 0):.2f}s" + ) + lines.append( + f" LLM Calls: {summary.get('llm_call_count', 0)} | " + f"Tool Calls: {summary.get('tool_call_count', 0)} | " + f"Cost: ${summary.get('total_cost', 0):.6f}" + ) + lines.append( + f" Tokens - Input: {summary.get('total_input_tokens', 0):,} | " + f"Output: {summary.get('total_output_tokens', 0):,} | " + f"Retrieved: {summary.get('total_retrieved_tokens', 0):,}" + ) + + return "\n".join(lines) + + def render_timing_breakdown(self) -> str: + """Render detailed timing breakdown.""" + summary = self.perf_data.get("summary", {}) + if not summary: + return "[No performance data]" + + llm_dist = summary.get("llm_duration_distribution", {}) + tool_dist = summary.get("tool_duration_distribution", {}) + + lines = [] + lines.append(f"\n LLM CALL DURATIONS") + lines.append("-" * 40) + lines.append( + f" Mean: {llm_dist.get('mean', 0):.3f}s | " + f"P50: {llm_dist.get('p50', 0):.3f}s | " + f"P90: {llm_dist.get('p90', 0):.3f}s" + ) + + lines.append(f"\n TOOL CALL DURATIONS") + lines.append("-" * 40) + lines.append( + f" Mean: {tool_dist.get('mean', 0):.3f}s | " + f"P50: {tool_dist.get('p50', 0):.3f}s | " + f"P90: {tool_dist.get('p90', 0):.3f}s" + ) + + return "\n".join(lines) + + def render_context_tokens(self) -> str: + """Render context token distribution across loops.""" + records = self.perf_data.get("summary", {}).get("context_token_records", []) + if not records: + return "[No context token data]" + + lines = [] + lines.append(f"\n CONTEXT TOKEN DISTRIBUTION") + lines.append("-" * 60) + lines.append(f" {'Loop':>4} {'Total':>8} {'System':>8} {'History':>8} {'Latest':>8}") + lines.append("-" * 60) + + for rec in records: + lines.append( + f" {rec.get('loop', 0):4d} " + f"{rec.get('message_tokens', 0):8,d} " + f"{rec.get('system_prompt_tokens', 0):8,d} " + f"{rec.get('history_tokens', 0):8,d} " + f"{rec.get('latest_tool_tokens', 0):8,d}" + ) + + return "\n".join(lines) + + +def load_results(jsonl_path: str) -> List[Dict[str, Any]]: + """Load results from JSONL file.""" + results = [] + with open(jsonl_path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + results.append(json.loads(line)) + return results + + +def visualize_result(result: Dict[str, Any]) -> str: + """Generate complete visualization for a single result.""" + trajectory = result.get("trajectory", []) + perf_data = result.get("perf", {}) + + viz = TrajectoryVisualizer(trajectory, perf_data) + dashboard = PerfDashboard(perf_data) + + parts = [ + viz.render_text(), + viz.render_timeline(), + viz.render_stats(), + dashboard.render_summary(), + dashboard.render_timing_breakdown(), + dashboard.render_context_tokens(), + ] + + return "\n".join(filter(None, parts)) + + +def compare_results(results: List[Dict[str, Any]], labels: List[str] = None) -> str: + """Compare multiple results side by side.""" + if not results: + return "[No results to compare]" + + labels = labels or [f"Result {i + 1}" for i in range(len(results))] + + lines = [] + lines.append("\n" + "=" * 70) + lines.append(" COMPARISON VIEW") + lines.append("=" * 70) + lines.append(f" {'Metric':<25} " + " ".join(f"{l:<15}" for l in labels)) + lines.append("-" * 70) + + metrics = [] + for r in results: + perf = r.get("perf", {}).get("summary", {}) + metrics.append( + { + "duration": perf.get("total_duration", 0), + "llm_calls": perf.get("llm_call_count", 0), + "tool_calls": perf.get("tool_call_count", 0), + "cost": perf.get("total_cost", 0), + "input_tokens": perf.get("total_input_tokens", 0), + "output_tokens": perf.get("total_output_tokens", 0), + "retrieved_tokens": perf.get("total_retrieved_tokens", 0), + } + ) + + lines.append(f" {'Duration (s)':<25} " + " ".join(f"{m['duration']:<15.2f}" for m in metrics)) + lines.append(f" {'LLM Calls':<25} " + " ".join(f"{m['llm_calls']:<15d}" for m in metrics)) + lines.append(f" {'Tool Calls':<25} " + " ".join(f"{m['tool_calls']:<15d}" for m in metrics)) + lines.append(f" {'Cost ($)':<25} " + " ".join(f"{m['cost']:<15.6f}" for m in metrics)) + lines.append( + f" {'Input Tokens':<25} " + " ".join(f"{m['input_tokens']:<15,}" for m in metrics) + ) + lines.append( + f" {'Output Tokens':<25} " + " ".join(f"{m['output_tokens']:<15,}" for m in metrics) + ) + lines.append( + f" {'Retrieved Tokens':<25} " + " ".join(f"{m['retrieved_tokens']:<15,}" for m in metrics) + ) + + lines.append("=" * 70) + return "\n".join(lines)