feat: enhance ARAG framework with performance tracking and visualization#4
Open
markson14 wants to merge 1 commit into
Open
feat: enhance ARAG framework with performance tracking and visualization#4markson14 wants to merge 1 commit into
markson14 wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds first-class performance instrumentation to the A-RAG agent runtime and introduces text-based visualization utilities to analyze trajectories and perf metrics.
Changes:
- Introduce
QueryPerfTracker/BatchPerfAggregatorfor per-query and batch performance reporting. - Add visualization helpers (
TrajectoryVisualizer,PerfDashboard, JSONL loaders/comparison) and document them. - Refactor
LLMClient.chat()with retry logic and phase timing metadata; updateBaseAgentto track token budgets + perf per loop/tool call.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| src/arag/visualizer.py | New visualization and comparison utilities for trajectories/perf output. |
| src/arag/core/perf_tracker.py | New structured perf tracking + summary/distribution computation. |
| src/arag/core/llm.py | Adds retries, logging, and phase timing data to chat calls. |
| src/arag/core/init.py | Re-exports perf tracker types from arag.core. |
| src/arag/agent/base.py | Integrates perf tracking into agent loops and tool execution. |
| src/arag/init.py | Exposes new perf + visualization utilities at package top-level. |
| README.md | Marks visualization tools as complete and links docs. |
| docs/visualization.md | Adds usage documentation/examples for visualization and perf. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+194
to
+201
| 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 |
Comment on lines
+159
to
+204
| 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 | ||
|
|
Comment on lines
+163
to
+165
| request_build_t0 = time.time() | ||
| request_build_duration = time.time() - request_build_t0 | ||
|
|
| if self.verbose: | ||
| print(f"LLM error: {e}") | ||
| logger.exception("LLM call failed in loop %s", loop_count) | ||
| break |
Comment on lines
260
to
266
| 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, | ||
| } |
Comment on lines
+4
to
+21
| 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 | ||
|
|
||
|
|
Comment on lines
+252
to
+264
| 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) |
Comment on lines
+14
to
+27
| ```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 | ||
| ) |
- 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
Author
|
@copilot apply changes based on the comments in this thread |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
btw, great work. I replaced your ARAG with my previous GraphRAG setup, which used to require a lot of meticulous, hand-crafted indexing and graph construction. Easy indexing could give a great boost on my datasets. More broadly, I think the old predefined, workflow-like RAG approach just doesn’t fit modern agent systems anymore. We now need more agentic tools that can adapt to how agents actually operate.
Looking forward to your future upgrades !