diff --git a/.github/workflows/pre_commit.yaml b/.github/workflows/pre_commit.yaml index 60a00792..1a19fdf0 100644 --- a/.github/workflows/pre_commit.yaml +++ b/.github/workflows/pre_commit.yaml @@ -10,7 +10,7 @@ jobs: - uses: actions/checkout@v3 with: ref: ${{ github.event.pull_request.head.ref }} - - uses: actions/setup-python@v3 + - uses: actions/setup-python@v3.14 - uses: pre-commit/action@v3.0.0 - uses: EndBug/add-and-commit@v9.1.3 # Only need to try and commit if the action failed diff --git a/CHANGELOG.md b/CHANGELOG.md index 707b2117..eef38452 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [develop] +### AgentiGrid +- Count tokens of every LLM call (objective extraction, search, post-search analysis, + steering, report regeneration), with a per-phase breakdown in the journal + (`token_usage`), the console summary, the PDF report and the GUI. Token totals now + include all phases; search-only totals are `token_usage.by_phase.search`. + ## [2.1] ### General diff --git a/agentigrid/README.md b/agentigrid/README.md index cb1f0ebb..91b98bd0 100644 --- a/agentigrid/README.md +++ b/agentigrid/README.md @@ -265,6 +265,34 @@ agentigrid ./data/exago/examples/case_ACTIVSg200.m \ In stress test mode, the LLM always uses fresh mode (each contingency tested independently from the base case), starts with the most loaded lines, and can escalate to N-2 combinations. The post-search report ranks contingencies by severity: infeasibility > voltage violations > high line loading > cost increase. +## Token Accounting + +Every LLM call of a session is counted (ExaGO and GridKit), each under one phase: + +| Phase | Calls | +|---|---| +| `objective_extraction` | objectives read from the goal at startup | +| `search` | every call inside the iteration loop, including explore/select, the `analyze` fallback and retries after a reply that was not JSON (all GridKit calls) | +| `analysis` | post-search analysis and goal classification (CLI, or the GUI when it runs it) | +| `steering` | objectives read from a steering directive | +| `report` | analysis regenerated for a PDF from a saved journal | +| `other` | anything else | + +The per-call record (phase, iteration, model, backend, prompt and completion tokens, and the +prompt/completion length in characters, `prompt_chars` / `completion_chars`) and the +per-phase totals are written as a top-level `token_usage` block to the journal JSON (ExaGO +`workdir/journal_*.json`, GridKit `workdir/gridkit/session_*/journal.json`) and to saved +sessions. Usage a backend did not report stays `null` per call and is counted in +`calls_without_usage`. The final console summary shows the total, the number of calls and +the breakdown, for example +`Tokens: 26,105 total (24,100 prompt + 2,005 completion) in 4 LLM calls; objective extraction 1,190, search 19,497, analysis 5,418`; +the PDF report and the GUI show the same. + +From this version on, `total_prompt_tokens` / `total_completion_tokens` and every printed +total include **all phases**. Earlier versions counted search calls only; that number is +`token_usage.by_phase.search` (sessions saved by an earlier version resume with their old +total as one `search` entry). + ## Session Save/Resume Searches can be saved to disk and resumed later — useful for long runs, interrupted sessions, or exploring different strategies from the same checkpoint. @@ -367,7 +395,7 @@ agentigrid ./data/exago/examples/case_ACTIVSg200.m "Find the maximum load scalin # With options agentigrid ./data/exago/examples/case_ACTIVSg200.m "Minimize generation cost" \ - --backend anthropic --model claude-sonnet-4-20250514 \ + --backend anthropic --model claude-sonnet-4-6 \ --app opflow --max-iter 30 --verbose # Quiet mode (only show final summary) @@ -424,7 +452,7 @@ python -m agentigrid ./data/exago/examples/case_ACTIVSg200.m "test goal" --dry-r Version 0.1.0 ============================================================ Backend: anthropic - Model: claude-sonnet-4-20250514 + Model: claude-sonnet-4-6 Application: opflow Base case: data/exago/examples/case_ACTIVSg200.m Goal: Find the maximum load scaling factor @@ -435,7 +463,7 @@ python -m agentigrid ./data/exago/examples/case_ACTIVSg200.m "test goal" --dry-r [Iter 0] Running base case simulation... [Iter 0] Base case: CONVERGED, cost=$27,557.57 -[Iter 1] Sending prompt to anthropic (claude-sonnet-4-20250514)... +[Iter 1] Sending prompt to anthropic (claude-sonnet-4-6)... [Iter 1] LLM action: modify — "Scale all loads +20%" [Iter 1] Applied 1 command(s), 0 skipped [Iter 1] Simulation completed in 0.04s — CONVERGED, cost=$33,019.55 @@ -450,7 +478,7 @@ python -m agentigrid ./data/exago/examples/case_ACTIVSg200.m "test goal" --dry-r ============================================================ Goal: Find the maximum load scaling factor Application: opflow - Backend: anthropic (claude-sonnet-4-20250514) + Backend: anthropic (claude-sonnet-4-6) Iterations: 6 (of max 20) Duration: 18.3 seconds Tokens used: ~12,450 (prompt: 9,200, completion: 3,250) diff --git a/agentigrid/agentigrid/backends/anthropic_backend.py b/agentigrid/agentigrid/backends/anthropic_backend.py index 10bbea96..a6b2ab99 100644 --- a/agentigrid/agentigrid/backends/anthropic_backend.py +++ b/agentigrid/agentigrid/backends/anthropic_backend.py @@ -6,6 +6,7 @@ import os from typing import Optional +from agentigrid.logging_setup import log_scope from agentigrid.backends.base import LLMBackend, LLMResponse from agentigrid.backends.json_extract import extract_json from agentigrid.config import LLMConfig @@ -27,6 +28,9 @@ def __init__(self, config: LLMConfig) -> None: raise ImportError("anthropic package is required for the Anthropic backend") self._config = config + # None = not yet known; True/False once a call has succeeded / the + # model has rejected the parameter (see complete()). + self.temperature_sent: Optional[bool] = None api_key = os.environ.get(config.api_key_env) if not api_key: logger.warning( @@ -40,6 +44,26 @@ def name(self) -> str: """Return the backend name.""" return "anthropic" + def _system_param(self, system_prompt: str): + """System prompt in the form the Messages API expects. + + With ``prompt_cache`` on, the system prompt is sent as a single text block + carrying an ephemeral ``cache_control`` breakpoint. AgentiGrid builds the + system prompt once per session, so every later iteration reads it from the + cache instead of paying the full input price again. This affects billing + only; the model sees identical text either way. Prompts below the model's + minimum cacheable length are silently processed uncached by the API. + """ + if not getattr(self._config, "prompt_cache", False): + return system_prompt + return [ + { + "type": "text", + "text": system_prompt, + "cache_control": {"type": "ephemeral"}, + } + ] + def supports_json_mode(self) -> bool: """Anthropic does not have a native JSON output mode.""" return False @@ -54,23 +78,50 @@ def complete( temp = temperature if temperature is not None else self._config.temperature try: - response = self._client.messages.create( + kwargs = dict( model=self._config.model, max_tokens=self._config.max_tokens, - temperature=temp, - system=system_prompt, + system=self._system_param(system_prompt), messages=[{"role": "user", "content": user_prompt}], ) - - raw_text = "" - for block in response.content: - if block.type == "text": - raw_text += block.text - - prompt_tokens = getattr(response.usage, "input_tokens", None) - completion_tokens = getattr(response.usage, "output_tokens", None) - - json_data, json_error = extract_json(raw_text) + if self.temperature_sent is not False: + kwargs["temperature"] = temp + with log_scope(logger, "LLM API request", backend=self.name(), model=self._config.model): + try: + response = self._client.messages.create(**kwargs) + if "temperature" in kwargs: + self.temperature_sent = True + except Exception as exc: + # Newer models reject the sampling-temperature parameter + # ("`temperature` is deprecated for this model"). Retry once + # without it and stop sending it for the rest of the session; + # the journal records temperature_sent=False. + if "temperature" in kwargs and "temperature" in str(exc).lower(): + logger.warning("Model %s rejects `temperature`; retrying without it.", self._config.model) + kwargs.pop("temperature") + self.temperature_sent = False + response = self._client.messages.create(**kwargs) + else: + raise + + with log_scope(logger, "LLM response parsing", backend=self.name(), model=self._config.model): + raw_text = "" + content_types = [] + for block in response.content: + content_types.append(getattr(block, "type", None)) + if block.type == "text": + raw_text += block.text + stop_reason = getattr(response, "stop_reason", None) + if stop_reason == "max_tokens": + logger.warning("Response stopped at max_tokens=%s (blocks: %s, %d text chars)", + self._config.max_tokens, content_types, len(raw_text)) + + prompt_tokens, cache_creation, cache_read = _input_token_breakdown( + response.usage + ) + completion_tokens = getattr(response.usage, "output_tokens", None) + + json_data, json_error = extract_json(raw_text) return LLMResponse( raw_text=raw_text, @@ -80,6 +131,10 @@ def complete( backend=self.name(), prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, + cache_creation_tokens=cache_creation, + cache_read_tokens=cache_read, + stop_reason=stop_reason, + content_types=content_types, ) except Exception as exc: @@ -93,4 +148,23 @@ def complete( backend=self.name(), prompt_tokens=None, completion_tokens=None, + api_error=True, ) + + +def _input_token_breakdown(usage) -> tuple[Optional[int], Optional[int], Optional[int]]: + """Return (total_input, cache_creation, cache_read) from an API usage object. + + With caching, the API's ``input_tokens`` counts only tokens after the last + cache breakpoint. The total input is ``input_tokens + cache_creation_input_tokens + + cache_read_input_tokens``; returning that total keeps ``prompt_tokens`` + comparable with runs made before caching was enabled. + """ + base = getattr(usage, "input_tokens", None) + creation = getattr(usage, "cache_creation_input_tokens", None) + read = getattr(usage, "cache_read_input_tokens", None) + creation = creation if isinstance(creation, int) else None + read = read if isinstance(read, int) else None + if not isinstance(base, int): + return None, creation, read + return base + (creation or 0) + (read or 0), creation, read diff --git a/agentigrid/agentigrid/backends/base.py b/agentigrid/agentigrid/backends/base.py index 32c4581a..d2a34d38 100644 --- a/agentigrid/agentigrid/backends/base.py +++ b/agentigrid/agentigrid/backends/base.py @@ -18,6 +18,20 @@ class LLMResponse: backend: str prompt_tokens: Optional[int] completion_tokens: Optional[int] + # Prompt-cache accounting (Anthropic). ``prompt_tokens`` always reports the + # TOTAL input tokens (uncached + cache writes + cache reads), so its meaning + # is the same with or without caching; these break that total down. + cache_creation_tokens: Optional[int] = None + cache_read_tokens: Optional[int] = None + # True when the request itself failed (network, auth, invalid request...), + # as opposed to the model answering with unparseable text. Lets the agent + # loop record API failures separately from rejected proposals. + api_error: bool = False + # Why generation stopped ("end_turn", "max_tokens", ...) and the content + # block types returned, when the backend reports them. A "max_tokens" stop + # is an output-cap limit, not a model proposal, and is recorded as such. + stop_reason: Optional[str] = None + content_types: Optional[list] = None class LLMBackend(ABC): diff --git a/agentigrid/agentigrid/backends/ollama_backend.py b/agentigrid/agentigrid/backends/ollama_backend.py index f7fb1541..88914719 100644 --- a/agentigrid/agentigrid/backends/ollama_backend.py +++ b/agentigrid/agentigrid/backends/ollama_backend.py @@ -5,6 +5,7 @@ import logging from typing import Optional +from agentigrid.logging_setup import log_scope from agentigrid.backends.base import LLMBackend, LLMResponse from agentigrid.backends.json_extract import extract_json from agentigrid.config import LLMConfig @@ -83,14 +84,16 @@ def complete( kwargs["format"] = "json" try: - response = self._client.chat(**kwargs) + with log_scope(logger, "LLM API request", backend=self.name(), model=self._config.model): + response = self._client.chat(**kwargs) - raw_text = response.get("message", {}).get("content", "") + with log_scope(logger, "LLM response parsing", backend=self.name(), model=self._config.model): + raw_text = response.get("message", {}).get("content", "") - prompt_tokens = response.get("prompt_eval_count") - completion_tokens = response.get("eval_count") + prompt_tokens = response.get("prompt_eval_count") + completion_tokens = response.get("eval_count") - json_data, json_error = extract_json(raw_text) + json_data, json_error = extract_json(raw_text) return LLMResponse( raw_text=raw_text, @@ -120,4 +123,5 @@ def complete( backend=self.name(), prompt_tokens=None, completion_tokens=None, + api_error=True, ) diff --git a/agentigrid/agentigrid/backends/openai_backend.py b/agentigrid/agentigrid/backends/openai_backend.py index 0b4261ab..9e6d2c5b 100644 --- a/agentigrid/agentigrid/backends/openai_backend.py +++ b/agentigrid/agentigrid/backends/openai_backend.py @@ -6,6 +6,7 @@ import os from typing import Optional +from agentigrid.logging_setup import log_scope from agentigrid.backends.base import LLMBackend, LLMResponse from agentigrid.backends.json_extract import extract_json from agentigrid.config import LLMConfig @@ -75,13 +76,15 @@ def complete( if use_json_mode: kwargs["response_format"] = {"type": "json_object"} - response = self._client.chat.completions.create(**kwargs) + with log_scope(logger, "LLM API request", backend=self.name(), model=self._config.model): + response = self._client.chat.completions.create(**kwargs) - raw_text = response.choices[0].message.content or "" - prompt_tokens = getattr(response.usage, "prompt_tokens", None) - completion_tokens = getattr(response.usage, "completion_tokens", None) + with log_scope(logger, "LLM response parsing", backend=self.name(), model=self._config.model): + raw_text = response.choices[0].message.content or "" + prompt_tokens = getattr(response.usage, "prompt_tokens", None) + completion_tokens = getattr(response.usage, "completion_tokens", None) - json_data, json_error = extract_json(raw_text) + json_data, json_error = extract_json(raw_text) return LLMResponse( raw_text=raw_text, @@ -104,4 +107,5 @@ def complete( backend=self.name(), prompt_tokens=None, completion_tokens=None, + api_error=True, ) diff --git a/agentigrid/agentigrid/config.py b/agentigrid/agentigrid/config.py index 7f0cb8d5..bad5207e 100644 --- a/agentigrid/agentigrid/config.py +++ b/agentigrid/agentigrid/config.py @@ -36,13 +36,14 @@ }, "llm": { "backend": "anthropic", - "model": "claude-sonnet-4-20250514", + "model": "claude-sonnet-4-6", "api_key_env": "ANTHROPIC_API_KEY", "openai_base_url": None, "ollama_host": "http://localhost:11434", "ollama_cloud_host": None, "temperature": 0.3, "max_tokens": 4096, + "prompt_cache": True, }, "search": { "max_iterations": 20, @@ -159,6 +160,9 @@ class LLMConfig: ollama_cloud_host: Optional[str] temperature: float max_tokens: int + # Anthropic prompt caching of the (per-session constant) system prompt. + # Billing-only: does not change model outputs. Ignored by other backends. + prompt_cache: bool = True @dataclass(frozen=True) @@ -374,6 +378,15 @@ def load_config( target[parts[-1]] = value logger.debug("Applied CLI overrides: %s", cli_overrides) + # AGENTIGRID_MAX_TOKENS overrides the output-token cap (e.g. to fix it for a + # batch of runs); the value used is recorded in the journal's llm_settings. + _mt = os.environ.get("AGENTIGRID_MAX_TOKENS", "").strip() + if _mt: + try: + merged["llm"] = {**merged["llm"], "max_tokens": int(_mt)} + except ValueError: + logger.warning("Ignoring non-integer AGENTIGRID_MAX_TOKENS=%r", _mt) + # Build frozen dataclass sections exago_path_fields = { "binary_dir", "opflow_binary", "scopflow_binary", "tcopflow_binary", diff --git a/agentigrid/agentigrid/exago_engine/agent_loop.py b/agentigrid/agentigrid/exago_engine/agent_loop.py index 20e59c1e..354279b8 100644 --- a/agentigrid/agentigrid/exago_engine/agent_loop.py +++ b/agentigrid/agentigrid/exago_engine/agent_loop.py @@ -17,9 +17,11 @@ from pathlib import Path from typing import Callable, Optional +from agentigrid.logging_setup import log_scope from agentigrid.backends import create_backend from agentigrid.backends.base import LLMBackend, LLMResponse from agentigrid.config import AppConfig +from agentigrid.token_usage import TokenLedger from agentigrid.exago_engine.commands import parse_command from agentigrid.exago_engine.executor import SimulationExecutor, SimulationResult from agentigrid.exago_engine.explore import ( @@ -30,7 +32,7 @@ compute_pareto_labels, format_variant_results, ) -from agentigrid.exago_engine.journal import JournalEntry, ObjectiveEntry, SearchJournal +from agentigrid.exago_engine.journal import JournalEntry, ObjectiveEntry, SearchJournal, format_iteration_count from agentigrid.exago_engine.metric_extractor import available_metrics, available_metrics_for_app, extract_all_metrics from agentigrid.exago_engine.modifier import apply_modifications, build_index_maps from agentigrid.exago_engine import sweep_metrics @@ -66,6 +68,9 @@ def _count_scenario_rows(scenario_path: Path) -> int: logger = logging.getLogger("agentigrid.exago_engine.agent_loop") +# Cap on the raw model response kept for a discarded iteration (journal telemetry). +DISCARDED_RESPONSE_MAX = 4000 + _MAX_CONSECUTIVE_PARSE_FAILURES = 3 @@ -736,8 +741,9 @@ class SearchSession: end_time: Optional[str] = None termination_reason: str = "" final_findings: Optional[dict] = None - total_prompt_tokens: int = 0 + total_prompt_tokens: int = 0 # all LLM calls of the session (every phase) total_completion_tokens: int = 0 + token_ledger: TokenLedger = field(default_factory=TokenLedger) goal_classification: Optional[dict] = None analysis_text: Optional[str] = None enforced_vmin: Optional[float] = None @@ -755,6 +761,21 @@ class SearchSession: class AgentLoopController: """Drives the iterative LLM-driven search.""" + @property + def _total_prompt_tokens(self) -> int: + """Prompt tokens of every LLM call so far (all phases).""" + return self._ledger.total_prompt + + @property + def _total_completion_tokens(self) -> int: + """Completion tokens of every LLM call so far (all phases).""" + return self._ledger.total_completion + + @property + def token_ledger(self) -> TokenLedger: + """Per-call token record of this session.""" + return self._ledger + def __init__( self, config: AppConfig, @@ -798,8 +819,8 @@ def __init__( self._latest_results_text: Optional[str] = None self._error_feedback: Optional[str] = None self._consecutive_parse_failures = 0 - self._total_prompt_tokens = 0 - self._total_completion_tokens = 0 + self._ledger = TokenLedger() # every LLM call of the session, by phase + self._ledger_iteration: Optional[int] = None self._opflow_results_cache: dict[int, OPFLOWResult] = {} self._scopflow_num_contingencies: int = 0 self._tcopflow_num_steps: int = 0 @@ -990,6 +1011,7 @@ def _normalize_sopflow_wind_base(self) -> None: # Main entry point # ------------------------------------------------------------------ + @log_scope(logger, "ExaGO session") def run(self, base_case: Path, goal: str) -> SearchSession: """Execute the full search loop.""" session_start = time.monotonic() @@ -1000,22 +1022,24 @@ def run(self, base_case: Path, goal: str) -> SearchSession: config=self._config, journal=self._journal, start_time=datetime.now().isoformat(), + token_ledger=self._ledger, ) # 1. Parse base case logger.info("Parsing base case: %s", base_case) - self._base_network = parse_matpower(base_case) - self._current_network = self._base_network - # REVIEW: model base-case wind as curtailable for SOPFLOW - # so scenarios with sub-nameplate wind are feasible (see the method + the - # search.sopflow_curtailable_wind_base flag). No-op for non-SOPFLOW / disabled. - self._normalize_sopflow_wind_base() - net_summary_text = network_summary( - self._base_network, - max_generators=self._config.report.network_summary_max_generators, - ) - net_metadata_text = network_metadata(self._base_network) - self._network_metadata_text = net_metadata_text + with log_scope(logger, "ExaGO case preparation", case_file=base_case): + self._base_network = parse_matpower(base_case) + self._current_network = self._base_network + # REVIEW: model base-case wind as curtailable for SOPFLOW + # so scenarios with sub-nameplate wind are feasible (see the method + the + # search.sopflow_curtailable_wind_base flag). No-op for non-SOPFLOW / disabled. + self._normalize_sopflow_wind_base() + net_summary_text = network_summary( + self._base_network, + max_generators=self._config.report.network_summary_max_generators, + ) + net_metadata_text = network_metadata(self._base_network) + self._network_metadata_text = net_metadata_text # 1b. Run OPFLOW benchmark at session start (if enabled) benchmark_text: Optional[str] = None @@ -1223,17 +1247,68 @@ def run(self, base_case: Path, goal: str) -> SearchSession: # 3. Agent loop max_iter = self._config.search.max_iterations + # Optional stall breaker: stop after N consecutive no-progress iterations + # (iterations that record no new journal entry — e.g. the model repeating + # an invalid action). 0/unset disables it (the default), so batch runs always + # reach max_iter and stay comparable; the UI can set it. + try: + stall_limit = int(os.environ.get("AGENTIGRID_STALL_LIMIT", "0") or "0") + except ValueError: + stall_limit = 0 + stall = 0 for iteration in range(1, max_iter + 1): if self._stop_requested: session.termination_reason = "user_stopped" self._print("\nSearch stopped by user.") break - action_type, should_continue = self._iteration(iteration, goal) - # Notify callback after each iteration + self._iteration_api_error = False + self._iteration_response_text = None + self._iteration_stop_reason = None + self._iteration_content_types = None + try: + with log_scope(logger, "ExaGO iteration", iteration=iteration): + action_type, should_continue = self._iteration(iteration, goal) + except Exception as e: + import traceback + self._print( + f"\n[Iter {iteration}] Action failed with an internal error " + f"({type(e).__name__}: {e}); discarding and continuing." + ) + traceback.print_exc() + self._error_feedback = ( + f"The previous action raised an internal error and was discarded: " + f"{type(e).__name__}: {e}. Respond with a well-formed action that " + f"strictly follows the schema." + ) + action_type, should_continue = "error", True + internal_error = True + else: + internal_error = False + # Did THIS iteration record a journal entry? A discarded/invalid action + # (common with weak local models) appends nothing. + latest_entry = self._journal.latest + recorded = latest_entry is not None and latest_entry.iteration == iteration + if action_type == "error" and not recorded: + self._record_discarded(iteration, internal_error) + # Notify the UI only on a real new entry; otherwise show a transient + # discarded card (avoids re-emitting the previous entry as a phantom). if self._on_iteration: - latest_entry = self._journal.latest - if latest_entry: + if recorded: self._on_iteration(iteration, latest_entry, action_type, self._latest_opflow) + elif action_type == "error": + self._emit_discarded(iteration) + # Stall breaker: count consecutive no-progress iterations. + if recorded: + stall = 0 + elif stall_limit: + stall += 1 + if stall >= stall_limit: + session.termination_reason = "stalled_no_progress" + self._print( + f"\nStopping early: {stall} consecutive no-progress " + f"iterations (stall limit {stall_limit})." + ) + break if not should_continue: if not session.termination_reason: session.termination_reason = "completed" @@ -1245,16 +1320,12 @@ def run(self, base_case: Path, goal: str) -> SearchSession: if not session.termination_reason: session.termination_reason = "completed" - # Final notification with termination reason - if self._on_iteration: - latest_entry = self._journal.latest - if latest_entry: - self._on_iteration( - latest_entry.iteration, - latest_entry, - session.termination_reason, - self._latest_opflow, - ) + # No final on_iteration re-emit: the last real entry was already sent by + # the per-iteration callback above. Re-emitting it duplicated the last + # timeline card, and when the final iterations were discarded it made a + # stale, lower-numbered entry appear AFTER the final iteration (e.g. a + # phantom "Iteration 6" after 20). The UI reads session.termination_reason + # when the search finishes. session.end_time = datetime.now().isoformat() session.total_prompt_tokens = self._total_prompt_tokens @@ -1272,6 +1343,88 @@ def run(self, base_case: Path, goal: str) -> SearchSession: self._finalize(session, elapsed) return session + # ------------------------------------------------------------------ + # Discarded-iteration diagnostic + # ------------------------------------------------------------------ + + def _record_discarded(self, iteration: int, internal: bool = False) -> None: + """Journal telemetry for an iteration whose action produced no entry. + + Persisted: iteration, kind (rejected / internal / api_error / + truncated), stop_reason and content block types when known, the + truncated feedback sent back to the model and, when the model answered, + its raw response (the rejected proposal itself, capped at + DISCARDED_RESPONSE_MAX characters). The entry list and all summary + statistics are unaffected.""" + stop = getattr(self, "_iteration_stop_reason", None) + if internal: + kind = "internal" + elif getattr(self, "_iteration_api_error", False): + kind = "api_error" + elif stop == "max_tokens": + kind = "truncated" # output cap hit: a harness limit, not a rejected proposal + else: + kind = "rejected" + record = { + "iteration": iteration, + "kind": kind, + "feedback": (self._error_feedback or "")[:300], + } + if stop is not None: + record["stop_reason"] = stop + ctypes = getattr(self, "_iteration_content_types", None) + if ctypes is not None: + record["content_types"] = list(ctypes) + text = getattr(self, "_iteration_response_text", None) + if text is not None and kind != "api_error": # an API error has no model response + record["response"] = text[:DISCARDED_RESPONSE_MAX] + record["response_truncated"] = len(text) > DISCARDED_RESPONSE_MAX + self._journal.discarded_actions.append(record) + + def _emit_discarded(self, iteration: int) -> None: + """Emit a UI-only 'discarded' timeline card for an iteration whose action + was rejected (invalid/malformed proposal, unknown action, parse error) + and therefore recorded no journal entry. + + The synthetic entry is deliberately NOT added to ``self._journal``: it + never affects summary stats, best-iteration selection, or the persisted + record. ``mode="discarded"`` keeps ``is_solve_iteration`` + False, and ``feasible=False`` / ``objective_value=None`` keep it out of + the session_manager's best-cost tracking. It exists purely to tell the + user *why* the iteration produced nothing, instead of showing a phantom + duplicate or a blank card. + """ + if not self._on_iteration: + return + reason = (self._error_feedback or "The model returned an invalid proposal.").strip() + # The UI label must attribute the failure to the MODEL, not the code: + # the raw feedback ("Unknown action X. Valid actions: ...") is written for + # the LLM and reads like a code assertion in the timeline. Keep only the + # first clause for the headline and prefix it clearly; the full corrective + # text still shows in the card's expander (llm_reasoning). + headline = reason.split(". ")[0].splitlines()[0].strip().rstrip(".") + if not headline: + headline = "invalid proposal" + headline = headline if len(headline) <= 120 else headline[:117] + "..." + synthetic = JournalEntry( + iteration=iteration, + description=f"LLM output rejected — {headline} (no change applied)", + commands=[], + objective_value=None, + feasible=False, + convergence_status="FAILED", + violations_count=0, + voltage_min=0.0, + voltage_max=0.0, + max_line_loading_pct=0.0, + total_gen_mw=0.0, + total_load_mw=0.0, + llm_reasoning=reason, # full detail shown in the card's expander + mode="discarded", + elapsed_seconds=0.0, + ) + self._on_iteration(iteration, synthetic, "discarded", None) + # ------------------------------------------------------------------ # Single iteration # ------------------------------------------------------------------ @@ -1283,6 +1436,7 @@ def _iteration( ) -> tuple[str, bool]: """Execute one iteration. Returns (action_type, should_continue).""" self._current_goal = goal + self._ledger_iteration = iteration # for token records of calls made below # Drain the steering queue at the iteration boundary new_directives: list[dict] = [] while True: @@ -1328,19 +1482,29 @@ def _iteration( self._on_phase(iteration, "llm_request") response = self._backend.complete(system_prompt, user_prompt) + self._ledger.record("search", response, iteration, + prompt_chars=len(system_prompt) + len(user_prompt)) logger.debug("LLM raw response: %s", response.raw_text[:500]) + # Kept for journal telemetry if this iteration's action is discarded: + # the exact proposal the parser / validator rejected. + self._iteration_response_text = response.raw_text or "" + self._iteration_stop_reason = getattr(response, "stop_reason", None) + self._iteration_content_types = getattr(response, "content_types", None) # Track tokens pt = response.prompt_tokens or 0 ct = response.completion_tokens or 0 - self._total_prompt_tokens += pt - self._total_completion_tokens += ct if pt or ct: self._print( f"[Iter {iteration}] Tokens: {pt} prompt + {ct} completion " f"(cumulative: ~{self._total_prompt_tokens + self._total_completion_tokens:,})" ) + # A failed request (auth, invalid request, network) is not a model + # answer: journal it as an API error, not as a rejected action. + if getattr(response, "api_error", False): + self._iteration_api_error = True + # Parse JSON from response if response.json_data is None: self._consecutive_parse_failures += 1 @@ -1350,14 +1514,23 @@ def _iteration( _MAX_CONSECUTIVE_PARSE_FAILURES, response.json_error, ) - self._print(f"[Iter {iteration}] Failed to parse LLM response as JSON") + truncated = getattr(response, "stop_reason", None) == "max_tokens" + self._print(f"[Iter {iteration}] Failed to parse LLM response as JSON" + + (f" (output hit max_tokens={self._config.llm.max_tokens})" if truncated else "")) if self._consecutive_parse_failures >= _MAX_CONSECUTIVE_PARSE_FAILURES: self._print(f"[Iter {iteration}] Too many consecutive parse failures — aborting") return "error", False - self._error_feedback = ( - "Failed to parse JSON from your response. " - "Please respond with a valid JSON object." - ) + if truncated: + self._error_feedback = ( + "Your previous response was cut off at the output-token limit before " + "a complete JSON object was produced. Respond with the JSON action " + "only, keeping the reasoning brief." + ) + else: + self._error_feedback = ( + "Failed to parse JSON from your response. " + "Please respond with a valid JSON object." + ) return "error", True self._consecutive_parse_failures = 0 @@ -1369,7 +1542,8 @@ def _iteration( self._batch_best = None self._executor.begin_batch() try: - return self._dispatch_action(action, iteration, data) + with log_scope(logger, "ExaGO action", iteration=iteration, action=action): + return self._dispatch_action(action, iteration, data) finally: self._prune_saved_solutions(action, self._executor.end_batch()) @@ -1511,6 +1685,14 @@ def _handle_modify( f"{skipped_count} skipped" ) + # A modify that applied zero commands changes nothing — make the timeline + # say so instead of showing a blank "No description" card. + if applied_count == 0: + if not description or description == "No description": + description = "No commands applied (no change)" + else: + description = f"{description} — no commands applied" + if all_errors: self._error_feedback = "Command errors:\n" + "\n".join(all_errors) @@ -1626,7 +1808,7 @@ def _handle_modify( self._active_steering_directives[-1]["directive"] if self._active_steering_directives else None ) - self._journal.add_from_results( + _modify_entry = self._journal.add_from_results( iteration=iteration, description=description, commands=raw_commands, @@ -1640,6 +1822,7 @@ def _handle_modify( gencost=modified_net.gencost if self._config.search.application == "pflow" else None, exago_command=_single_call_record(sim_result), ) + _modify_entry.skipped_commands = list(all_errors) # Extract tracked metrics for multi-objective tracking if opflow is not None: @@ -1695,6 +1878,7 @@ def _variant_has_vlimits(raw_cmds: list[dict]) -> bool: return True return False + @log_scope(logger, "ExaGO variant exploration") def _handle_explore( self, iteration: int, data: dict ) -> tuple[str, bool]: @@ -2079,6 +2263,11 @@ def _resolve_candidate_buses(self, spec: dict) -> tuple[Optional[list[int]], str Returns (candidates, "") on success or (None, error_message) on failure. Shared by the feasibility sweep and the boundary sweep. """ + if not isinstance(spec, dict): + return None, ( + f"candidate_set must be an object with a 'type' field, " + f"got {type(spec).__name__}: {spec!r}" + ) ctype = spec.get("type", "") if ctype == "all_buses": candidates = [b.bus_i for b in self._base_network.buses] @@ -2275,6 +2464,8 @@ def _handle_sweep( candidate_set_spec = data.get("candidate_set", {}) mutation_template = data.get("mutation", {}) feasibility_spec = data.get("feasibility", {}) + if not isinstance(feasibility_spec, dict): + feasibility_spec = {} if not mutation_template or "action" not in mutation_template: self._error_feedback = "sweep requires a 'mutation' dict with an 'action' key." @@ -2693,6 +2884,7 @@ def solve_probe(delta_mw: float) -> _ProbeOutcome: result["pf_tan_phi"] = tan_phi return result + @log_scope(logger, "ExaGO boundary sweep") def _handle_boundary_sweep( self, iteration: int, data: dict ) -> tuple[str, bool]: @@ -2712,6 +2904,8 @@ def _handle_boundary_sweep( pf_spec = data.get("power_factor", s.boundary_power_factor_default) candidate_set_spec = data.get("candidate_set", {}) feasibility_spec = data.get("feasibility", {}) + if not isinstance(feasibility_spec, dict): + feasibility_spec = {} vmin = feasibility_spec.get("Vmin", 0.9) vmax = feasibility_spec.get("Vmax", 1.1) q_frac = s.boundary_gen_q_frac @@ -2901,6 +3095,7 @@ def _boundary_progress(done: int, total: int) -> None: return "sweep", True + @log_scope(logger, "ExaGO contingency screen") def _run_contingency_screen( self, net: MATNetwork, @@ -3346,6 +3541,7 @@ def _handle_contingency_sweep( self._print(f"[Iter {iteration}] {verdict}") return "sweep", True + @log_scope(logger, "ExaGO reserve screen") def _handle_reserve_screen( self, iteration: int, data: dict ) -> tuple[str, bool]: @@ -3370,6 +3566,8 @@ def _handle_reserve_screen( description = data.get("description", "Hot reserve / N-1 generator security") reasoning = data.get("reasoning", "") feasibility_spec = data.get("feasibility", {}) + if not isinstance(feasibility_spec, dict): + feasibility_spec = {} vmin = feasibility_spec.get("Vmin", 0.9) vmax = feasibility_spec.get("Vmax", 1.1) minimize = bool(data.get("minimize", False)) @@ -3544,6 +3742,7 @@ def _handle_reserve_screen( ) return "sweep", True + @log_scope(logger, "ExaGO reserve minimization") def _run_reserve_minimization( self, iteration: int, net: MATNetwork, vmin: float, vmax: float, ) -> "reserve_search.ReserveMinResult": @@ -3722,6 +3921,7 @@ def _build_reserve_llm_view( ) return "\n".join(lines) + @log_scope(logger, "ExaGO relief phase") def _run_relief_phase( self, iteration: int, @@ -4678,6 +4878,8 @@ def _run_analysis_query(self, query: str) -> str: context = self._latest_results_text or "No results available." user = f"Question: {query}\n\nCurrent results:\n{context}" response = self._backend.complete(system, user) + self._ledger.record("search", response, self._ledger_iteration, + prompt_chars=len(system) + len(user)) return response.raw_text except Exception as exc: logger.warning("LLM fallback for analysis failed: %s", exc) @@ -4698,6 +4900,7 @@ def _run_analysis_query(self, query: str) -> str: # Multi-objective helpers # ------------------------------------------------------------------ + @log_scope(logger, "Initial objective extraction") def _extract_initial_objectives(self, goal: str) -> None: """Use the LLM to extract tracked objectives from the initial goal.""" try: @@ -4707,6 +4910,8 @@ def _extract_initial_objectives(self, goal: str) -> None: context="initial_goal", ) response = self._backend.complete(sys_prompt, user_prompt) + self._ledger.record("objective_extraction", response, 0, + prompt_chars=len(sys_prompt) + len(user_prompt)) parsed = parse_objective_extraction(response.raw_text) if parsed: for obj_data in parsed: @@ -4730,6 +4935,7 @@ def _extract_initial_objectives(self, goal: str) -> None: except Exception as exc: logger.warning("Failed to extract initial objectives: %s", exc) + @log_scope(logger, "Steering objective extraction") def _extract_objectives_from_steering(self, directive: str, iteration: int) -> None: """Extract any new objectives from a steering directive.""" try: @@ -4739,6 +4945,8 @@ def _extract_objectives_from_steering(self, directive: str, iteration: int) -> N context="steering_directive", ) response = self._backend.complete(sys_prompt, user_prompt) + self._ledger.record("steering", response, iteration, + prompt_chars=len(sys_prompt) + len(user_prompt)) parsed = parse_objective_extraction(response.raw_text) if parsed: for obj_data in parsed: @@ -4777,6 +4985,7 @@ def _backfill_metrics(self) -> None: # Prompt assembly # ------------------------------------------------------------------ + @log_scope(logger, "ExaGO prompt assembly") def _assemble_prompt( self, goal: str, @@ -4825,9 +5034,10 @@ def _assemble_prompt( # Finalization # ------------------------------------------------------------------ + @log_scope(logger, "ExaGO finalization") def _finalize(self, session: SearchSession, elapsed_seconds: float) -> None: """Print summary and save journal.""" - total_tokens = self._total_prompt_tokens + self._total_completion_tokens + tokens_before_analysis = self._ledger.total if self._pending_explore is not None: # session ended right after an explore dirs, best = self._pending_explore self._pending_explore = None @@ -4844,13 +5054,15 @@ def _finalize(self, session: SearchSession, elapsed_seconds: float) -> None: stats=raw_stats, # Bounded digest avoids context overflow on large sweeps. journal_formatted=self._journal.format_for_classification(), - total_tokens=total_tokens, + total_tokens=tokens_before_analysis, objective_registry=self._journal.objective_registry.to_dict_list(), preference_history=self._journal.objective_registry.history, application=self._config.search.application, near_optimal_abs_tol=self._config.report.near_optimal_abs_tol, ) response = self._backend.complete(sys_prompt, user_prompt) + self._ledger.record("analysis", response, + prompt_chars=len(sys_prompt) + len(user_prompt)) analysis_text = response.raw_text valid_iters = {e.iteration for e in self._journal.entries} goal_classification = parse_goal_classification(analysis_text, valid_iters) @@ -4870,6 +5082,8 @@ def _finalize(self, session: SearchSession, elapsed_seconds: float) -> None: ) # Store on session for downstream consumers (GUI, tests) + session.total_prompt_tokens = self._total_prompt_tokens + session.total_completion_tokens = self._total_completion_tokens session.goal_classification = goal_classification session.analysis_text = analysis_text session.objective_registry_data = self._journal.objective_registry.to_dict_list() @@ -4889,16 +5103,12 @@ def _finalize(self, session: SearchSession, elapsed_seconds: float) -> None: print(f" Application: {session.application}") print(f" Backend: {self._backend.name()} ({self._config.llm.model})") print( - f" Iterations: {stats['total_iterations']} " + f" Iterations: {format_iteration_count(stats)} " f"(of max {self._config.search.max_iterations})" ) print(f" Duration: {elapsed_seconds:.1f} seconds") - if total_tokens: - print( - f" Tokens used: ~{total_tokens:,} " - f"(prompt: {self._total_prompt_tokens:,}, " - f"completion: {self._total_completion_tokens:,})" - ) + if len(self._ledger): + print(f" {self._ledger.summary_text()}") print(f" Termination: {session.termination_reason}") # Goal-aware best-solution line @@ -5014,6 +5224,14 @@ def _finalize(self, session: SearchSession, elapsed_seconds: float) -> None: self._config.output.workdir / f"journal_{timestamp}.{fmt}" ) journal_path.parent.mkdir(parents=True, exist_ok=True) + self._journal.token_usage = self._ledger.to_dict() + self._journal.llm_settings = { + "backend": self._backend.name(), + "model": self._config.llm.model, + "max_tokens": self._config.llm.max_tokens, + "prompt_cache": bool(getattr(self._config.llm, "prompt_cache", False)), + "temperature_sent": getattr(self._backend, "temperature_sent", None), + } if fmt == "csv": self._journal.export_csv(journal_path) else: @@ -5024,6 +5242,7 @@ def _finalize(self, session: SearchSession, elapsed_seconds: float) -> None: # Session save/resume # ------------------------------------------------------------------ + @log_scope(logger, "ExaGO session resume") def resume_from(self, save_dir: Path) -> SearchSession: """Resume a search from a saved session checkpoint. @@ -5080,25 +5299,30 @@ def resume_from(self, save_dir: Path) -> SearchSession: else: self._tcopflow_profile_overrides = {} - # Restore token counts - self._total_prompt_tokens = saved["total_prompt_tokens"] - self._total_completion_tokens = saved["total_completion_tokens"] + # Restore token counts (sessions saved before per-call accounting: totals only) + if saved.get("token_usage"): + self._ledger = TokenLedger.from_dict(saved["token_usage"]) + else: + self._ledger = TokenLedger.from_legacy_totals( + saved["total_prompt_tokens"], saved["total_completion_tokens"], + ) # Parse base case and restore networks base_case = saved["base_case_path"] goal = saved["goal"] self._current_goal = goal - self._base_network = parse_matpower(base_case) - self._current_network = saved["current_network"] or self._base_network + with log_scope(logger, "ExaGO case preparation", case_file=base_case): + self._base_network = parse_matpower(base_case) + self._current_network = saved["current_network"] or self._base_network - # Rebuild system prompt - net_summary_text = network_summary( - self._base_network, - max_generators=self._config.report.network_summary_max_generators, - ) - net_metadata_text = network_metadata(self._base_network) - self._network_metadata_text = net_metadata_text + # Rebuild system prompt + net_summary_text = network_summary( + self._base_network, + max_generators=self._config.report.network_summary_max_generators, + ) + net_metadata_text = network_metadata(self._base_network) + self._network_metadata_text = net_metadata_text # Restore load_factor from journal if saved if self._journal.load_factor is not None: self._session_load_factor = self._journal.load_factor @@ -5154,20 +5378,67 @@ def resume_from(self, save_dir: Path) -> SearchSession: config=self._config, journal=self._journal, start_time=datetime.now().isoformat(), + token_ledger=self._ledger, ) # Continue the agent loop from last_iteration + 1 max_iter = self._config.search.max_iterations + try: + stall_limit = int(os.environ.get("AGENTIGRID_STALL_LIMIT", "0") or "0") + except ValueError: + stall_limit = 0 + stall = 0 for iteration in range(last_iteration + 1, max_iter + 1): if self._stop_requested: session.termination_reason = "user_stopped" self._print("\nSearch stopped by user.") break - action_type, should_continue = self._iteration(iteration, goal) + self._iteration_api_error = False + self._iteration_response_text = None + self._iteration_stop_reason = None + self._iteration_content_types = None + try: + with log_scope(logger, "ExaGO iteration", iteration=iteration): + action_type, should_continue = self._iteration(iteration, goal) + except Exception as e: + import traceback + self._print( + f"\n[Iter {iteration}] Action failed with an internal error " + f"({type(e).__name__}: {e}); discarding and continuing." + ) + traceback.print_exc() + self._error_feedback = ( + f"The previous action raised an internal error and was discarded: " + f"{type(e).__name__}: {e}. Respond with a well-formed action that " + f"strictly follows the schema." + ) + action_type, should_continue = "error", True + internal_error = True + else: + internal_error = False + + # Only emit when this iteration recorded a new entry (see run() above). + latest_entry = self._journal.latest + recorded = latest_entry is not None and latest_entry.iteration == iteration + if action_type == "error" and not recorded: + self._record_discarded(iteration, internal_error) if self._on_iteration: - latest_entry = self._journal.latest - if latest_entry: + if recorded: self._on_iteration(iteration, latest_entry, action_type, self._latest_opflow) + elif action_type == "error": + self._emit_discarded(iteration) + # Stall breaker: count consecutive no-progress iterations. + if recorded: + stall = 0 + elif stall_limit: + stall += 1 + if stall >= stall_limit: + session.termination_reason = "stalled_no_progress" + self._print( + f"\nStopping early: {stall} consecutive no-progress " + f"iterations (stall limit {stall_limit})." + ) + break if not should_continue: if not session.termination_reason: session.termination_reason = "completed" @@ -5179,14 +5450,8 @@ def resume_from(self, save_dir: Path) -> SearchSession: if not session.termination_reason: session.termination_reason = "completed" - # Final notification - if self._on_iteration: - latest_entry = self._journal.latest - if latest_entry: - self._on_iteration( - latest_entry.iteration, latest_entry, - session.termination_reason, self._latest_opflow, - ) + # No final on_iteration re-emit (see run() above) — it only duplicated + # the last timeline card. session.end_time = datetime.now().isoformat() session.total_prompt_tokens = self._total_prompt_tokens @@ -5258,6 +5523,7 @@ def save_session(self, save_dir: Path, config_path: Path | str | None = None) -> current_network=self._current_network, total_prompt_tokens=self._total_prompt_tokens, total_completion_tokens=self._total_completion_tokens, + token_usage=self._ledger.to_dict(), last_iteration=last_iteration, enforced_vmin=enforced_vmin, enforced_vmax=enforced_vmax, diff --git a/agentigrid/agentigrid/exago_engine/benchmark.py b/agentigrid/agentigrid/exago_engine/benchmark.py index 50edaf49..aa35697d 100644 --- a/agentigrid/agentigrid/exago_engine/benchmark.py +++ b/agentigrid/agentigrid/exago_engine/benchmark.py @@ -7,6 +7,7 @@ from pathlib import Path from typing import Optional +from agentigrid.logging_setup import log_scope from agentigrid.exago_engine.executor import SimulationExecutor, SimulationResult from agentigrid.exago_engine.journal import SearchJournal from agentigrid.exago_parsers import parse_simulation_result_for_app @@ -50,6 +51,7 @@ class BenchmarkResult: error: str | None = None +@log_scope(logger, "OPFLOW baseline benchmark") def _run_opflow_on_base_case( base_case_path: Path, config: AppConfig, @@ -84,6 +86,7 @@ def _run_opflow_on_base_case( return parsed +@log_scope(logger, "PFLOW versus OPFLOW benchmark") def run_pflow_vs_opflow_benchmark( base_case_path: Path, pflow_journal: SearchJournal, diff --git a/agentigrid/agentigrid/exago_engine/executor.py b/agentigrid/agentigrid/exago_engine/executor.py index 70e86cfa..51e4461e 100644 --- a/agentigrid/agentigrid/exago_engine/executor.py +++ b/agentigrid/agentigrid/exago_engine/executor.py @@ -14,6 +14,7 @@ from pathlib import Path from typing import Callable, Optional +from agentigrid.logging_setup import log_scope from agentigrid.config import ExagoConfig, OutputConfig from agentigrid.exago_parsers.matpower_model import MATNetwork from agentigrid.exago_parsers.matpower_writer import write_matpower @@ -288,29 +289,35 @@ def run( t0 = time.monotonic() try: - if self._env_script: - shell_cmd = _shell_command - proc = subprocess.run( - shell_cmd, - shell=True, - executable="/bin/bash", - capture_output=True, - text=True, - timeout=self._exago.timeout, - cwd=run_dir, - ) - else: - run_env = None - if thread_limit and thread_limit > 0: - run_env = {**os.environ, **{v: str(thread_limit) for v in _THREAD_ENV_VARS}} - proc = subprocess.run( - cmd, - capture_output=True, - text=True, - timeout=self._exago.timeout, - cwd=run_dir, - env=run_env, - ) + with log_scope( + logger, "ExaGO subprocess", + application=application, + iteration=iteration, + run_dir=run_dir, + ): + if self._env_script: + shell_cmd = _shell_command + proc = subprocess.run( + shell_cmd, + shell=True, + executable="/bin/bash", + capture_output=True, + text=True, + timeout=self._exago.timeout, + cwd=run_dir, + ) + else: + run_env = None + if thread_limit and thread_limit > 0: + run_env = {**os.environ, **{v: str(thread_limit) for v in _THREAD_ENV_VARS}} + proc = subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=self._exago.timeout, + cwd=run_dir, + env=run_env, + ) except subprocess.TimeoutExpired: elapsed = time.monotonic() - t0 error_msg = f"Simulation timed out after {self._exago.timeout} seconds" @@ -415,7 +422,10 @@ def run_parallel( results: dict[int, SimulationResult] = {} done = 0 - with ThreadPoolExecutor(max_workers=max_workers) as pool: + with ( + log_scope(logger, "ExaGO simulation batch", tasks=len(tasks), workers=max_workers), + ThreadPoolExecutor(max_workers=max_workers) as pool, + ): future_to_idx: dict = {} for idx, (network, application, iteration, extra_args) in enumerate(tasks): future = pool.submit( @@ -490,7 +500,10 @@ def map_callables( results: dict[int, object] = {} done = 0 - with ThreadPoolExecutor(max_workers=max_workers) as pool: + with ( + log_scope(logger, "ExaGO candidate batch", tasks=len(fns), workers=max_workers), + ThreadPoolExecutor(max_workers=max_workers) as pool, + ): future_to_idx = {pool.submit(fn): idx for idx, fn in enumerate(fns)} for future in as_completed(future_to_idx): idx = future_to_idx[future] diff --git a/agentigrid/agentigrid/exago_engine/goal_classifier.py b/agentigrid/agentigrid/exago_engine/goal_classifier.py index d9edf13a..9ce3a01f 100644 --- a/agentigrid/agentigrid/exago_engine/goal_classifier.py +++ b/agentigrid/agentigrid/exago_engine/goal_classifier.py @@ -58,7 +58,8 @@ def build_classification_prompts( termination_reason: How the search ended (completed, max_iterations, etc.). stats: dict from ``SearchJournal.summary_stats()`` (before any override). journal_formatted: Output of ``SearchJournal.format_detailed()``. - total_tokens: Total prompt + completion tokens used during the search. + total_tokens: Prompt + completion tokens of every LLM call made before this + one (all phases; the classification call cannot count itself). objective_registry: Optional list of tracked objective dicts. preference_history: Optional list of preference change dicts. application: ExaGO application name (e.g., "opflow", "tcopflow", "sopflow"). @@ -138,7 +139,7 @@ def build_classification_prompts( f"\n" f"Lowest-cost feasible: {stats['best_objective']} " f"({_best_loc})\n" - f"Tokens used: ~{total_tokens:,}\n" + f"Tokens used before this analysis: ~{total_tokens:,}\n" f"\n" f"=== Detailed Journal ===\n" f"{journal_formatted}\n" diff --git a/agentigrid/agentigrid/exago_engine/journal.py b/agentigrid/agentigrid/exago_engine/journal.py index f86d4bc5..1d07bb5c 100644 --- a/agentigrid/agentigrid/exago_engine/journal.py +++ b/agentigrid/agentigrid/exago_engine/journal.py @@ -55,6 +55,10 @@ class JournalEntry: contingency_meta: Optional[dict] = None # C.5: mutation, changed element, affected elements, order, pass/fail counts reserve_meta: Optional[dict] = None # C.8: hot-reserve / N-1 generator security accounting findings: Optional[dict] = None # complete: the LLM's full findings (summary, details, ...) + # Modify iterations: commands the LLM proposed that were NOT applied (parse + # errors or validator skips), as messages. [] = every proposed command was + # applied; None = not recorded (non-modify entries, older journals). + skipped_commands: Optional[list[str]] = None def describe_contingency_scope(meta: dict) -> str: @@ -217,6 +221,16 @@ def __init__(self) -> None: self.session_best: Optional[dict] = None # load_factor: session-level load scaling factor (updated by set_load_factor action). self.load_factor: Optional[float] = None + # token_usage: TokenLedger.to_dict() of the session, set before export. + self.token_usage: Optional[dict] = None + # llm_settings: the LLM settings this run actually used (backend, model, + # max_tokens, prompt_cache, temperature_sent), set before export. + self.llm_settings: Optional[dict] = None + # discarded_actions: iterations whose LLM action recorded no entry -- + # rejected by the deterministic validator / parser ("rejected"), a failed + # API request ("api_error"), output cut at max_tokens ("truncated") or + # an internal error ("internal"). Exported with the journal. + self.discarded_actions: list[dict] = [] def update_session_best( self, @@ -1122,6 +1136,12 @@ def export_json(self, path: Path) -> None: data["session_best"] = self.session_best if self.load_factor is not None: data["load_factor"] = self.load_factor + if self.token_usage is not None: + data["token_usage"] = self.token_usage + if self.llm_settings is not None: + data["llm_settings"] = self.llm_settings + if getattr(self, "discarded_actions", None): + data["discarded_actions"] = self.discarded_actions path.write_text(json.dumps(data, indent=2), encoding="utf-8") logger.info("Journal exported to %s (%d entries)", path, len(self._entries)) @@ -1139,7 +1159,7 @@ def export_csv(self, path: Path) -> None: "mode", "elapsed_seconds", "timestamp", "steering_directive", "tracked_metrics", "feasibility_detail", "solver", "num_steps", "num_scenarios", "explored_variants", "candidate_count", "feasible_buses", "exago_command", - "contingency_meta", "reserve_meta", "findings", + "contingency_meta", "reserve_meta", "findings", "skipped_commands", ] with open(path, "w", newline="", encoding="utf-8") as f: @@ -1155,6 +1175,7 @@ def export_csv(self, path: Path) -> None: row["contingency_meta"] = json.dumps(row.get("contingency_meta") or None) row["reserve_meta"] = json.dumps(row.get("reserve_meta") or None) row["findings"] = json.dumps(row.get("findings") or None) + row["skipped_commands"] = json.dumps(row.get("skipped_commands")) writer.writerow(row) logger.info("Journal CSV exported to %s (%d entries)", path, len(self._entries)) @@ -1163,6 +1184,22 @@ def export_csv(self, path: Path) -> None: # Summary statistics # ------------------------------------------------------------------ + def _iteration_counts(self) -> tuple[int, bool]: + """Agent iterations actually run, and whether a base case was solved. + + ``total_iterations`` counts journal entries: it includes the base-case + solve (iteration 0) and leaves out iterations whose LLM output was + rejected (they are kept in ``discarded_actions``, not as entries). + This counts distinct iteration numbers >= 1 across both, which is the + number of agent iterations the user asked for and the run performed. + """ + numbers = {e.iteration for e in self._entries} + for d in getattr(self, "discarded_actions", None) or []: + it = d.get("iteration") if isinstance(d, dict) else None + if isinstance(it, int): + numbers.add(it) + return sum(1 for n in numbers if n >= 1), 0 in numbers + def summary_stats( self, best_iteration_override: int | None = None, @@ -1182,9 +1219,12 @@ def summary_stats( feasible_count, infeasible_count, objective_trend, voltage_range_trend, goal_type. """ + llm_iterations, has_base_case = self._iteration_counts() if not self._entries: return { "total_iterations": 0, + "llm_iterations": llm_iterations, + "has_base_case": has_base_case, "best_objective": None, "best_iteration": None, "best_bus": None, @@ -1242,6 +1282,8 @@ def summary_stats( return { "total_iterations": len(self._entries), + "llm_iterations": llm_iterations, + "has_base_case": has_base_case, "best_objective": best_objective, "best_iteration": best_iteration, "best_bus": best_bus, @@ -1296,3 +1338,15 @@ def _best_cost_in_sweeps(self) -> tuple[float, int, int | None] | None: if best is None or c < best[0]: best = (c, e.iteration, v.get("bus")) return best + + +def format_iteration_count(stats: dict) -> str: + """Display form of the iteration count, e.g. ``"2 + base case"``. + + Falls back to ``total_iterations`` for stats dicts written before + ``llm_iterations`` existed. + """ + if "llm_iterations" not in stats: + return str(stats.get("total_iterations", 0)) + n = stats["llm_iterations"] + return f"{n} + base case" if stats.get("has_base_case") else str(n) diff --git a/agentigrid/agentigrid/exago_engine/regenerate_report.py b/agentigrid/agentigrid/exago_engine/regenerate_report.py index a18c66fa..0c97aa4f 100644 --- a/agentigrid/agentigrid/exago_engine/regenerate_report.py +++ b/agentigrid/agentigrid/exago_engine/regenerate_report.py @@ -18,6 +18,7 @@ from typing import Any, Optional from agentigrid.config import AppConfig +from agentigrid.token_usage import TokenLedger from agentigrid.exago_engine.goal_classifier import ( build_classification_prompts, parse_goal_classification, @@ -56,7 +57,7 @@ def load_journal_export(path: Path) -> SearchJournal: Rebuilds every current :class:`JournalEntry` field (dict.get with dataclass defaults, so older exports still load), the :class:`ObjectiveRegistry` (with its saved preference history), and the optional - ``benchmark_result`` / ``session_best`` / ``load_factor`` fields. + ``benchmark_result`` / ``session_best`` / ``load_factor`` / ``token_usage`` fields. Raises: ValueError: If *path* looks like a ``session.json`` (has ``format_version`` @@ -110,6 +111,7 @@ def load_journal_export(path: Path) -> SearchJournal: exago_command=ed.get("exago_command"), contingency_meta=ed.get("contingency_meta"), reserve_meta=ed.get("reserve_meta"), + skipped_commands=ed.get("skipped_commands"), )) # Rebuild the objective registry, then overwrite its history with the saved @@ -133,6 +135,8 @@ def load_journal_export(path: Path) -> SearchJournal: journal.session_best = raw["session_best"] if raw.get("load_factor") is not None: journal.load_factor = raw["load_factor"] + if raw.get("token_usage") is not None: # absent in journals before token accounting + journal.token_usage = raw["token_usage"] logger.info("Loaded journal export %s (%d entries)", path, len(journal.entries)) return journal @@ -145,17 +149,21 @@ def build_session_stub( termination_reason: str = "regenerated", start_time: Optional[str] = None, end_time: Optional[str] = None, + token_ledger: Optional[TokenLedger] = None, ) -> Any: """Build a :class:`SearchSession` carrying exactly what ``generate()`` reads. ``SearchSession`` is a plain dataclass constructible offline (no controller, - no ExaGO), so the real object is used. Token counters are zeroed; timestamps - default to now(); ``benchmark_result`` is carried over from the journal. + no ExaGO), so the real object is used. Token counters come from + *token_ledger*, else from the journal's ``token_usage`` (zero when the journal + predates token accounting); timestamps default to now(); ``benchmark_result`` + is carried over from the journal. """ # Imported lazily so this module has no import-time coupling to the loop. from agentigrid.exago_engine.agent_loop import SearchSession now = datetime.now().isoformat() + ledger = token_ledger if token_ledger is not None else TokenLedger.from_dict(journal.token_usage) return SearchSession( goal=goal, application=config.search.application, @@ -165,8 +173,9 @@ def build_session_stub( start_time=start_time or now, end_time=end_time or now, termination_reason=termination_reason, - total_prompt_tokens=0, - total_completion_tokens=0, + total_prompt_tokens=ledger.total_prompt, + total_completion_tokens=ledger.total_completion, + token_ledger=ledger, objective_registry_data=journal.objective_registry.to_dict_list(), preference_history=journal.objective_registry.history, benchmark_result=journal.benchmark_result, @@ -198,6 +207,8 @@ def regenerate_from_journal( Path to the written PDF. """ journal = load_journal_export(journal_json_path) + # The session's own calls (if recorded), plus the regeneration call below. + ledger = TokenLedger.from_dict(journal.token_usage) analysis_text: Optional[str] = None gc: Optional[dict] = None @@ -217,6 +228,7 @@ def regenerate_from_journal( ) try: resp = backend.complete(sys_p, user_p) + ledger.record("report", resp, prompt_chars=len(sys_p) + len(user_p)) analysis_text = resp.raw_text gc = parse_goal_classification( analysis_text, {e.iteration for e in journal.entries} @@ -228,7 +240,7 @@ def regenerate_from_journal( analysis_text = None gc = None - session = build_session_stub(journal, config, goal) + session = build_session_stub(journal, config, goal, token_ledger=ledger) # Imported here (not at module top) so the deterministic --no-analysis path # and unit tests don't pay the ReportLab/Plotly import cost unless rendering. diff --git a/agentigrid/agentigrid/exago_engine/session_io.py b/agentigrid/agentigrid/exago_engine/session_io.py index cfd65e4b..0a5c0ec4 100644 --- a/agentigrid/agentigrid/exago_engine/session_io.py +++ b/agentigrid/agentigrid/exago_engine/session_io.py @@ -9,6 +9,7 @@ from pathlib import Path from typing import Any, Optional +from agentigrid.logging_setup import log_scope from agentigrid.exago_engine.journal import ( JournalEntry, ObjectiveEntry, ObjectiveRegistry, SearchJournal, ) @@ -18,7 +19,7 @@ logger = logging.getLogger("agentigrid.exago_engine.session_io") -SESSION_FORMAT_VERSION = "1.1" +SESSION_FORMAT_VERSION = "1.2" # 1.2: token_usage (per-call token ledger) def save_session( @@ -34,6 +35,7 @@ def save_session( total_prompt_tokens: int, total_completion_tokens: int, last_iteration: int, + token_usage: dict | None = None, termination_reason: str = "paused", enforced_vmin: float | None = None, enforced_vmax: float | None = None, @@ -59,9 +61,10 @@ def save_session( steering_history: List of steering directive dicts. active_steering_directives: Currently active steering directives. current_network: The modified network state (last iteration's network). - total_prompt_tokens: Cumulative prompt tokens used. - total_completion_tokens: Cumulative completion tokens used. + total_prompt_tokens: Cumulative prompt tokens used (all LLM calls). + total_completion_tokens: Cumulative completion tokens used (all LLM calls). last_iteration: The last completed iteration number. + token_usage: ``TokenLedger.to_dict()`` of the session (per-call breakdown). termination_reason: Why the session was saved (e.g., "paused", "user_stopped"). enforced_vmin: Minimum voltage limit enforced in current network. enforced_vmax: Maximum voltage limit enforced in current network. @@ -69,51 +72,53 @@ def save_session( Returns: Path to the saved session directory. """ - save_dir.mkdir(parents=True, exist_ok=True) - - # Save current network state - network_file = None - if current_network is not None: - network_file = "current_network.m" - write_matpower(current_network, save_dir / network_file) - - # Build session data - session_data = { - "format_version": SESSION_FORMAT_VERSION, - "saved_at": datetime.now().isoformat(), - "goal": goal, - "application": application, - "base_case_path": str(base_case_path), - "config_path": str(config_path) if config_path else None, - "last_iteration": last_iteration, - "termination_reason": termination_reason, - "total_prompt_tokens": total_prompt_tokens, - "total_completion_tokens": total_completion_tokens, - "enforced_vmin": enforced_vmin, - "enforced_vmax": enforced_vmax, - "current_network_file": network_file, - "steering_history": steering_history, - "active_steering_directives": active_steering_directives, - "tcopflow_period_data": tcopflow_period_data, - "tcopflow_dT_min": tcopflow_dT_min, - "tcopflow_duration_min": tcopflow_duration_min, - "tcopflow_is_coupling": tcopflow_is_coupling, - "sopflow_num_scenarios": sopflow_num_scenarios, - "sopflow_scenario_override": sopflow_scenario_override, - "tcopflow_profile_overrides": tcopflow_profile_overrides, - "benchmark_result": benchmark_result, - "explore_cache_info": explore_cache_info, - "journal": { - "entries": [asdict(e) for e in journal.entries], - "objective_registry": journal.objective_registry.to_dict_list(), - "preference_history": journal.objective_registry.history, - }, - } - - session_path = save_dir / "session.json" - session_path.write_text(json.dumps(session_data, indent=2), encoding="utf-8") - logger.info("Session saved to %s (%d entries)", save_dir, len(journal)) - return save_dir + with log_scope(logger, "Session save", save_dir=save_dir): + save_dir.mkdir(parents=True, exist_ok=True) + + # Save current network state + network_file = None + if current_network is not None: + network_file = "current_network.m" + write_matpower(current_network, save_dir / network_file) + + # Build session data + session_data = { + "format_version": SESSION_FORMAT_VERSION, + "saved_at": datetime.now().isoformat(), + "goal": goal, + "application": application, + "base_case_path": str(base_case_path), + "config_path": str(config_path) if config_path else None, + "last_iteration": last_iteration, + "termination_reason": termination_reason, + "total_prompt_tokens": total_prompt_tokens, + "total_completion_tokens": total_completion_tokens, + "token_usage": token_usage, + "enforced_vmin": enforced_vmin, + "enforced_vmax": enforced_vmax, + "current_network_file": network_file, + "steering_history": steering_history, + "active_steering_directives": active_steering_directives, + "tcopflow_period_data": tcopflow_period_data, + "tcopflow_dT_min": tcopflow_dT_min, + "tcopflow_duration_min": tcopflow_duration_min, + "tcopflow_is_coupling": tcopflow_is_coupling, + "sopflow_num_scenarios": sopflow_num_scenarios, + "sopflow_scenario_override": sopflow_scenario_override, + "tcopflow_profile_overrides": tcopflow_profile_overrides, + "benchmark_result": benchmark_result, + "explore_cache_info": explore_cache_info, + "journal": { + "entries": [asdict(e) for e in journal.entries], + "objective_registry": journal.objective_registry.to_dict_list(), + "preference_history": journal.objective_registry.history, + }, + } + + session_path = save_dir / "session.json" + session_path.write_text(json.dumps(session_data, indent=2), encoding="utf-8") + logger.info("Session saved to %s (%d entries)", save_dir, len(journal)) + return save_dir def load_session(save_dir: Path) -> dict[str, Any]: @@ -130,111 +135,115 @@ def load_session(save_dir: Path) -> dict[str, Any]: - "steering_history", "active_steering_directives" - "current_network" (MATNetwork or None) - "last_iteration", "total_prompt_tokens", "total_completion_tokens" + - "token_usage" (TokenLedger.to_dict(), None for sessions saved before format 1.2) - "enforced_vmin", "enforced_vmax" Raises: FileNotFoundError: If session.json doesn't exist. ValueError: If the format version is unsupported. """ - session_path = save_dir / "session.json" - if not session_path.exists(): - raise FileNotFoundError(f"Session file not found: {session_path}") - - raw = json.loads(session_path.read_text(encoding="utf-8")) - - version = raw.get("format_version", "0") - if version not in ("1.0", SESSION_FORMAT_VERSION): - raise ValueError( - f"Unsupported session format version '{version}' " - f"(expected '{SESSION_FORMAT_VERSION}')" - ) - - # Reconstruct journal entries - journal_entries = [] - for entry_data in raw.get("journal", {}).get("entries", []): - entry = JournalEntry( - iteration=entry_data["iteration"], - description=entry_data["description"], - commands=entry_data.get("commands", []), - objective_value=entry_data.get("objective_value"), - feasible=entry_data.get("feasible", False), - convergence_status=entry_data.get("convergence_status", "UNKNOWN"), - violations_count=entry_data.get("violations_count", 0), - voltage_min=entry_data.get("voltage_min", 0.0), - voltage_max=entry_data.get("voltage_max", 0.0), - max_line_loading_pct=entry_data.get("max_line_loading_pct", 0.0), - total_gen_mw=entry_data.get("total_gen_mw", 0.0), - total_load_mw=entry_data.get("total_load_mw", 0.0), - llm_reasoning=entry_data.get("llm_reasoning", ""), - mode=entry_data.get("mode", "fresh"), - elapsed_seconds=entry_data.get("elapsed_seconds", 0.0), - timestamp=entry_data.get("timestamp", ""), - steering_directive=entry_data.get("steering_directive"), - tracked_metrics=entry_data.get("tracked_metrics"), - feasibility_detail=entry_data.get("feasibility_detail", ""), - solver=entry_data.get("solver", ""), - num_steps=entry_data.get("num_steps", 0), - num_scenarios=entry_data.get("num_scenarios", 0), - explored_variants=entry_data.get("explored_variants"), - # Previously dropped on resume — restore so sweep/contingency/reserve - # metadata survives a save/load round-trip (dataclass defaults keep - # older sessions loadable). - candidate_count=entry_data.get("candidate_count", 0), - feasible_buses=entry_data.get("feasible_buses"), - exago_command=entry_data.get("exago_command"), - contingency_meta=entry_data.get("contingency_meta"), - reserve_meta=entry_data.get("reserve_meta"), - findings=entry_data.get("findings"), - ) - journal_entries.append(entry) - - # Reconstruct objective registry - registry = ObjectiveRegistry() - for obj_data in raw.get("journal", {}).get("objective_registry", []): - registry.register(ObjectiveEntry( - name=obj_data["name"], - direction=obj_data["direction"], - threshold=obj_data.get("threshold"), - priority=obj_data.get("priority", "primary"), - introduced_at=obj_data.get("introduced_at", 0), - source=obj_data.get("source", "initial"), - )) - # Restore history (register() creates its own history entries, so we - # need to replace with the saved history) - saved_history = raw.get("journal", {}).get("preference_history", []) - registry._history = saved_history - - # Load current network if saved - current_network = None - network_file = raw.get("current_network_file") - if network_file: - network_path = save_dir / network_file - if network_path.exists(): - current_network = parse_matpower(network_path) - - return { - "goal": raw["goal"], - "application": raw["application"], - "base_case_path": Path(raw["base_case_path"]), - "config_path": raw.get("config_path"), - "journal_entries": journal_entries, - "objective_registry": registry, - "steering_history": raw.get("steering_history", []), - "active_steering_directives": raw.get("active_steering_directives", []), - "current_network": current_network, - "last_iteration": raw.get("last_iteration", 0), - "total_prompt_tokens": raw.get("total_prompt_tokens", 0), - "total_completion_tokens": raw.get("total_completion_tokens", 0), - "enforced_vmin": raw.get("enforced_vmin"), - "enforced_vmax": raw.get("enforced_vmax"), - "termination_reason": raw.get("termination_reason", "paused"), - "tcopflow_period_data": raw.get("tcopflow_period_data"), - "tcopflow_dT_min": raw.get("tcopflow_dT_min", 0.0), - "tcopflow_duration_min": raw.get("tcopflow_duration_min", 0.0), - "tcopflow_is_coupling": raw.get("tcopflow_is_coupling", True), - "sopflow_num_scenarios": raw.get("sopflow_num_scenarios", 0), - "sopflow_scenario_override": raw.get("sopflow_scenario_override"), - "tcopflow_profile_overrides": raw.get("tcopflow_profile_overrides"), - "benchmark_result": raw.get("benchmark_result"), - "explore_cache_info": raw.get("explore_cache_info"), - } + with log_scope(logger, "Session load", save_dir=save_dir): + session_path = save_dir / "session.json" + if not session_path.exists(): + raise FileNotFoundError(f"Session file not found: {session_path}") + + raw = json.loads(session_path.read_text(encoding="utf-8")) + + version = raw.get("format_version", "0") + if version not in ("1.0", "1.1", SESSION_FORMAT_VERSION): + raise ValueError( + f"Unsupported session format version '{version}' " + f"(expected '{SESSION_FORMAT_VERSION}')" + ) + + # Reconstruct journal entries + journal_entries = [] + for entry_data in raw.get("journal", {}).get("entries", []): + entry = JournalEntry( + iteration=entry_data["iteration"], + description=entry_data["description"], + commands=entry_data.get("commands", []), + objective_value=entry_data.get("objective_value"), + feasible=entry_data.get("feasible", False), + convergence_status=entry_data.get("convergence_status", "UNKNOWN"), + violations_count=entry_data.get("violations_count", 0), + voltage_min=entry_data.get("voltage_min", 0.0), + voltage_max=entry_data.get("voltage_max", 0.0), + max_line_loading_pct=entry_data.get("max_line_loading_pct", 0.0), + total_gen_mw=entry_data.get("total_gen_mw", 0.0), + total_load_mw=entry_data.get("total_load_mw", 0.0), + llm_reasoning=entry_data.get("llm_reasoning", ""), + mode=entry_data.get("mode", "fresh"), + elapsed_seconds=entry_data.get("elapsed_seconds", 0.0), + timestamp=entry_data.get("timestamp", ""), + steering_directive=entry_data.get("steering_directive"), + tracked_metrics=entry_data.get("tracked_metrics"), + feasibility_detail=entry_data.get("feasibility_detail", ""), + solver=entry_data.get("solver", ""), + num_steps=entry_data.get("num_steps", 0), + num_scenarios=entry_data.get("num_scenarios", 0), + explored_variants=entry_data.get("explored_variants"), + # Previously dropped on resume — restore so sweep/contingency/reserve + # metadata survives a save/load round-trip (dataclass defaults keep + # older sessions loadable). + candidate_count=entry_data.get("candidate_count", 0), + feasible_buses=entry_data.get("feasible_buses"), + exago_command=entry_data.get("exago_command"), + contingency_meta=entry_data.get("contingency_meta"), + reserve_meta=entry_data.get("reserve_meta"), + findings=entry_data.get("findings"), + skipped_commands=entry_data.get("skipped_commands"), + ) + journal_entries.append(entry) + + # Reconstruct objective registry + registry = ObjectiveRegistry() + for obj_data in raw.get("journal", {}).get("objective_registry", []): + registry.register(ObjectiveEntry( + name=obj_data["name"], + direction=obj_data["direction"], + threshold=obj_data.get("threshold"), + priority=obj_data.get("priority", "primary"), + introduced_at=obj_data.get("introduced_at", 0), + source=obj_data.get("source", "initial"), + )) + # Restore history (register() creates its own history entries, so we + # need to replace with the saved history) + saved_history = raw.get("journal", {}).get("preference_history", []) + registry._history = saved_history + + # Load current network if saved + current_network = None + network_file = raw.get("current_network_file") + if network_file: + network_path = save_dir / network_file + if network_path.exists(): + current_network = parse_matpower(network_path) + + return { + "goal": raw["goal"], + "application": raw["application"], + "base_case_path": Path(raw["base_case_path"]), + "config_path": raw.get("config_path"), + "journal_entries": journal_entries, + "objective_registry": registry, + "steering_history": raw.get("steering_history", []), + "active_steering_directives": raw.get("active_steering_directives", []), + "current_network": current_network, + "last_iteration": raw.get("last_iteration", 0), + "total_prompt_tokens": raw.get("total_prompt_tokens", 0), + "total_completion_tokens": raw.get("total_completion_tokens", 0), + "token_usage": raw.get("token_usage"), # None before format 1.2 + "enforced_vmin": raw.get("enforced_vmin"), + "enforced_vmax": raw.get("enforced_vmax"), + "termination_reason": raw.get("termination_reason", "paused"), + "tcopflow_period_data": raw.get("tcopflow_period_data"), + "tcopflow_dT_min": raw.get("tcopflow_dT_min", 0.0), + "tcopflow_duration_min": raw.get("tcopflow_duration_min", 0.0), + "tcopflow_is_coupling": raw.get("tcopflow_is_coupling", True), + "sopflow_num_scenarios": raw.get("sopflow_num_scenarios", 0), + "sopflow_scenario_override": raw.get("sopflow_scenario_override"), + "tcopflow_profile_overrides": raw.get("tcopflow_profile_overrides"), + "benchmark_result": raw.get("benchmark_result"), + "explore_cache_info": raw.get("explore_cache_info"), + } diff --git a/agentigrid/agentigrid/exago_parsers/__init__.py b/agentigrid/agentigrid/exago_parsers/__init__.py index 077b7433..7a687b62 100644 --- a/agentigrid/agentigrid/exago_parsers/__init__.py +++ b/agentigrid/agentigrid/exago_parsers/__init__.py @@ -4,6 +4,7 @@ import logging +from agentigrid.logging_setup import log_scope from agentigrid.exago_parsers.matpower_model import ( Branch, Bus, @@ -68,38 +69,43 @@ def parse_simulation_result_for_app(sim_result, application: str, bus_limits=Non Returns the OPFLOWResult portion of the tuple (metadata available separately). Falls back to OPFLOW parser for unknown applications with a warning. """ - if application in _OPFLOW_COMPATIBLE_APPS: + with log_scope( + _logger, "ExaGO result parsing", + application=application, + run_dir=getattr(sim_result, "workdir", None), + ): + if application in _OPFLOW_COMPATIBLE_APPS: + return parse_simulation_result(sim_result, bus_limits=bus_limits) + if application == "scopflow": + parsed = parse_scopflow_simulation_result(sim_result, bus_limits=bus_limits) + if parsed is None: + return None + opflow_result, _metadata = parsed + return opflow_result + if application == "tcopflow": + parsed = parse_tcopflow_simulation_result(sim_result, bus_limits=bus_limits) + if parsed is None: + return None + opflow_result, _metadata = parsed + return opflow_result + if application == "sopflow": + parsed = parse_sopflow_simulation_result(sim_result, bus_limits=bus_limits) + if parsed is None: + return None + opflow_result, _metadata = parsed + return opflow_result + if application == "pflow": + parsed = parse_pflow_simulation_result(sim_result, bus_limits=bus_limits) + if parsed is None: + return None + opflow_result, _metadata = parsed + return opflow_result + _logger.warning( + "Unknown application '%s' for parse_simulation_result_for_app — " + "falling back to OPFLOW parser.", + application, + ) return parse_simulation_result(sim_result, bus_limits=bus_limits) - if application == "scopflow": - parsed = parse_scopflow_simulation_result(sim_result, bus_limits=bus_limits) - if parsed is None: - return None - opflow_result, _metadata = parsed - return opflow_result - if application == "tcopflow": - parsed = parse_tcopflow_simulation_result(sim_result, bus_limits=bus_limits) - if parsed is None: - return None - opflow_result, _metadata = parsed - return opflow_result - if application == "sopflow": - parsed = parse_sopflow_simulation_result(sim_result, bus_limits=bus_limits) - if parsed is None: - return None - opflow_result, _metadata = parsed - return opflow_result - if application == "pflow": - parsed = parse_pflow_simulation_result(sim_result, bus_limits=bus_limits) - if parsed is None: - return None - opflow_result, _metadata = parsed - return opflow_result - _logger.warning( - "Unknown application '%s' for parse_simulation_result_for_app — " - "falling back to OPFLOW parser.", - application, - ) - return parse_simulation_result(sim_result, bus_limits=bus_limits) def parse_scopflow_metadata(sim_result) -> dict | None: diff --git a/agentigrid/agentigrid/exago_parsers/matpower_parser.py b/agentigrid/agentigrid/exago_parsers/matpower_parser.py index 49576d1d..78d48f9f 100644 --- a/agentigrid/agentigrid/exago_parsers/matpower_parser.py +++ b/agentigrid/agentigrid/exago_parsers/matpower_parser.py @@ -180,7 +180,9 @@ def parse_matpower(path: Path) -> MATNetwork: if not path.exists(): raise FileNotFoundError(f"MATPOWER file not found: {path}") - text = path.read_text(encoding="utf-8") + # Exported cases may carry non-UTF-8 bytes in comments (e.g. Windows curly + # quotes in citations); replace them rather than reject the file. + text = path.read_text(encoding="utf-8", errors="replace") # --- Function name --- m = _FUNC_RE.search(text) diff --git a/agentigrid/agentigrid/gridkit_engine/agent_loop.py b/agentigrid/agentigrid/gridkit_engine/agent_loop.py index 3b350d77..6de92328 100644 --- a/agentigrid/agentigrid/gridkit_engine/agent_loop.py +++ b/agentigrid/agentigrid/gridkit_engine/agent_loop.py @@ -20,6 +20,7 @@ from pathlib import Path from typing import Optional +from agentigrid.logging_setup import log_scope from agentigrid.config import AppConfig from agentigrid.gridkit_engine import criteria from agentigrid.gridkit_engine.commands import FaultStudy, RecordVariables, parse_fault_study @@ -40,6 +41,7 @@ from agentigrid.gridkit_parsers.network_summary import network_summary from agentigrid.gridkit_parsers.results_parser import read_results from agentigrid.prompts.gridkit import build_system_prompt +from agentigrid.token_usage import TokenLedger logger = logging.getLogger("agentigrid.gridkit_engine.agent_loop") @@ -58,8 +60,9 @@ class GridkitSession: error: Optional[str] = None final_answer: Optional[dict] = None journal_path: Optional[Path] = None - total_prompt_tokens: int = 0 + total_prompt_tokens: int = 0 # all LLM calls of the session total_completion_tokens: int = 0 + token_ledger: TokenLedger = field(default_factory=TokenLedger) screens: list[str] = field(default_factory=list) # verdict lines, in order steady_state: Optional[str] = None # ExaGO source, when started from it @@ -101,34 +104,36 @@ def _print(self, msg: str) -> None: logger.info(msg) # ------------------------------------------------------------------ + @log_scope(logger, "GridKit session") def run(self, case_file: Path | str, goal: str, steady_state: Optional[SteadyStateSource] = None) -> GridkitSession: """Run the study on *case_file*. With *steady_state*, the case's dynamic models start from that ExaGO solution (written as from_exago.case.json in the session folder) instead of the case's own operating point.""" - case_file = Path(case_file) - binary = self._settings.dynamic_simulation - if not binary.exists(): - raise FileNotFoundError( - f"GridKit DynamicSimulation not found at {binary} " - "(see applications/gridkit/README.md)" - ) - stamp = datetime.now().strftime("%Y%m%d_%H%M%S") - session_dir = self._settings.workdir / f"session_{stamp}" - self._goal = goal - self._source_case_file = case_file - source_text = None - if steady_state is not None: - case, notes = build_from_exago(load_case(case_file), steady_state) - case_file = write_case(case, session_dir / "from_exago.case.json") - source_text = "\n".join([f"Operating point: {steady_state.describe()}.", - *(f" {n}" for n in notes)]) - self._case = load_case(case_file) - self._emap, _ = create_element_map(case_file, session_dir) - self._case_name = self._emap.get("case_name") or case_file.stem - self._session_dir = session_dir - self._journal = GridkitJournal(goal, case_file, session_dir) - session = GridkitSession(goal=goal, case_file=case_file, session_dir=session_dir) + with log_scope(logger, "GridKit case preparation", case_file=case_file): + case_file = Path(case_file) + binary = self._settings.dynamic_simulation + if not binary.exists(): + raise FileNotFoundError( + f"GridKit DynamicSimulation not found at {binary} " + "(see applications/gridkit/README.md)" + ) + stamp = datetime.now().strftime("%Y%m%d_%H%M%S") + session_dir = self._settings.workdir / f"session_{stamp}" + self._goal = goal + self._source_case_file = case_file + source_text = None + if steady_state is not None: + case, notes = build_from_exago(load_case(case_file), steady_state) + case_file = write_case(case, session_dir / "from_exago.case.json") + source_text = "\n".join([f"Operating point: {steady_state.describe()}.", + *(f" {n}" for n in notes)]) + self._case = load_case(case_file) + self._emap, _ = create_element_map(case_file, session_dir) + self._case_name = self._emap.get("case_name") or case_file.stem + self._session_dir = session_dir + self._journal = GridkitJournal(goal, case_file, session_dir) + session = GridkitSession(goal=goal, case_file=case_file, session_dir=session_dir) self._session = session if source_text: session.steady_state = steady_state.describe() @@ -162,41 +167,44 @@ def run(self, case_file: Path | str, goal: str, feedback: Optional[str] = None parse_failures = 0 for iteration in range(1, self._settings.max_iterations + 1): - user_prompt = self._user_prompt(goal, latest, feedback, iteration) - feedback = None - self._print(f"[Iter {iteration}] Asking {self._backend.name()} ({self._cfg.llm.model})...") - response = self._backend.complete(system_prompt, user_prompt) - session.total_prompt_tokens += response.prompt_tokens or 0 - session.total_completion_tokens += response.completion_tokens or 0 - data = response.json_data - if data is None and _BACKEND_ERROR_RE.match(response.raw_text or ""): - session.termination_reason = "llm_error" - session.error = response.raw_text.strip() - self._journal.add(JournalEntry(iteration, "error", "LLM call failed", - result_text=session.error)) - break - if data is None: - parse_failures += 1 - self._journal.add(JournalEntry(iteration, "error", "reply was not JSON", - result_text=response.raw_text[:2000])) - if parse_failures >= _MAX_CONSECUTIVE_PARSE_FAILURES: - session.termination_reason = "parse_failures" + with log_scope(logger, "GridKit iteration", iteration=iteration): + user_prompt = self._user_prompt(goal, latest, feedback, iteration) + feedback = None + self._print(f"[Iter {iteration}] Asking {self._backend.name()} ({self._cfg.llm.model})...") + response = self._backend.complete(system_prompt, user_prompt) + session.token_ledger.record("search", response, iteration, + prompt_chars=len(system_prompt) + len(user_prompt)) + session.total_prompt_tokens = session.token_ledger.total_prompt + session.total_completion_tokens = session.token_ledger.total_completion + data = response.json_data + if data is None and _BACKEND_ERROR_RE.match(response.raw_text or ""): + session.termination_reason = "llm_error" + session.error = response.raw_text.strip() + self._journal.add(JournalEntry(iteration, "error", "LLM call failed", + result_text=session.error)) break - feedback = "Your reply was not a valid JSON object. Reply with exactly one JSON object." - continue - parse_failures = 0 - action = str(data.get("action", "")).lower() - if action == "fault_screen": - latest, feedback = self._handle_fault_screen(iteration, data) - if latest: - session.screens.append(latest.splitlines()[0]) - elif action == "complete": - session.final_answer = self._handle_complete(iteration, data) - session.termination_reason = "completed" - break - else: - feedback = f"Unknown action '{action}'. Valid actions: fault_screen, complete." - self._journal.add(JournalEntry(iteration, "error", feedback, request=data)) + if data is None: + parse_failures += 1 + self._journal.add(JournalEntry(iteration, "error", "reply was not JSON", + result_text=response.raw_text[:2000])) + if parse_failures >= _MAX_CONSECUTIVE_PARSE_FAILURES: + session.termination_reason = "parse_failures" + break + feedback = "Your reply was not a valid JSON object. Reply with exactly one JSON object." + continue + parse_failures = 0 + action = str(data.get("action", "")).lower() + if action == "fault_screen": + latest, feedback = self._handle_fault_screen(iteration, data) + if latest: + session.screens.append(latest.splitlines()[0]) + elif action == "complete": + session.final_answer = self._handle_complete(iteration, data) + session.termination_reason = "completed" + break + else: + feedback = f"Unknown action '{action}'. Valid actions: fault_screen, complete." + self._journal.add(JournalEntry(iteration, "error", feedback, request=data)) else: session.termination_reason = "max_iterations" return self._finalize(session) @@ -205,32 +213,39 @@ def run(self, case_file: Path | str, goal: str, def _run_base_case(self, case: dict, emap: dict, run_dir: Path, iteration: int = 0, what: str = "no fault") -> tuple[str, bool]: """No-fault run of a copy of *case*: it must solve and stay still.""" - work, report = apply_modifications( - case, [RecordVariables(machines=["delta", "omega"], buses=["Vm"])], emap, - ) - case_path = write_case(work, run_dir / "case.json") - job = RunJob(run_dir, case_path, tmax_s=self._settings.tmax_s, - dt_monitor_s=self._settings.dt_monitor_s) - self._print(f"[Iter {iteration}] Base case: {what}, {self._settings.tmax_s:g} s ...") - res = run_one(job, self._settings) - if res.status.solved: - series = read_results(res.csv_path, emap) - steady = criteria.check_steady_state( - series.t, series.device_variable("delta", MACHINE_CLASSES), series.bus_variable("Vm"), + with log_scope(logger, "GridKit base case", iteration=iteration, run_dir=run_dir): + work, report = apply_modifications( + case, [RecordVariables(machines=["delta", "omega"], buses=["Vm"])], emap, ) - ok = steady.status == "pass" - text = f"GridKit {res.status.reason}. Steady state: {steady.status} - {steady.detail}." - else: - ok = False - text = f"GridKit base case did not solve: {res.status.reason}." - if not ok: - text += " Fault screens are not meaningful until the base case solves and stays steady." - self._print(f"[Iter {iteration}] {text}") - self._journal.add(JournalEntry( - iteration, "base_case", what, verdict=text, result_text=text, - elapsed_s=res.elapsed_s, run_dir=str(run_dir), - )) - return text, ok + case_path = write_case(work, run_dir / "case.json") + job = RunJob(run_dir, case_path, tmax_s=self._settings.tmax_s, + dt_monitor_s=self._settings.dt_monitor_s) + self._print(f"[Iter {iteration}] Base case: {what}, {self._settings.tmax_s:g} s ...") + res = run_one(job, self._settings) + rows = [] + if res.status.solved: + series = read_results(res.csv_path, emap) + steady = criteria.check_steady_state( + series.t, series.device_variable("delta", MACHINE_CLASSES), series.bus_variable("Vm"), + ) + ok = steady.status == "pass" + text = f"GridKit {res.status.reason}. Steady state: {steady.status} - {steady.detail}." + rows = [{"bus": None, "status": steady.status, "solved": True, "run": res.status.reason, + "checks": {steady.name: {"status": steady.status, "value": steady.value, + "limit": steady.limit, "detail": steady.detail, + "data": steady.data}}, + "run_dir": str(run_dir)}] + else: + ok = False + text = f"GridKit base case did not solve: {res.status.reason}." + if not ok: + text += " Fault screens are not meaningful until the base case solves and stays steady." + self._print(f"[Iter {iteration}] {text}") + self._journal.add(JournalEntry( + iteration, "base_case", what, verdict=text, result_text=text, rows=rows, + elapsed_s=res.elapsed_s, run_dir=str(run_dir), + )) + return text, ok def _start(self, start_from: str, iteration: int) -> _Start: """The operating point a screen asked for, built (and checked) once.""" @@ -276,95 +291,97 @@ def _user_prompt(self, goal: str, latest: Optional[str], feedback: Optional[str] # ------------------------------------------------------------------ def _handle_fault_screen(self, iteration: int, data: dict) -> tuple[Optional[str], Optional[str]]: """Returns (result text for the LLM, error feedback).""" - description = data.get("description", "fault screen") - try: - study = parse_fault_study(data, self._settings.study_defaults()) - except ValueError as exc: - self._journal.add(JournalEntry(iteration, "error", str(exc), request=data)) - return None, str(exc) - if study.start_from == "case" and goal_wants_steady_state(self._goal) \ - and "latest_steady_state" not in self._starts: - msg = ('fault_screen rejected: the goal asks for the most recent steady state, ' - 'so set "start_from": "latest_steady_state".') - self._journal.add(JournalEntry(iteration, "error", msg, request=data)) - return None, msg - ca = study.application == "ContingencyAnalysis" - if ca and not self._settings.contingency_analysis.exists(): - msg = (f"fault_screen rejected: GridKit ContingencyAnalysis not found at " - f"{self._settings.contingency_analysis}; use DynamicSimulation.") - self._journal.add(JournalEntry(iteration, "error", msg, request=data)) - return None, msg - start = self._start(study.start_from, iteration) - if not start.ok: - text = (f"The fault screen was NOT run: {start.label} is not usable.\n{start.text}") - self._journal.add(JournalEntry(iteration, "fault_screen", f"{description} (not run)", - request=data, verdict=text.splitlines()[0], - result_text=text)) + with log_scope(logger, "GridKit action", iteration=iteration, action="fault_screen"): + description = data.get("description", "fault screen") + try: + study = parse_fault_study(data, self._settings.study_defaults()) + except ValueError as exc: + self._journal.add(JournalEntry(iteration, "error", str(exc), request=data)) + return None, str(exc) + if study.start_from == "case" and goal_wants_steady_state(self._goal) \ + and "latest_steady_state" not in self._starts: + msg = ('fault_screen rejected: the goal asks for the most recent steady state, ' + 'so set "start_from": "latest_steady_state".') + self._journal.add(JournalEntry(iteration, "error", msg, request=data)) + return None, msg + ca = study.application == "ContingencyAnalysis" + if ca and not self._settings.contingency_analysis.exists(): + msg = (f"fault_screen rejected: GridKit ContingencyAnalysis not found at " + f"{self._settings.contingency_analysis}; use DynamicSimulation.") + self._journal.add(JournalEntry(iteration, "error", msg, request=data)) + return None, msg + start = self._start(study.start_from, iteration) + if not start.ok: + text = (f"The fault screen was NOT run: {start.label} is not usable.\n{start.text}") + self._journal.add(JournalEntry(iteration, "fault_screen", f"{description} (not run)", + request=data, verdict=text.splitlines()[0], + result_text=text)) + return text, None + if study.start_from == "latest_steady_state" and not self._session.steady_state: + self._session.steady_state = start.label + check = validate_fault_study(study, start.emap) + if not check.valid: + msg = "fault_screen rejected: " + "; ".join(check.errors) + self._journal.add(JournalEntry(iteration, "error", msg, request=data)) + return None, msg + + key = json.dumps(asdict(study), sort_keys=True) + if key in self._screen_cache: + first_iter, text = self._screen_cache[key] + text = (f"This identical fault screen already ran in iteration {first_iter}; " + f"its result is repeated below. Do not run it again.\n{text}") + self._journal.add(JournalEntry(iteration, "fault_screen", f"{description} (cached)", + request=data, verdict=text.splitlines()[1], + result_text=text)) + return text, None + + iter_dir = self._session_dir / f"iter_{iteration:03d}" + work, wmap, fault_ids, report = prepare_fault_case(start.case, study, start.emap, + only_study_faults=ca) + if report.errors: + msg = "fault_screen could not build the case copy: " + "; ".join(report.errors) + self._journal.add(JournalEntry(iteration, "error", msg, request=data)) + return None, msg + case_path = write_case(work, iter_dir / "case.json") + (iter_dir / "element_map.json").write_text(json.dumps(wmap, indent=1)) + + clear = study.clear_time_s() + device_ids = {f["element_id"]: f["id"] for f in wmap["bus_faults"]} + jobs = [ + RunJob(iter_dir / ("contingency" if ca else f"bus_{bus}"), case_path, + study.tmax_s, study.dt_monitor_s, element_id=eid, start_s=study.start_s, + clear_s=clear, bus=bus, fault_id=device_ids.get(eid)) + for bus, eid in fault_ids.items() + ] + how = ("ContingencyAnalysis, one process" if ca + else f"{self._settings.resolved_workers(len(jobs))} at a time") + self._print(f'[Iter {iteration}] fault_screen "{description}": {len(jobs)} bus faults, {how} ...') + t0 = time.monotonic() + results = run_contingency(jobs, self._settings) if ca else run_many(jobs, self._settings) + with log_scope(logger, "GridKit result evaluation", iteration=iteration, runs=len(results)): + outcomes = [self._judge(r, wmap, study, clear) for r in results] + elapsed = time.monotonic() - t0 + + cycles = study.clearing_cycles + study.margin_cycles + note = (f"bus short circuit to ground at {study.start_s:g} s, cleared after {cycles:g} cycles " + f"({criteria.num(study.duration_s(), '.3f', ' s')}), R = {study.R:g} pu, X = {study.X:g} pu, " + f"{study.tmax_s:g} s simulated, GridKit {study.application}, " + f"starting from {start.label}") + text = screen_verdict(self._subject(study), outcomes, note) + if check.warnings: + text += "\n\nWARNINGS:\n" + "\n".join(f" {w}" for w in check.warnings) + if len(outcomes) <= _FULL_TABLE_MAX: + text += "\n\n" + _outcome_table(outcomes) + self._print(f"[Iter {iteration}] {text.splitlines()[0]} ({elapsed:.1f} s)") + + self._screen_cache[key] = (iteration, text) + self._journal.add(JournalEntry( + iteration, "fault_screen", description, request=data, + verdict=text.splitlines()[0], result_text=text, + rows=[outcome_row(o, r.job.run_dir) for o, r in zip(outcomes, results)], + elapsed_s=elapsed, llm_reasoning=str(data.get("reasoning", "")), run_dir=str(iter_dir), + )) return text, None - if study.start_from == "latest_steady_state" and not self._session.steady_state: - self._session.steady_state = start.label - check = validate_fault_study(study, start.emap) - if not check.valid: - msg = "fault_screen rejected: " + "; ".join(check.errors) - self._journal.add(JournalEntry(iteration, "error", msg, request=data)) - return None, msg - - key = json.dumps(asdict(study), sort_keys=True) - if key in self._screen_cache: - first_iter, text = self._screen_cache[key] - text = (f"This identical fault screen already ran in iteration {first_iter}; " - f"its result is repeated below. Do not run it again.\n{text}") - self._journal.add(JournalEntry(iteration, "fault_screen", f"{description} (cached)", - request=data, verdict=text.splitlines()[1], - result_text=text)) - return text, None - - iter_dir = self._session_dir / f"iter_{iteration:03d}" - work, wmap, fault_ids, report = prepare_fault_case(start.case, study, start.emap, - only_study_faults=ca) - if report.errors: - msg = "fault_screen could not build the case copy: " + "; ".join(report.errors) - self._journal.add(JournalEntry(iteration, "error", msg, request=data)) - return None, msg - case_path = write_case(work, iter_dir / "case.json") - (iter_dir / "element_map.json").write_text(json.dumps(wmap, indent=1)) - - clear = study.clear_time_s() - device_ids = {f["element_id"]: f["id"] for f in wmap["bus_faults"]} - jobs = [ - RunJob(iter_dir / ("contingency" if ca else f"bus_{bus}"), case_path, - study.tmax_s, study.dt_monitor_s, element_id=eid, start_s=study.start_s, - clear_s=clear, bus=bus, fault_id=device_ids.get(eid)) - for bus, eid in fault_ids.items() - ] - how = ("ContingencyAnalysis, one process" if ca - else f"{self._settings.resolved_workers(len(jobs))} at a time") - self._print(f'[Iter {iteration}] fault_screen "{description}": {len(jobs)} bus faults, {how} ...') - t0 = time.monotonic() - results = run_contingency(jobs, self._settings) if ca else run_many(jobs, self._settings) - outcomes = [self._judge(r, wmap, study, clear) for r in results] - elapsed = time.monotonic() - t0 - - cycles = study.clearing_cycles + study.margin_cycles - note = (f"bus short circuit to ground at {study.start_s:g} s, cleared after {cycles:g} cycles " - f"({study.duration_s():.3f} s), R = {study.R:g} pu, X = {study.X:g} pu, " - f"{study.tmax_s:g} s simulated, GridKit {study.application}, " - f"starting from {start.label}") - text = screen_verdict(self._subject(study), outcomes, note) - if check.warnings: - text += "\n\nWARNINGS:\n" + "\n".join(f" {w}" for w in check.warnings) - if len(outcomes) <= _FULL_TABLE_MAX: - text += "\n\n" + _outcome_table(outcomes) - self._print(f"[Iter {iteration}] {text.splitlines()[0]} ({elapsed:.1f} s)") - - self._screen_cache[key] = (iteration, text) - self._journal.add(JournalEntry( - iteration, "fault_screen", description, request=data, - verdict=text.splitlines()[0], result_text=text, - rows=[outcome_row(o, r.job.run_dir) for o, r in zip(outcomes, results)], - elapsed_s=elapsed, llm_reasoning=str(data.get("reasoning", "")), run_dir=str(iter_dir), - )) - return text, None def _judge(self, res: RunResult, wmap: dict, study: FaultStudy, clear: float) -> FaultOutcome: checks = [] @@ -386,7 +403,9 @@ def _handle_complete(self, iteration: int, data: dict) -> dict: verdict=str(answer["summary"]), findings=answer)) return answer + @log_scope(logger, "GridKit finalization") def _finalize(self, session: GridkitSession) -> GridkitSession: + self._journal.token_usage = session.token_ledger.to_dict() session.journal_path = self._journal.export_json() print() print("=" * 60) @@ -406,8 +425,7 @@ def _finalize(self, session: GridkitSession) -> GridkitSession: print(" Details:") for line in details.splitlines(): print(f" {line}") - print(f" Tokens: {session.total_prompt_tokens} prompt + " - f"{session.total_completion_tokens} completion") + print(f" {session.token_ledger.summary_text()}") print(f" Journal: {session.journal_path}") return session @@ -420,18 +438,18 @@ def _exago_label(source: SteadyStateSource) -> str: def _outcome_table(outcomes: list[FaultOutcome]) -> str: """Every fault, one row: result and the main check values.""" - def val(o: FaultOutcome, name: str, fmt: str) -> str: + def val(o: FaultOutcome, name: str, fmt: str, unit: str = "", scale: float = 1.0) -> str: c = next((c for c in o.checks if c.name == name), None) if c is None or c.value is None: return "-" if c is None else c.status - return fmt.format(c.value) + return criteria.num(scale * c.value, fmt, unit) - lines = ["ALL BUS FAULTS:", - f"{'Bus':>6} | {'Result':<10} | {'Angle spread':>12} | {'Vmin 2.5s':>9} | {'Damping':>8}", - "-" * 58] + lines = ["ALL BUS FAULTS (rounded value, full value in brackets):", + f"{'Bus':>6} | {'Result':<10} | {'Angle spread':>32} | {'Vmin 2.5s':>28} | {'Damping':>28}", + "-" * 116] for o in sorted(outcomes, key=lambda o: o.bus): lines.append( - f"{o.bus:>6} | {o.status:<10} | {val(o, 'angle_stability', '{:.1f} deg'):>12} | " - f"{val(o, 'voltage_recovery', '{:.3f}'):>9} | {val(o, 'damping', '{:.1%}'):>8}" + f"{o.bus:>6} | {o.status:<10} | {val(o, 'angle_stability', '.1f', ' deg'):>32} | " + f"{val(o, 'voltage_recovery', '.3f', ' pu'):>28} | {val(o, 'damping', '.1f', '%', 100):>28}" ) return "\n".join(lines) diff --git a/agentigrid/agentigrid/gridkit_engine/criteria.py b/agentigrid/agentigrid/gridkit_engine/criteria.py index 77b90564..2b69539b 100644 --- a/agentigrid/agentigrid/gridkit_engine/criteria.py +++ b/agentigrid/agentigrid/gridkit_engine/criteria.py @@ -28,7 +28,7 @@ from __future__ import annotations import math -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Mapping, Optional, Sequence # --------------------------------------------------------------------------- @@ -105,12 +105,19 @@ class CheckResult: limit: Optional[float] detail: str source: str + data: dict = field(default_factory=dict) # every number in *detail*, unrounded (journal) @property def passed(self) -> bool: return self.status == "pass" +def num(x: float, fmt: str, unit: str = "") -> str: + """*x* rounded with *fmt* for reading, then its full value in brackets, so a + rounded 1.100 that is really 1.10036 stays visible.""" + return f"{x:{fmt}}{unit} [{float(x)!r}{unit}]" + + def _series_after(t: Sequence[float], x: Sequence[float], t0: float) -> list[tuple[float, float]]: return [(ti, xi) for ti, xi in zip(t, x) if ti >= t0] @@ -135,9 +142,10 @@ def check_angle_stability( ok = worst <= limit_deg return CheckResult( "angle_stability", "pass" if ok else "fail", worst, limit_deg, - f"largest machine angle spread {worst:.1f} deg at t = {worst_t:g} s" + f"largest machine angle spread {num(worst, '.1f', ' deg')} at t = {worst_t:g} s" + ("" if ok else " (loss of synchronism)"), src, + {"spread_deg": worst, "t_s": worst_t}, ) @@ -165,9 +173,10 @@ def check_voltage_recovery( ok = worst_v >= min_pu return CheckResult( "voltage_recovery", "pass" if ok else "fail", worst_v, min_pu, - f"lowest voltage from {after_s:g} s after clearing: {worst_v:.3f} pu " + f"lowest voltage from {after_s:g} s after clearing: {num(worst_v, '.3f', ' pu')} " f"at bus {worst_bus}, t = {worst_t:g} s", src, + {"lowest_pu": worst_v, "bus": worst_bus, "t_s": worst_t}, ) @@ -196,8 +205,9 @@ def check_voltage_envelope( bus, ti, vi, vmin = worst return CheckResult( name, "pass" if ok else "fail", vi, vmin, - f"closest to envelope: bus {bus} at t = {ti:g} s, {vi:.3f} pu vs minimum {vmin:.2f} pu", + f"closest to envelope: bus {bus} at t = {ti:g} s, {num(vi, '.3f', ' pu')} vs minimum {vmin:.2f} pu", source, + {"v_pu": vi, "bus": bus, "t_s": ti, "envelope_min_pu": vmin, "margin_pu": worst_margin}, ) @@ -216,9 +226,11 @@ def check_final_voltage( ok = finals[lo_bus] >= vmin and finals[hi_bus] <= vmax return CheckResult( "final_voltage", "pass" if ok else "fail", finals[lo_bus], vmin, - f"end of run: lowest {finals[lo_bus]:.3f} pu (bus {lo_bus}), " - f"highest {finals[hi_bus]:.3f} pu (bus {hi_bus}), band {vmin}-{vmax} pu", + f"end of run: lowest {num(finals[lo_bus], '.3f', ' pu')} (bus {lo_bus}), " + f"highest {num(finals[hi_bus], '.3f', ' pu')} (bus {hi_bus}), band {vmin}-{vmax} pu", src, + {"lowest_pu": finals[lo_bus], "lowest_bus": lo_bus, + "highest_pu": finals[hi_bus], "highest_bus": hi_bus, "vmin_pu": vmin, "vmax_pu": vmax}, ) @@ -302,11 +314,12 @@ def check_damping( f"all swings below {min_amplitude_deg} deg after clearing", src) zeta, nm = worst ok = zeta >= min_ratio and not unmeasured - detail = f"lowest damping ratio {100 * zeta:.1f}% (machine {nm})" + detail = f"lowest damping ratio {num(100 * zeta, '.1f', '%')} (machine {nm})" if unmeasured: detail += f"; not measurable for {', '.join(unmeasured[:5])}" status = "pass" if ok else ("fail" if zeta < min_ratio else "not_tested") - return CheckResult("damping", status, zeta, min_ratio, detail, src) + return CheckResult("damping", status, zeta, min_ratio, detail, src, + {"ratio": zeta, "machine": nm, "unmeasured": unmeasured}) # Pre-fault steady state: the no-fault run of the case copy must stay still, @@ -341,9 +354,10 @@ def check_steady_state( ok = angle_move <= angle_tol_deg and v_move <= v_tol_pu return CheckResult( "steady_state", "pass" if ok else "fail", angle_move, angle_tol_deg, - f"without a fault: angles move {angle_move:.3f} deg (limit {angle_tol_deg}), " - f"voltages move {v_move:.4f} pu (limit {v_tol_pu})", + f"without a fault: angles move {num(angle_move, '.3f', ' deg')} (limit {angle_tol_deg}), " + f"voltages move {num(v_move, '.4f', ' pu')} (limit {v_tol_pu})", src, + {"angle_move_deg": angle_move, "v_move_pu": v_move}, ) # PJM tests the agent does not run, because GridKit lacks the element they diff --git a/agentigrid/agentigrid/gridkit_engine/executor.py b/agentigrid/agentigrid/gridkit_engine/executor.py index ebef55de..7b1b872b 100644 --- a/agentigrid/agentigrid/gridkit_engine/executor.py +++ b/agentigrid/agentigrid/gridkit_engine/executor.py @@ -14,6 +14,7 @@ from pathlib import Path from typing import Optional +from agentigrid.logging_setup import log_scope from agentigrid.gridkit_engine.settings import GridkitSettings from agentigrid.gridkit_parsers.status_parser import RunStatus, parse_run_status @@ -78,11 +79,17 @@ def run_one(job: RunJob, settings: GridkitSettings) -> RunResult: binary = settings.dynamic_simulation.resolve() start = time.monotonic() try: - proc = subprocess.run( - [str(binary), SOLVER_FILE], cwd=job.run_dir, - capture_output=True, text=True, timeout=settings.timeout_s, - ) - exit_code, stdout, stderr = proc.returncode, proc.stdout, proc.stderr + with log_scope( + logger, "GridKit subprocess", + application="DynamicSimulation", + run_dir=job.run_dir, + bus=job.bus, + ): + proc = subprocess.run( + [str(binary), SOLVER_FILE], cwd=job.run_dir, + capture_output=True, text=True, timeout=settings.timeout_s, + ) + exit_code, stdout, stderr = proc.returncode, proc.stdout, proc.stderr except subprocess.TimeoutExpired as exc: exit_code = None stdout = exc.stdout.decode() if isinstance(exc.stdout, bytes) else (exc.stdout or "") @@ -103,7 +110,10 @@ def run_many(jobs: list[RunJob], settings: GridkitSettings) -> list[RunResult]: if not jobs: return [] workers = settings.resolved_workers(len(jobs)) - with ThreadPoolExecutor(max_workers=workers) as pool: + with ( + log_scope(logger, "GridKit simulation batch", tasks=len(jobs), workers=workers), + ThreadPoolExecutor(max_workers=workers) as pool, + ): return list(pool.map(lambda j: run_one(j, settings), jobs)) @@ -133,11 +143,17 @@ def run_contingency(jobs: list[RunJob], settings: GridkitSettings) -> list[RunRe binary = settings.contingency_analysis.resolve() start = time.monotonic() try: - proc = subprocess.run( - [str(binary), SOLVER_FILE], cwd=run_dir, - capture_output=True, text=True, timeout=settings.timeout_s * len(jobs), - ) - exit_code, stdout, stderr = proc.returncode, proc.stdout, proc.stderr + with log_scope( + logger, "GridKit subprocess", + application="ContingencyAnalysis", + run_dir=run_dir, + faults=len(jobs), + ): + proc = subprocess.run( + [str(binary), SOLVER_FILE], cwd=run_dir, + capture_output=True, text=True, timeout=settings.timeout_s * len(jobs), + ) + exit_code, stdout, stderr = proc.returncode, proc.stdout, proc.stderr except subprocess.TimeoutExpired as exc: exit_code = None stdout = exc.stdout.decode() if isinstance(exc.stdout, bytes) else (exc.stdout or "") diff --git a/agentigrid/agentigrid/gridkit_engine/from_exago.py b/agentigrid/agentigrid/gridkit_engine/from_exago.py index 246464b7..5f9738a6 100644 --- a/agentigrid/agentigrid/gridkit_engine/from_exago.py +++ b/agentigrid/agentigrid/gridkit_engine/from_exago.py @@ -23,6 +23,7 @@ import copy import json +import logging import math import re from collections import defaultdict @@ -30,10 +31,13 @@ from pathlib import Path from typing import Optional +from agentigrid.logging_setup import log_scope from agentigrid.exago_parsers.matpower_model import MATNetwork from agentigrid.exago_parsers.matpower_parser import parse_matpower from agentigrid.gridkit_parsers.case_parser import MACHINE_CLASSES, _is_bus_port +logger = logging.getLogger("agentigrid.gridkit_engine.from_exago") + # AC steady states that ExaGO saves as one solved MATPOWER case (-save_output). STEADY_STATE_APPS = {"opflow": "opflowout.m", "pflow": "pflowout.m"} @@ -76,40 +80,41 @@ def find_latest_steady_state(exago_workdir: Path | str) -> SteadyStateSource: Raises SteadyStateError when that simulation is not a single converged AC steady state with a saved solution (no silent fallback to an older one). """ - exago_workdir = Path(exago_workdir) - journals = sorted(exago_workdir.glob("journal_*.json")) - if not journals: - raise SteadyStateError(f"No ExaGO journal in {exago_workdir}; run an ExaGO study first.") - journal = journals[-1] # names carry the timestamp, so sorting by name is by time - entries = json.loads(journal.read_text()).get("entries", []) - simulated = [e for e in entries if e.get("exago_command")] - if not simulated: - raise SteadyStateError(f"{journal.name} has no simulation.") - entry = simulated[-1] - cmd = entry["exago_command"] - it = entry.get("iteration") - app = cmd.get("application", "") - where = f"The most recent ExaGO iteration ({it}, {journal.name})" - if cmd.get("mode") != "single": - raise SteadyStateError(f"{where} is a '{cmd.get('mode')}' run of many solves, " - "not a single steady state.") - if app not in STEADY_STATE_APPS: - raise SteadyStateError(f"{where} ran {app}; only {sorted(STEADY_STATE_APPS)} give an " - "AC steady state GridKit can start from.") - # PFLOW's saved file always has converged = 0 (ExaGO writes the OPFLOW - # flag), so the journal's status is the one to trust. - status = str(entry.get("convergence_status", "")) - if not status.upper().startswith("CONVERGED"): - raise SteadyStateError(f"{where} did not converge (status: {status or 'unknown'}).") - netfile = _netfile(cmd.get("argv") or []) - if netfile is None: - raise SteadyStateError(f"{where}: no -netfile in the recorded command.") - solved = netfile.parent / STEADY_STATE_APPS[app] - if not solved.exists(): - raise SteadyStateError(f"{where}: solution file {solved} is missing " - "(runs before ExaGO saved its output cannot be used).") - return SteadyStateSource(journal, int(it), str(entry.get("description", "")), app, - netfile.parent, solved) + with log_scope(logger, "ExaGO steady state lookup", workdir=exago_workdir): + exago_workdir = Path(exago_workdir) + journals = sorted(exago_workdir.glob("journal_*.json")) + if not journals: + raise SteadyStateError(f"No ExaGO journal in {exago_workdir}; run an ExaGO study first.") + journal = journals[-1] # names carry the timestamp, so sorting by name is by time + entries = json.loads(journal.read_text()).get("entries", []) + simulated = [e for e in entries if e.get("exago_command")] + if not simulated: + raise SteadyStateError(f"{journal.name} has no simulation.") + entry = simulated[-1] + cmd = entry["exago_command"] + it = entry.get("iteration") + app = cmd.get("application", "") + where = f"The most recent ExaGO iteration ({it}, {journal.name})" + if cmd.get("mode") != "single": + raise SteadyStateError(f"{where} is a '{cmd.get('mode')}' run of many solves, " + "not a single steady state.") + if app not in STEADY_STATE_APPS: + raise SteadyStateError(f"{where} ran {app}; only {sorted(STEADY_STATE_APPS)} give an " + "AC steady state GridKit can start from.") + # PFLOW's saved file always has converged = 0 (ExaGO writes the OPFLOW + # flag), so the journal's status is the one to trust. + status = str(entry.get("convergence_status", "")) + if not status.upper().startswith("CONVERGED"): + raise SteadyStateError(f"{where} did not converge (status: {status or 'unknown'}).") + netfile = _netfile(cmd.get("argv") or []) + if netfile is None: + raise SteadyStateError(f"{where}: no -netfile in the recorded command.") + solved = netfile.parent / STEADY_STATE_APPS[app] + if not solved.exists(): + raise SteadyStateError(f"{where}: solution file {solved} is missing " + "(runs before ExaGO saved its output cannot be used).") + return SteadyStateSource(journal, int(it), str(entry.get("description", "")), app, + netfile.parent, solved) # --------------------------------------------------------------------------- @@ -120,14 +125,47 @@ def _signal_ports(device: dict) -> set: return {v for k, v in device.get("ports", {}).items() if not _is_bus_port(k)} -def _machine_unit(device: dict) -> str: - """Machine id within its bus ('49_1_genrou' -> '1'), for ordering.""" - parts = str(device.get("id", "")).split("_") - return parts[1] if len(parts) > 2 else str(device.get("id", "")) +# GridKit cases whose renewables are "_static_gen_" LoadZIP injections +# (no dynamic model). The derived ExaGO case carries them as negative load. +STATIC_GEN_CASES = {"ACTIVSg10k"} + +# A dispatch this close to Pmin/Pmax is taken as at the limit (ExaGO -save_output +# writes Pg with 2 decimals). +PG_LIMIT_SNAP_MW = 0.01 + + +def static_gen_loads(case: dict, loads: list[dict]) -> list[dict]: + """The renewable injections to keep unchanged; empty unless *case* is in + STATIC_GEN_CASES. They are fixed at the case's own output, which the ExaGO + case netted from its load, so only the ordinary load is rebuilt from ExaGO + and each keeps its own load model (constant current vs. impedance).""" + if case.get("header", {}).get("case_name") not in STATIC_GEN_CASES: + return [] + return [d for d in loads if "_static_gen_" in d["id"]] -def _unit_key(unit: str): - return (0, int(unit), "") if unit.isdigit() else (1, 0, unit) +def _snap_to_valve_limits(devices: list[dict], base: float) -> list[str]: + """Machines at a Tgov1 valve limit, set exactly inside it. + + Tgov1 starts from valve position p0 * base / Trate and rejects it outside + [Pvmin, Pvmax] with a 4-epsilon tolerance. Its stored limits differ from the + MATPOWER Pmin/Pmax in the last digits, so a unit dispatched at Pmax can + fail; within PG_LIMIT_SNAP_MW of a limit it is set 1e-10 (relative) inside. + """ + machines = {d["ports"].get("pmech"): d for d in devices if d["class"] in MACHINE_CLASSES} + snapped = [] + for gov in devices: + if gov["class"] != "Tgov1" or gov["ports"].get("pmech") not in machines: + continue + m, p = machines[gov["ports"]["pmech"]], gov["params"] + pg = m["params"]["p0"] * base + for limit, inward in ((p["Pvmax"], -1e-10), (p["Pvmin"], 1e-10)): + mw = limit * p["Trate"] + if abs(pg - mw) < PG_LIMIT_SNAP_MW: + m["params"]["p0"] = mw * (1.0 + inward) / base + snapped.append(m["id"]) + break + return snapped def apply_steady_state(case: dict, net: MATNetwork) -> tuple[dict, list[str]]: @@ -169,11 +207,14 @@ def apply_steady_state(case: dict, net: MATNetwork) -> tuple[dict, list[str]]: for d in devices: if d["class"] in MACHINE_CLASSES: machines_at[d["ports"]["bus"]].append(d) - no_model, ambiguous, removed_machines = [], [], [] + no_model, ambiguous, removed_machines, at_limit = [], [], [], [] for bus in sorted(set(gens_at) | set(machines_at)): gens = gens_at.get(bus, []) online = [g for g in gens if g.status] - machines = sorted(machines_at.get(bus, []), key=lambda d: _unit_key(_machine_unit(d))) + # Machines pair with the MATPOWER rows in file order: both files list a + # bus's units in the source order (ACTIVSg10k: 1, 10, 2, ... 9), which a + # numeric sort of the unit ids would break at buses with 10+ units. + machines = machines_at.get(bus, []) if len(online) > len(machines): no_model.append(bus) continue @@ -190,13 +231,28 @@ def apply_steady_state(case: dict, net: MATNetwork) -> tuple[dict, list[str]]: if g is None or not g.status: removed_machines.append(m) else: - m["params"]["p0"] = g.Pg / base + # ExaGO writes Pg to 2 decimals; a unit dispatched at Pmin/Pmax + # can land just outside it, which GridKit's governors reject. + pg = g.Pg + for limit in (g.Pmax, g.Pmin): + if pg != limit and abs(pg - limit) < PG_LIMIT_SNAP_MW: + at_limit.append(m["id"]) + pg = limit + break + m["params"]["p0"] = pg / base m["params"]["q0"] = g.Qg / base if no_model: raise SteadyStateError( f"Online ExaGO generator(s) at bus(es) {no_model} have no dynamic model " "in the GridKit case; GridKit cannot simulate them." ) + if at_limit: + notes.append(f"{len(at_limit)} machine(s) within {PG_LIMIT_SNAP_MW} MW of Pmin/Pmax set " + f"to the limit (ExaGO's output rounding): {', '.join(at_limit[:10])}") + at_valve = _snap_to_valve_limits(devices, base) + if at_valve: + notes.append(f"{len(at_valve)} machine(s) within {PG_LIMIT_SNAP_MW} MW of their Tgov1 valve " + f"limit set just inside it: {', '.join(at_valve[:10])}") if ambiguous: raise SteadyStateError( f"Cannot match ExaGO generators to GridKit machines at bus(es) {ambiguous} " @@ -219,6 +275,16 @@ def apply_steady_state(case: dict, net: MATNetwork) -> tuple[dict, list[str]]: # --- loads and shunts -------------------------------------------------- old_loads = [d for d in devices if d["class"] == "LoadZIP"] + static = static_gen_loads(case, old_loads) + injection: dict[int, list[float]] = defaultdict(lambda: [0.0, 0.0]) # MW, MVAr + for d in static: + injection[d["ports"]["bus"]][0] += d["params"]["Pnom"] * base + injection[d["ports"]["bus"]][1] += d["params"]["Qnom"] * base + if static: + static_ids = {id(d) for d in static} + old_loads = [d for d in old_loads if id(d) not in static_ids] + notes.append(f"{len(static)} static_gen injection(s) at {len(injection)} bus(es) kept " + "unchanged; ExaGO load at those buses rebuilt net of them.") nonz = sorted({d["ports"]["bus"] for d in old_loads if d["params"].get("alphaI", 0) or d["params"].get("alphaP", 0)}) if nonz: @@ -228,14 +294,19 @@ def apply_steady_state(case: dict, net: MATNetwork) -> tuple[dict, list[str]]: alphas = {d["ports"]["bus"]: (d["params"].get("alphaI", 0.0), d["params"].get("alphaP", 0.0)) for d in old_loads} load_ids = {d["ports"]["bus"]: d["id"] for d in old_loads} - devices[:] = [d for d in devices if d["class"] != "LoadZIP"] + devices[:] = [d for d in devices if d["class"] != "LoadZIP"] + static n_shunts = 0 for n, b in sorted(buses.items()): - if b.Pd or b.Qd: + # ExaGO's Pd/Qd is net of any kept injection, so add it back. + pd = b.Pd - injection[n][0] if n in injection else b.Pd + qd = b.Qd - injection[n][1] if n in injection else b.Qd + if n in injection and n not in load_ids and abs(pd) < 0.05 and abs(qd) < 0.05: + pd = qd = 0.0 # rounding left over from netting; the bus has no load + if pd or qd: ai, ap = alphas.get(n, (0.0, 0.0)) devices.append({"class": "LoadZIP", "id": load_ids.get(n, f"load_{n}"), "ports": {"bus": n}, - "params": {"Pnom": b.Pd / base, "Qnom": b.Qd / base, + "params": {"Pnom": pd / base, "Qnom": qd / base, "alphaI": ai, "alphaP": ap}}) if b.Gs or b.Bs: vm2 = b.Vm * b.Vm # Pnom/Qnom are taken at the initial voltage @@ -282,7 +353,76 @@ def apply_steady_state(case: dict, net: MATNetwork) -> tuple[dict, list[str]]: return out, notes +def exago_shifts_to_degrees(net: MATNetwork) -> list[str]: + """Convert the branch phase shifts of an ExaGO saved solution to degrees. + + ExaGO reads MATPOWER shift angles in degrees and converts them to radians + (psreaddata.cpp), but -save_output writes them back still in radians + (psoutput.cpp), so opflowout.m/pflowout.m break the MATPOWER convention. + Bus angles (Va) are written in degrees and are left as they are. + """ + shifted = [br for br in net.branches if br.angle] + for br in shifted: + br.angle = math.degrees(br.angle) + if not shifted: + return [] + return [f"{len(shifted)} phase shift(s) in ExaGO's solution converted from radians to degrees: " + + ", ".join(f"{br.fbus}-{br.tbus} {br.angle:.3f} deg [{br.angle!r} deg]" for br in shifted[:10])] + + +def _end_flows(br, vf: complex, vt: complex, base: float) -> tuple[complex, complex]: + """MVA entering *br* at fbus with the tap and phase shift on the fbus side, + and at fbus if they were on the tbus side instead.""" + ys = 1.0 / complex(br.r, br.x) + ysh = 0.5j * br.b + tap = (br.ratio or 1.0) * complex(math.cos(math.radians(br.angle)), math.sin(math.radians(br.angle))) + i_tap_here = (ys + ysh) / abs(tap) ** 2 * vf - ys / tap.conjugate() * vt + i_tap_there = (ys + ysh) * vf - ys / tap * vt + return vf * i_tap_here.conjugate() * base, vf * i_tap_there.conjugate() * base + + +def exago_branch_directions(net: MATNetwork) -> list[str]: + """Put back the from/to buses of transformers that ExaGO wrote reversed. + + ExaGO's -save_output (psoutput.cpp) writes some branches with fbus and tbus + swapped but keeps the tap ratio and phase shift, which belong to the + original from bus, so a reader puts them on the wrong winding (ACTIVSg10k: + 3562 of 12706 branches, 476 of them off-nominal transformers). Its reported + flows are right for the real direction, so for each transformer the + direction whose computed flow matches the reported Pf/Qf at fbus is kept. + Run after exago_shifts_to_degrees. Lines (tap 1, no shift) are the same + either way and are left as written. + """ + buses = {b.bus_i: b.Vm * complex(math.cos(math.radians(b.Va)), math.sin(math.radians(b.Va))) + for b in net.buses} + base = net.baseMVA + flipped = [] + for br in net.branches: + if not br.status or ((br.ratio in (0.0, 1.0)) and not br.angle) or len(br.extra) < 2: + continue + reported = complex(br.extra[0], br.extra[1]) + as_written, reversed_ = (abs(s - reported) for s in + _end_flows(br, buses[br.fbus], buses[br.tbus], base)) + if reversed_ < as_written: + br.fbus, br.tbus = br.tbus, br.fbus + br.extra[:4] = br.extra[2:4] + br.extra[:2] + flipped.append(br) + if not flipped: + return [] + return [f"{len(flipped)} transformer(s) written by ExaGO with from/to reversed put back in " + "their real direction (tap on the right side): " + + ", ".join(f"{br.fbus}-{br.tbus} tap {br.ratio:g}" for br in flipped[:10])] + + def build_from_exago(case: dict, source: SteadyStateSource) -> tuple[dict, list[str]]: """Read *source*'s solved case and apply it to the GridKit *case*.""" - net = parse_matpower(source.solved_file) - return apply_steady_state(case, net) + with log_scope( + logger, "ExaGO to GridKit conversion", + application=source.application, + run_dir=source.run_dir, + ): + net = parse_matpower(source.solved_file) + notes = exago_shifts_to_degrees(net) + notes += exago_branch_directions(net) + out, more = apply_steady_state(case, net) + return out, notes + more diff --git a/agentigrid/agentigrid/gridkit_engine/journal.py b/agentigrid/agentigrid/gridkit_engine/journal.py index fa0813b4..45e2478f 100644 --- a/agentigrid/agentigrid/gridkit_engine/journal.py +++ b/agentigrid/agentigrid/gridkit_engine/journal.py @@ -34,7 +34,8 @@ def outcome_row(o: FaultOutcome, run_dir: Path | None = None) -> dict: "solved": o.run.solved, "run": o.run.reason, "checks": { - c.name: {"status": c.status, "value": c.value, "limit": c.limit, "detail": c.detail} + c.name: {"status": c.status, "value": c.value, "limit": c.limit, "detail": c.detail, + "data": c.data} for c in o.checks }, "run_dir": str(run_dir) if run_dir else None, @@ -47,6 +48,7 @@ def __init__(self, goal: str, case_file: Path, session_dir: Path): self.case_file = case_file self.session_dir = session_dir self.entries: list[JournalEntry] = [] + self.token_usage: dict | None = None # TokenLedger.to_dict(), set before export def add(self, entry: JournalEntry) -> JournalEntry: self.entries.append(entry) @@ -72,6 +74,7 @@ def to_dict(self) -> dict: "session_dir": str(self.session_dir), "element_map": str(self.session_dir / "element_map.json"), "entries": [asdict(e) for e in self.entries], + **({"token_usage": self.token_usage} if self.token_usage is not None else {}), } def export_json(self, path: Path | None = None) -> Path: diff --git a/agentigrid/agentigrid/gridkit_engine/validation.py b/agentigrid/agentigrid/gridkit_engine/validation.py index 4b11e625..897c8f09 100644 --- a/agentigrid/agentigrid/gridkit_engine/validation.py +++ b/agentigrid/agentigrid/gridkit_engine/validation.py @@ -100,13 +100,13 @@ def validate_fault_study( if study.clearing_cycles <= 0 or study.margin_cycles < 0: errors.append("clearing_cycles must be positive and margin_cycles not negative") if clear >= study.tmax_s: - errors.append(f"Fault clears at {clear:.3f} s, not before tmax_s = {study.tmax_s} s") + errors.append(f"Fault clears at {criteria.num(clear, '.3f', ' s')}, not before tmax_s = {study.tmax_s} s") if study.dt_monitor_s <= 0: errors.append("dt_monitor_s must be positive") elif study.dt_monitor_s > duration / 2: warnings.append( f"dt_monitor_s = {study.dt_monitor_s} s records fewer than two points during the " - f"{duration:.3f} s fault" + f"{criteria.num(duration, '.3f', ' s')} fault" ) unknown = sorted(set(study.checks) - set(CHECKS)) diff --git a/agentigrid/agentigrid/gridkit_parsers/results_parser.py b/agentigrid/agentigrid/gridkit_parsers/results_parser.py index 461b56f8..63bd4504 100644 --- a/agentigrid/agentigrid/gridkit_parsers/results_parser.py +++ b/agentigrid/agentigrid/gridkit_parsers/results_parser.py @@ -12,6 +12,7 @@ from dataclasses import dataclass, field from pathlib import Path +from agentigrid.logging_setup import log_scope from agentigrid.gridkit_parsers.case_parser import column_owners logger = logging.getLogger("agentigrid.gridkit_parsers.results") @@ -42,25 +43,26 @@ def device_variable(self, name: str, classes: frozenset[str] | None = None) -> d def read_results(csv_path: Path | str, element_map: dict) -> TimeSeries: """Read a GridKit CSV into a :class:`TimeSeries`.""" - with Path(csv_path).open(encoding="utf-8") as fh: - rows = [r for r in csv.reader(fh) if r] - if not rows: - return TimeSeries(t=[]) - header, data = rows[0], rows[1:] - columns = [[float(r[i]) for r in data] for i in range(len(header))] - series = TimeSeries(t=columns[0]) - owners = column_owners(element_map, header[1:]) - for i, name in enumerate(header[1:], start=1): - owner = owners[name] - if owner is None: - series.unmatched.append(name) - elif owner["kind"] == "bus": - series.buses.setdefault(owner["bus"], {})[owner["variable"]] = columns[i] - else: - dev = owner["id"] - series.devices.setdefault(dev, {})[owner["variable"]] = columns[i] - series.device_class[dev] = owner["class"] - series.device_bus[dev] = owner["bus"] - if series.unmatched: - logger.warning("CSV columns not in the element map: %s", series.unmatched[:10]) - return series + with log_scope(logger, "GridKit CSV parsing", csv_path=csv_path): + with Path(csv_path).open(encoding="utf-8") as fh: + rows = [r for r in csv.reader(fh) if r] + if not rows: + return TimeSeries(t=[]) + header, data = rows[0], rows[1:] + columns = [[float(r[i]) for r in data] for i in range(len(header))] + series = TimeSeries(t=columns[0]) + owners = column_owners(element_map, header[1:]) + for i, name in enumerate(header[1:], start=1): + owner = owners[name] + if owner is None: + series.unmatched.append(name) + elif owner["kind"] == "bus": + series.buses.setdefault(owner["bus"], {})[owner["variable"]] = columns[i] + else: + dev = owner["id"] + series.devices.setdefault(dev, {})[owner["variable"]] = columns[i] + series.device_class[dev] = owner["class"] + series.device_bus[dev] = owner["bus"] + if series.unmatched: + logger.warning("CSV columns not in the element map: %s", series.unmatched[:10]) + return series diff --git a/agentigrid/agentigrid/logging_setup.py b/agentigrid/agentigrid/logging_setup.py index ef59b122..09c594d2 100644 --- a/agentigrid/agentigrid/logging_setup.py +++ b/agentigrid/agentigrid/logging_setup.py @@ -3,8 +3,33 @@ from __future__ import annotations import logging +import time +from contextlib import contextmanager from datetime import datetime +from itertools import count from pathlib import Path +from typing import Iterator + + +_scope_ids = count(1) + + +@contextmanager +def log_scope(logger: logging.Logger, name: str, **context: object) -> Iterator[None]: + """Log on scope enter and exit, including time spent.""" + scope_id = next(_scope_ids) + details = " ".join(f"{key}={value}" for key, value in context.items()) + label = f"{name} [scope={scope_id}{' ' + details if details else ''}]" + start = time.monotonic() + logger.debug("Scope start: %s", label) + try: + yield + except BaseException as exc: + logger.debug("Scope raised: %s elapsed=%.3fs exception=%s", + label, time.monotonic() - start, type(exc).__name__) + raise + else: + logger.debug("Scope end: %s elapsed=%.3fs", label, time.monotonic() - start) def setup_logging(logs_dir: Path, verbose: bool = False) -> None: @@ -26,7 +51,7 @@ def setup_logging(logs_dir: Path, verbose: bool = False) -> None: # --- Console handler --- console = logging.StreamHandler() console.setLevel(logging.DEBUG if verbose else logging.INFO) - console_fmt = logging.Formatter("[%(asctime)s] %(levelname)s: %(message)s", datefmt="%H:%M:%S") + console_fmt = logging.Formatter("[%(asctime)s.%(msecs)03d] %(levelname)s: %(message)s", datefmt="%H:%M:%S") console.setFormatter(console_fmt) root_logger.addHandler(console) @@ -37,7 +62,7 @@ def setup_logging(logs_dir: Path, verbose: bool = False) -> None: file_handler = logging.FileHandler(log_file, encoding="utf-8") file_handler.setLevel(logging.DEBUG) file_fmt = logging.Formatter( - "[%(asctime)s] %(levelname)s [%(name)s]: %(message)s", + "[%(asctime)s.%(msecs)03d] %(levelname)s [%(name)s]: %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) file_handler.setFormatter(file_fmt) diff --git a/agentigrid/agentigrid/token_usage.py b/agentigrid/agentigrid/token_usage.py new file mode 100644 index 00000000..3f80462d --- /dev/null +++ b/agentigrid/agentigrid/token_usage.py @@ -0,0 +1,251 @@ +"""Token accounting for every LLM call of a session. + +One :class:`TokenLedger` per session records each ``LLMBackend.complete()`` +response under a phase, so totals cover every call and the breakdown can be +inspected (journal ``token_usage`` block, console summary, report, GUI). +""" + +from __future__ import annotations + +from typing import Any, Optional + +from agentigrid.backends.base import LLMResponse + +# Phase names, in display order. +# objective_extraction objectives read from the initial goal +# search every call inside the iteration loop (incl. explore/select, +# the 'analyze' fallback and retries after parse failures) +# analysis post-search analysis / goal classification +# steering objectives read from a steering directive +# report analysis regenerated for a report from a saved journal +# other anything else +PHASES = ("objective_extraction", "search", "analysis", "steering", "report", "other") + +TOKEN_USAGE_VERSION = 1 + + +def _label(phase: str) -> str: + return phase.replace("_", " ") + + +class TokenLedger: + """Per-call record of LLM token usage, with per-phase and grand totals. + + Usage a backend did not report stays ``None`` in the per-call list; the + aggregates count it as 0 and report it in ``calls_without_usage``. + + ``prompt_tokens`` is the total input of a call. With Anthropic prompt + caching it already includes ``cache_creation_tokens`` (cache writes) and + ``cache_read_tokens`` (cache reads), which break that total down; they are + ``None`` for backends without caching. ``api_error`` marks a failed request + and ``stop_reason`` why generation stopped (``"max_tokens"`` = cut off). + """ + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def record( + self, + phase: str, + response: LLMResponse, + iteration: Optional[int] = None, + prompt_chars: Optional[int] = None, + ) -> None: + """Add one ``complete()`` response under *phase*. + + *prompt_chars* is the length of system + user prompt (``None`` if the + call site cannot provide it); ``completion_chars`` is the length of the + response's ``raw_text``. Fields missing on *response* are recorded as + ``None``, so bookkeeping never fails on a partial response object. + """ + raw_text = getattr(response, "raw_text", None) + if phase not in PHASES: + raise ValueError(f"Unknown token phase '{phase}' (expected one of {PHASES})") + self.calls.append({ + "phase": phase, + "iteration": iteration, + "model": _str_or_none(getattr(response, "model", None)), + "backend": _str_or_none(getattr(response, "backend", None)), + "prompt_tokens": _int_or_none(getattr(response, "prompt_tokens", None)), + "completion_tokens": _int_or_none(getattr(response, "completion_tokens", None)), + "prompt_chars": _int_or_none(prompt_chars), + "completion_chars": len(raw_text) if isinstance(raw_text, str) else None, + "cache_creation_tokens": _int_or_none(getattr(response, "cache_creation_tokens", None)), + "cache_read_tokens": _int_or_none(getattr(response, "cache_read_tokens", None)), + "api_error": getattr(response, "api_error", False) is True, + "stop_reason": _str_or_none(getattr(response, "stop_reason", None)), + }) + + # ------------------------------------------------------------------ + # Totals + # ------------------------------------------------------------------ + + @property + def total_prompt(self) -> int: + return sum(c["prompt_tokens"] or 0 for c in self.calls) + + @property + def total_completion(self) -> int: + return sum(c["completion_tokens"] or 0 for c in self.calls) + + @property + def total(self) -> int: + return self.total_prompt + self.total_completion + + @property + def total_prompt_chars(self) -> int: + return sum(c.get("prompt_chars") or 0 for c in self.calls) + + @property + def total_completion_chars(self) -> int: + return sum(c.get("completion_chars") or 0 for c in self.calls) + + @property + def total_cache_creation(self) -> int: + return sum(c.get("cache_creation_tokens") or 0 for c in self.calls) + + @property + def total_cache_read(self) -> int: + return sum(c.get("cache_read_tokens") or 0 for c in self.calls) + + @property + def api_errors(self) -> int: + return sum(1 for c in self.calls if c.get("api_error")) + + @property + def truncated(self) -> int: + """Calls whose output was cut off at the max_tokens cap.""" + return sum(1 for c in self.calls if c.get("stop_reason") == "max_tokens") + + @property + def calls_without_usage(self) -> int: + return sum(1 for c in self.calls if _missing_usage(c)) + + def __len__(self) -> int: + return len(self.calls) + + @property + def by_phase(self) -> dict[str, dict[str, int]]: + """Aggregates per phase that has at least one call, in :data:`PHASES` order.""" + out: dict[str, dict[str, int]] = {} + for phase in PHASES: + calls = [c for c in self.calls if c["phase"] == phase] + if not calls: + continue + prompt = sum(c["prompt_tokens"] or 0 for c in calls) + completion = sum(c["completion_tokens"] or 0 for c in calls) + out[phase] = { + "calls": len(calls), + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": prompt + completion, + "calls_without_usage": sum(1 for c in calls if _missing_usage(c)), + "prompt_chars": sum(c.get("prompt_chars") or 0 for c in calls), + "completion_chars": sum(c.get("completion_chars") or 0 for c in calls), + } + return out + + # ------------------------------------------------------------------ + # Export / import + # ------------------------------------------------------------------ + + def to_dict(self) -> dict[str, Any]: + """JSON-ready block (journal ``token_usage``, session checkpoint).""" + return { + "version": TOKEN_USAGE_VERSION, + "calls": len(self.calls), + "prompt_tokens": self.total_prompt, + "completion_tokens": self.total_completion, + "total_tokens": self.total, + "calls_without_usage": self.calls_without_usage, + "prompt_chars": self.total_prompt_chars, + "completion_chars": self.total_completion_chars, + "cache_creation_tokens": self.total_cache_creation, + "cache_read_tokens": self.total_cache_read, + "api_errors": self.api_errors, + "truncated": self.truncated, + "by_phase": self.by_phase, + "per_call": [dict(c) for c in self.calls], + } + + @classmethod + def from_dict(cls, data: Optional[dict]) -> "TokenLedger": + """Rebuild from :meth:`to_dict` output; ``None`` or ``{}`` gives an empty ledger.""" + ledger = cls() + for c in (data or {}).get("per_call", []): + ledger.calls.append({ + "phase": c.get("phase", "other"), + "iteration": c.get("iteration"), + "model": c.get("model"), + "backend": c.get("backend"), + "prompt_tokens": c.get("prompt_tokens"), + "completion_tokens": c.get("completion_tokens"), + "prompt_chars": c.get("prompt_chars"), # absent in older journals + "completion_chars": c.get("completion_chars"), + # absent before the cache / error fields were added + "cache_creation_tokens": c.get("cache_creation_tokens"), + "cache_read_tokens": c.get("cache_read_tokens"), + "api_error": bool(c.get("api_error", False)), + "stop_reason": c.get("stop_reason"), + **({"carried_over": True} if c.get("carried_over") else {}), + }) + return ledger + + @classmethod + def from_legacy_totals(cls, prompt_tokens: int, completion_tokens: int) -> "TokenLedger": + """Ledger for a session saved before per-call accounting existed. + + Such sessions counted search calls only, so the saved totals become one + ``search`` entry marked ``carried_over``. + """ + ledger = cls() + if prompt_tokens or completion_tokens: + ledger.calls.append({ + "phase": "search", "iteration": None, "model": None, "backend": None, + "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, + "prompt_chars": None, "completion_chars": None, + "cache_creation_tokens": None, "cache_read_tokens": None, + "api_error": False, "stop_reason": None, + "carried_over": True, + }) + return ledger + + # ------------------------------------------------------------------ + # Text + # ------------------------------------------------------------------ + + def breakdown_text(self) -> str: + """``search 19,497, objective extraction 1,210, analysis 5,021`` (phases with calls only).""" + return ", ".join(f"{_label(p)} {a['total_tokens']:,}" for p, a in self.by_phase.items()) + + def summary_text(self) -> str: + """One line for the console summary.""" + n = len(self.calls) + text = ( + f"Tokens: {self.total:,} total ({self.total_prompt:,} prompt + " + f"{self.total_completion:,} completion) in {n} LLM call{'s' if n != 1 else ''}" + ) + if self.calls: + text += f"; {self.breakdown_text()}" + cached = self.total_cache_read + if cached: + text += f"; {cached:,} prompt tokens read from cache" + missing = self.calls_without_usage + if missing: + text += f"; {missing} call{'s' if missing != 1 else ''} without usage data" + return text + + +def _str_or_none(value: Any) -> Optional[str]: + return value if isinstance(value, str) else None + + +def _int_or_none(value: Any) -> Optional[int]: + """Token count as int; anything that is not a number (None, a mock) is None.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return int(value) + + +def _missing_usage(call: dict) -> bool: + return call["prompt_tokens"] is None or call["completion_tokens"] is None diff --git a/agentigrid/claude_code_changes.md b/agentigrid/claude_code_changes.md index 83a815a8..252f7819 100644 --- a/agentigrid/claude_code_changes.md +++ b/agentigrid/claude_code_changes.md @@ -4,6 +4,41 @@ This document records significant changes made by Claude Code, grouped by the pr --- +## Agent loop and LLM backends: robust iterations, honest UI, prompt caching (2026-10-03; on token accounting 2026-11-05) + +**Motivation.** Long runs with local models (Ollama) showed iterations that went wrong in ways the UI hid or misreported: a malformed action crashed the run; a rejected iteration re-emitted the previous card, so the timeline showed stale or out-of-order cards (e.g. "Iteration 6" after iteration 20); analysis and completion entries were labelled FAILED; a modify that applied no command showed "No description"; and the summary counted the base case as an iteration ("Iterations 3" for max 2). Separately, Anthropic runs re-sent the same system prompt at full price on every iteration, and some newer models reject the `temperature` parameter. + +### Agent loop (`agentigrid/exago_engine/agent_loop.py`) + +- Each iteration runs inside a guard: an exception in an action handler is discarded with feedback to the model, and the run continues. Handlers that read `feasibility` / `candidate_set` tolerate non-dict values. +- The UI callback fires only when the iteration recorded a journal entry. Otherwise a UI-only "discarded" card explains why ("LLM output rejected — … (no change applied)"); it is never added to the journal. The final re-emit of the last entry is gone (it duplicated the last card). Same in `resume`. +- `journal.discarded_actions` records every iteration that produced no entry: kind (`rejected`, `api_error`, `truncated`, `internal`), the feedback sent to the model, `stop_reason`, content block types and the raw response (capped at 4,000 characters). Entries and summary statistics are unchanged. +- A response cut at `max_tokens` gets its own message and feedback asking for the JSON action only. +- A modify that applies zero commands says so in its description; modify entries record `skipped_commands`. +- Optional stall breaker: `AGENTIGRID_STALL_LIMIT=N` ends the run after N consecutive iterations without a new entry (`termination_reason = "stalled_no_progress"`). Off by default. + +### Backends (`agentigrid/backends/`) + +- Anthropic prompt caching of the per-session system prompt (`llm.prompt_cache`, default on). Billing only; the model sees the same text. `prompt_tokens` stays the total input; cache writes and reads are reported separately. +- Token counts go through `TokenLedger` (samimk/token-accounting); there is no second counter. Each per-call record also keeps `cache_creation_tokens`, `cache_read_tokens`, `api_error` and `stop_reason`, and `token_usage` adds the totals `cache_creation_tokens`, `cache_read_tokens`, `api_errors` and `truncated`. Older records load with these empty. The console summary adds the tokens read from cache. +- The journal records the run's LLM settings as `llm_settings`: backend, model, `max_tokens`, `prompt_cache`, `temperature_sent`. +- Models that reject `temperature` are retried once without it, and it is not sent again in that session. +- `LLMResponse` gains `api_error`, `stop_reason`, `content_types` and cache token fields; failed requests in all backends set `api_error`. +- `AGENTIGRID_MAX_TOKENS` overrides the output-token cap. + +### UI and reports + +- Iterations shown as "N + base case" (UI summary and live panel, CLI summary, PDF report). `summary_stats()` adds `llm_iterations` and `has_base_case`; `total_iterations` is unchanged. +- Timeline cards show FAILED only for solves that did not converge, not for analysis / complete / sweep / contingency / explore entries or discarded proposals. +- Sidebar checkbox for the stall breaker (unchecked by default). +- Default Anthropic model `claude-sonnet-4-20250514` (retired) → `claude-sonnet-4-6`. + +### Tests + +`tests/llm/test_prompt_cache.py` (new: caching, token breakdown and its `TokenLedger` records, `llm_settings`, temperature retry, API errors, max_tokens, env override), `tests/journal/test_iteration_count.py` (new), `tests/agent_loop/test_agent_loop.py` (+8: discarded telemetry, API errors, truncation, response cap, no stale UI cards, stall breaker off/on, internal errors). The 4 new agent-loop callback/stall/error tests fail on the previous agent loop. + +--- + ## C.2/C.3 corrections — metric gating, summary aggregation, predicate audit, sweep dedup (2026-06-23) Four contained fixes from the prompt-13/15/17 runs. No capability behavior changed (capacity numbers, dispatch results, metric/predicate definitions, and the C.1 binding identifier are untouched) — only how results are gated, summarized, audited, and deduplicated. diff --git a/agentigrid/configs/default_config.yaml.template b/agentigrid/configs/default_config.yaml.template index dd5dd092..23908c84 100644 --- a/agentigrid/configs/default_config.yaml.template +++ b/agentigrid/configs/default_config.yaml.template @@ -30,7 +30,7 @@ data: # LLM backend settings llm: backend: "anthropic" # openai | anthropic | ollama | ollama-cloud - model: "claude-sonnet-4-20250514" + model: "claude-sonnet-4-6" api_key_env: "ANTHROPIC_API_KEY" # env var name containing the API key # OpenAI-specific openai_base_url: null @@ -40,6 +40,7 @@ llm: # Generation parameters temperature: 0.3 max_tokens: 4096 + prompt_cache: true # Anthropic only: cache the system prompt (billing-only; outputs unchanged) # Search settings search: diff --git a/agentigrid/configs/local_config.yaml.template b/agentigrid/configs/local_config.yaml.template index 3f4d5764..9e5cf608 100644 --- a/agentigrid/configs/local_config.yaml.template +++ b/agentigrid/configs/local_config.yaml.template @@ -29,6 +29,7 @@ llm: # ollama_cloud_host: null # set for Ollama Cloud temperature: 0.3 max_tokens: 4096 + prompt_cache: true # Anthropic only: cache the system prompt (billing-only; outputs unchanged) # Search settings search: diff --git a/agentigrid/launcher/ARCHITECTURE.md b/agentigrid/launcher/ARCHITECTURE.md index 3a8a2d98..23263b9b 100644 --- a/agentigrid/launcher/ARCHITECTURE.md +++ b/agentigrid/launcher/ARCHITECTURE.md @@ -224,7 +224,7 @@ Always visible. Contains all search parameters. **Section: LLM Backend** - Backend selector: `anthropic` | `openai` | `ollama` | `ollama-cloud` -- Model name text input (pre-filled with defaults per backend: `claude-sonnet-4-20250514` for Anthropic, `gpt-4o` for OpenAI, etc.) +- Model name text input (pre-filled with defaults per backend: `claude-sonnet-4-6` for Anthropic, `gpt-4o` for OpenAI, etc.) - Temperature slider (0.0 – 1.0, default 0.3) **Section: Search Parameters** @@ -274,7 +274,7 @@ Shown while a search is running. Two-column layout: - Key metrics: voltage range, max line loading, gen/load totals - Mode used (fresh/accumulative) - Current iteration shows a live status indicator: - - Spinner with phase text: "Sending prompt to Claude (claude-sonnet-4-20250514)..." → "Applying 3 modifications..." → "Running OPFLOW simulation..." → "Parsing results..." + - Spinner with phase text: "Sending prompt to Claude (claude-sonnet-4-6)..." → "Applying 3 modifications..." → "Running OPFLOW simulation..." → "Parsing results..." **Right column (~35%)** — Live Charts: - **Convergence chart** (Plotly line chart): Objective value vs. iteration number diff --git a/agentigrid/launcher/app.py b/agentigrid/launcher/app.py index 13613f1f..b592e149 100644 --- a/agentigrid/launcher/app.py +++ b/agentigrid/launcher/app.py @@ -36,7 +36,11 @@ reserve_trajectory_chart, ) -from agentigrid.exago_engine.journal import describe_contingency_scope +from agentigrid.exago_engine.journal import ( + NON_SOLVE_STATUSES, + describe_contingency_scope, + format_iteration_count, +) from agentigrid.exago_parsers import parse_matpower, network_summary # ── Page Configuration ─────────────────────────────────────────────────────── @@ -395,6 +399,23 @@ def render_sidebar() -> dict: max_iterations = st.slider( "Max iterations", 1, 50, 20, disabled=disabled, ) + stall_stop = st.checkbox( + "Stop early if the model stalls", + value=False, + disabled=disabled, + help=( + "End the search after N consecutive no-progress iterations " + "(the model repeating an invalid action or applying no change). " + "Saves iterations when a model is stuck. Leave unchecked for " + "comparable runs that always reach Max iterations." + ), + ) + stall_limit = st.number_input( + "Stop after N no-progress iterations", + min_value=2, max_value=20, value=5, step=1, + disabled=disabled or not stall_stop, + ) + os.environ["AGENTIGRID_STALL_LIMIT"] = str(int(stall_limit)) if stall_stop else "0" mpi_disabled = disabled or application not in ("scopflow", "sopflow") mpi_np = st.number_input( "MPI processes", @@ -798,15 +819,23 @@ def render_live_monitor(): icon = _iteration_icon(entry) obj = entry.get("objective_value") is_pflow_live = st.session_state.get("current_application") == "pflow" - if is_pflow_live: - obj_str = "analysis" if obj is not None else "FAILED" + # An entry with no objective is only "FAILED" if it was a real solve + # that did not converge. Analysis/complete/sweep/contingency/explore + # markers and discarded proposals carry no cost — they are not failures. + cs = entry.get("convergence_status") or "" + _no_solve = cs in NON_SOLVE_STATUSES or entry.get("mode") == "discarded" + if obj is not None: + obj_str = "analysis" if is_pflow_live else f"${obj:,.2f}" + elif _no_solve: + obj_str = "" # no cost, but not a failure — omit the segment else: - obj_str = f"${obj:,.2f}" if obj is not None else "FAILED" + obj_str = "FAILED" # a solve that genuinely did not converge elapsed = entry.get("sim_elapsed") time_str = f"{elapsed:.1f}s" if elapsed is not None else "—" + _cost_seg = f" — {obj_str}" if obj_str else "" label = ( f"{icon} Iteration {entry['iteration']}: " - f"{entry['description']} — {obj_str} ({time_str})" + f"{entry['description']}{_cost_seg} ({time_str})" ) with st.expander(label): @@ -934,7 +963,9 @@ def render_live_monitor(): # Progress Stats st.markdown("---") - n_iters = len(st.session_state.iteration_log) + _its = {e.get("iteration") for e in st.session_state.iteration_log} + _n_agent = sum(1 for i in _its if isinstance(i, int) and i >= 1) + n_iters = f"{_n_agent} + base case" if 0 in _its else str(_n_agent) feasible_count = sum( 1 for e in st.session_state.iteration_log if e.get("feasible") ) @@ -1326,10 +1357,12 @@ def _render_overview_tab(session): mc1, mc2, mc3, mc4, mc5 = st.columns(5) mc1.metric("Application", f"{session.application}") - mc2.metric("Iterations", f"{stats['total_iterations']}") + mc2.metric("Iterations", format_iteration_count(stats)) mc3.metric("Duration", duration_str) mc4.metric("Termination", session.termination_reason) mc5.metric("Tokens", f"{total_tokens:,}" if total_tokens > 0 else "—") + if len(session.token_ledger): + st.caption(session.token_ledger.summary_text()) base_entry = session.journal.entries[0] if session.journal.entries else None best_entry = None diff --git a/agentigrid/launcher/config_builder.py b/agentigrid/launcher/config_builder.py index 4e13f207..a410bb61 100644 --- a/agentigrid/launcher/config_builder.py +++ b/agentigrid/launcher/config_builder.py @@ -253,7 +253,7 @@ def load_example_goals() -> list[dict[str, str]]: # Default model names per backend DEFAULT_MODELS: dict[str, str] = { - "anthropic": "claude-sonnet-4-20250514", + "anthropic": "claude-sonnet-4-6", "openai": "gpt-4o", "ollama": "qwen2.5:7b", "ollama-cloud": "qwen2.5:7b", diff --git a/agentigrid/launcher/report_generator.py b/agentigrid/launcher/report_generator.py index 8719b60f..6c99da8e 100644 --- a/agentigrid/launcher/report_generator.py +++ b/agentigrid/launcher/report_generator.py @@ -22,8 +22,13 @@ from reportlab.pdfbase import pdfmetrics from reportlab.pdfbase.ttfonts import TTFont +from agentigrid.logging_setup import log_scope from agentigrid.exago_engine.agent_loop import SearchSession -from agentigrid.exago_engine.journal import describe_contingency_scope, is_solve_iteration +from agentigrid.exago_engine.journal import ( + describe_contingency_scope, + format_iteration_count, + is_solve_iteration, +) from agentigrid.exago_parsers.opflow_results import OPFLOWResult try: @@ -39,7 +44,20 @@ reserve_trajectory_chart, ) -logger = logging.getLogger("launcher.report_generator") + +def _token_usage_line(session: SearchSession) -> str: + """'Token usage: in LLM calls ()' for the summary.""" + total = session.total_prompt_tokens + session.total_completion_tokens + if total <= 0: + return "Token usage: N/A" + ledger = getattr(session, "token_ledger", None) + if ledger is not None and len(ledger) and ledger.total == total: + n = len(ledger) + return f"Token usage: {total:,} in {n} LLM call{'s' if n != 1 else ''} ({ledger.breakdown_text()})" + return f"Token usage: {total:,}" + + +logger = logging.getLogger("agentigrid.launcher.report_generator") # ── Font Registration ──────────────────────────────────────────────────────── @@ -82,6 +100,7 @@ def _register_fonts(): # ── Chart Export Helper ────────────────────────────────────────────────────── +@log_scope(logger, "Report chart export") def _export_chart_image(fig, width_px: int = 800, height_px: int = 400) -> bytes | None: """Export a Plotly figure to PNG bytes. @@ -215,6 +234,7 @@ def _build_styles(self) -> dict[str, ParagraphStyle]: ), } + @log_scope(logger, "PDF report generation") def generate( self, session: SearchSession, @@ -555,7 +575,6 @@ def _build_executive_summary( start = datetime.fromisoformat(session.start_time) end = datetime.fromisoformat(session.end_time) if session.end_time else datetime.now() duration = end - start - total_tokens = session.total_prompt_tokens + session.total_completion_tokens # Reserve shortcut — checked before the contingency/sweep shortcuts if contingency_entry is not None and getattr(contingency_entry, "reserve_meta", None): @@ -599,10 +618,10 @@ def _build_executive_summary( "is not deliverable for these losses (required reserve is mis-located)." ) res_lines.extend([ - f"Total iterations: {stats['total_iterations']}", + f"Iterations: {format_iteration_count(stats)}", f"Duration: {duration.total_seconds():.0f}s", f"Termination: {session.termination_reason}", - (f"Token usage: {total_tokens:,}" if total_tokens > 0 else "Token usage: N/A"), + _token_usage_line(session), ]) for line in res_lines: elements.append(Paragraph(self._escape_xml(line), s["body"])) @@ -645,10 +664,10 @@ def _build_executive_summary( relief_line += f"; {unresolved_count} unresolved." cont_lines.append(relief_line) cont_lines.extend([ - f"Total iterations: {stats['total_iterations']}", + f"Iterations: {format_iteration_count(stats)}", f"Duration: {duration.total_seconds():.0f}s", f"Termination: {session.termination_reason}", - (f"Token usage: {total_tokens:,}" if total_tokens > 0 else "Token usage: N/A"), + _token_usage_line(session), ]) for line in cont_lines: elements.append(Paragraph(self._escape_xml(line), s["body"])) @@ -683,10 +702,10 @@ def _build_executive_summary( _headline = f"Sweep complete: {n_feas} of {n_cand} candidate buses are feasible." sweep_lines = [ _headline, - f"Total iterations: {stats['total_iterations']}", + f"Iterations: {format_iteration_count(stats)}", f"Duration: {duration.total_seconds():.0f}s", f"Termination: {session.termination_reason}", - f"Token usage: {total_tokens:,}" if total_tokens > 0 else "Token usage: N/A", + _token_usage_line(session), ] for line in sweep_lines: elements.append(Paragraph(self._escape_xml(line), s["body"])) @@ -707,7 +726,7 @@ def _build_executive_summary( 1 for e in solve_entries if e.feasibility_detail == "marginal" ) lines = [ - f"Total iterations: {stats['total_iterations']}", + f"Iterations: {format_iteration_count(stats)}", f"Feasible solutions: {solve_feasible}", f"Infeasible: {solve_infeasible}", ] @@ -726,7 +745,7 @@ def _build_executive_summary( lines.extend([ f"Duration: {duration.total_seconds():.0f}s", f"Termination: {session.termination_reason}", - f"Token usage: {total_tokens:,}" if total_tokens > 0 else "Token usage: N/A", + _token_usage_line(session), ]) if stats["best_objective"] is not None: @@ -1520,8 +1539,9 @@ def _build_convergence_section( ) # Auto-generated text - n = stats["total_iterations"] - text = f"The search ran for {n} iterations." + n = stats.get("llm_iterations", stats["total_iterations"]) + text = (f"The search ran for {n} iteration(s) after the base case." + if stats.get("has_base_case") else f"The search ran for {n} iteration(s).") if stats["best_objective"] is not None: base_entry = session.journal.entries[0] if session.journal.entries else None if base_entry and base_entry.objective_value and base_entry.objective_value != 0: diff --git a/agentigrid/launcher/session_manager.py b/agentigrid/launcher/session_manager.py index bb691eaf..c07e0d1b 100644 --- a/agentigrid/launcher/session_manager.py +++ b/agentigrid/launcher/session_manager.py @@ -239,6 +239,12 @@ def get_summary_analysis(self, session: SearchSession) -> str: ) response = backend.complete(system_prompt, user_prompt) + # Same ledger as the session the analysis belongs to. + session.token_ledger.record("analysis", response, + prompt_chars=len(system_prompt) + len(user_prompt)) + session.total_prompt_tokens = session.token_ledger.total_prompt + session.total_completion_tokens = session.token_ledger.total_completion + session.journal.token_usage = session.token_ledger.to_dict() raw_text = response.raw_text valid_iters = {e.iteration for e in session.journal.entries} @@ -398,6 +404,7 @@ def save_current_session(self, save_dir: Path | None = None) -> Path | None: current_network=None, total_prompt_tokens=self._session.total_prompt_tokens, total_completion_tokens=self._session.total_completion_tokens, + token_usage=self._session.token_ledger.to_dict(), last_iteration=len(self._session.journal) - 1, termination_reason=self._session.termination_reason, ) diff --git a/agentigrid/tests/agent_loop/test_agent_loop.py b/agentigrid/tests/agent_loop/test_agent_loop.py index c7df4a30..b6f80a57 100644 --- a/agentigrid/tests/agent_loop/test_agent_loop.py +++ b/agentigrid/tests/agent_loop/test_agent_loop.py @@ -254,6 +254,152 @@ def test_invalid_json_then_valid(self, tmp_path: Path): assert len(session.journal) == 3 # base + modify + complete assert session.termination_reason == "completed" + def test_discarded_iterations_are_journaled_as_telemetry(self, tmp_path: Path): + """An unparseable action and an unknown action record no entry, but are + kept in journal.discarded_actions (validator-rejection telemetry).""" + cfg = _make_config(tmp_path, max_iterations=10) + responses = [ + _objectives_response(), # objective parser + _make_llm_response(None, raw_text="This is not JSON"), # iter 1 + _make_llm_response({"action": "dance", "reasoning": "x"}), # iter 2 + _make_llm_response({"action": "complete", "reasoning": "Done", + "findings": {"summary": "ok"}}), # iter 3 + ] + mock_backend = MockBackend(responses) + sim_result = _make_sim_result(stdout=_sample_stdout(), success=True) + with patch("agentigrid.exago_engine.agent_loop.create_backend", return_value=mock_backend), \ + patch("agentigrid.exago_engine.agent_loop.SimulationExecutor") as mock_exec_cls: + mock_executor = MagicMock() + mock_executor.run.return_value = sim_result + mock_exec_cls.return_value = mock_executor + controller = AgentLoopController(cfg) + session = controller.run(BASE_CASE, "Test discard telemetry") + + assert len(session.journal) == 2 # base + complete + disc = session.journal.discarded_actions + assert [(d["iteration"], d["kind"]) for d in disc] == [(1, "rejected"), (2, "rejected")] + assert "parse JSON" in disc[0]["feedback"] and "Unknown action" in disc[1]["feedback"] + # the rejected proposal itself is kept (raw model output) + assert disc[0]["response"] == "This is not JSON" and disc[0]["response_truncated"] is False + assert '"dance"' in disc[1]["response"] or "dance" in disc[1]["response"] + + def test_api_errors_are_not_counted_as_rejections(self, tmp_path: Path): + cfg = _make_config(tmp_path, max_iterations=10) + api_err = LLMResponse(raw_text="Anthropic API error: 400", json_data=None, json_error="x", + model="m", backend="test", prompt_tokens=None, completion_tokens=None, + api_error=True) + responses = [_objectives_response(), api_err, + _make_llm_response({"action": "complete", "reasoning": "Done", + "findings": {"summary": "ok"}})] + mock_backend = MockBackend(responses) + sim_result = _make_sim_result(stdout=_sample_stdout(), success=True) + with patch("agentigrid.exago_engine.agent_loop.create_backend", return_value=mock_backend), \ + patch("agentigrid.exago_engine.agent_loop.SimulationExecutor") as mock_exec_cls: + mock_executor = MagicMock() + mock_executor.run.return_value = sim_result + mock_exec_cls.return_value = mock_executor + session = AgentLoopController(cfg).run(BASE_CASE, "Test api error") + disc = session.journal.discarded_actions + assert [(d["iteration"], d["kind"]) for d in disc] == [(1, "api_error")] + assert "response" not in disc[0] # no model response for a failed request + + def test_truncated_output_is_journaled_as_truncated(self, tmp_path: Path): + cfg = _make_config(tmp_path, max_iterations=10) + cut = LLMResponse(raw_text="", json_data=None, json_error="no JSON", model="m", + backend="test", prompt_tokens=10, completion_tokens=4096, + stop_reason="max_tokens", content_types=["thinking"]) + responses = [_objectives_response(), cut, + _make_llm_response({"action": "complete", "reasoning": "Done", + "findings": {"summary": "ok"}})] + mock_backend = MockBackend(responses) + sim_result = _make_sim_result(stdout=_sample_stdout(), success=True) + with patch("agentigrid.exago_engine.agent_loop.create_backend", return_value=mock_backend), \ + patch("agentigrid.exago_engine.agent_loop.SimulationExecutor") as mock_exec_cls: + mock_executor = MagicMock() + mock_executor.run.return_value = sim_result + mock_exec_cls.return_value = mock_executor + session = AgentLoopController(cfg).run(BASE_CASE, "Test truncation") + d = session.journal.discarded_actions[0] + assert (d["iteration"], d["kind"], d["stop_reason"]) == (1, "truncated", "max_tokens") + assert d["content_types"] == ["thinking"] and d["response"] == "" + assert "output-token limit" in d["feedback"] + + def test_long_discarded_response_is_capped(self, tmp_path: Path): + from agentigrid.exago_engine import agent_loop as al + cfg = _make_config(tmp_path, max_iterations=10) + long_text = "x" * (al.DISCARDED_RESPONSE_MAX + 50) + responses = [_objectives_response(), _make_llm_response(None, raw_text=long_text), + _make_llm_response({"action": "complete", "reasoning": "Done", + "findings": {"summary": "ok"}})] + mock_backend = MockBackend(responses) + sim_result = _make_sim_result(stdout=_sample_stdout(), success=True) + with patch("agentigrid.exago_engine.agent_loop.create_backend", return_value=mock_backend), \ + patch("agentigrid.exago_engine.agent_loop.SimulationExecutor") as mock_exec_cls: + mock_executor = MagicMock() + mock_executor.run.return_value = sim_result + mock_exec_cls.return_value = mock_executor + session = AgentLoopController(cfg).run(BASE_CASE, "Test cap") + d = session.journal.discarded_actions[0] + assert len(d["response"]) == al.DISCARDED_RESPONSE_MAX and d["response_truncated"] is True + + def _run(self, tmp_path, responses, goal="Test", max_iterations=10, **ctl_kw): + cfg = _make_config(tmp_path, max_iterations=max_iterations) + mock_backend = MockBackend([_objectives_response()] + responses) + sim_result = _make_sim_result(stdout=_sample_stdout(), success=True) + with patch("agentigrid.exago_engine.agent_loop.create_backend", return_value=mock_backend), \ + patch("agentigrid.exago_engine.agent_loop.SimulationExecutor") as mock_exec_cls: + mock_executor = MagicMock() + mock_executor.run.return_value = sim_result + mock_exec_cls.return_value = mock_executor + return AgentLoopController(cfg, **ctl_kw).run(BASE_CASE, goal) + + _COMPLETE = {"action": "complete", "reasoning": "Done", "findings": {"summary": "ok"}} + + def test_ui_callback_never_replays_a_stale_entry(self, tmp_path: Path): + """A rejected iteration produces a 'discarded' card for that iteration, + never a re-emit of the previous entry, and the run ends without a + final duplicate of the last card.""" + calls = [] + session = self._run( + tmp_path, + [_make_llm_response(None, raw_text="not JSON"), # iter 1: rejected + _make_llm_response({"action": "dance", "reasoning": "x"}), # iter 2: rejected + _make_llm_response(self._COMPLETE)], # iter 3 + on_iteration=lambda it, entry, kind, opf: calls.append((it, kind, entry.mode)), + ) + assert [c[0] for c in calls] == [0, 1, 2, 3] + assert [c[1] for c in calls[1:3]] == ["discarded", "discarded"] + assert all(c[2] == "discarded" for c in calls[1:3]) + assert session.termination_reason == "completed" + # the synthetic cards are UI-only: the journal holds base case + complete + assert len(session.journal) == 2 + + def test_stall_breaker_is_off_by_default(self, tmp_path: Path, monkeypatch): + monkeypatch.delenv("AGENTIGRID_STALL_LIMIT", raising=False) + bad = {"action": "dance", "reasoning": "x"} + session = self._run(tmp_path, [_make_llm_response(bad)] * 3 + [_make_llm_response(self._COMPLETE)]) + assert session.termination_reason == "completed" + assert len(session.journal.discarded_actions) == 3 + + def test_stall_breaker_stops_after_n_no_progress_iterations(self, tmp_path: Path, monkeypatch): + monkeypatch.setenv("AGENTIGRID_STALL_LIMIT", "2") + bad = {"action": "dance", "reasoning": "x"} + session = self._run(tmp_path, [_make_llm_response(bad)] * 3 + [_make_llm_response(self._COMPLETE)]) + assert session.termination_reason == "stalled_no_progress" + assert [d["iteration"] for d in session.journal.discarded_actions] == [1, 2] + + def test_internal_error_in_an_action_is_discarded_and_the_run_continues(self, tmp_path: Path): + with patch.object(AgentLoopController, "_handle_analyze", side_effect=RuntimeError("boom")): + session = self._run( + tmp_path, + [_make_llm_response({"action": "analyze", "query": "max loading", "reasoning": "x"}), + _make_llm_response(self._COMPLETE)], + ) + assert session.termination_reason == "completed" + d = session.journal.discarded_actions[0] + assert (d["iteration"], d["kind"]) == (1, "internal") + assert "RuntimeError: boom" in d["feedback"] + def test_unknown_action(self, tmp_path: Path): """LLM returns unknown action, then completes.""" cfg = _make_config(tmp_path, max_iterations=10) diff --git a/agentigrid/tests/agent_loop/test_token_accounting.py b/agentigrid/tests/agent_loop/test_token_accounting.py new file mode 100644 index 00000000..f63f2e78 --- /dev/null +++ b/agentigrid/tests/agent_loop/test_token_accounting.py @@ -0,0 +1,278 @@ +"""Every LLM call of an ExaGO session is counted, with its phase. + +Covers: +- a session with objective extraction, a steering directive, search iterations + (incl. the 'analyze' LLM fallback) and the post-search analysis: ledger calls + == backend calls, totals == sum of all calls, phases with the right counts +- the journal export carries the ``token_usage`` block; an old journal without + it still loads +- session checkpoints save ``token_usage``; old checkpoints give totals only +- report regeneration records its call under ``report`` on top of the + journal's own usage +""" + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from agentigrid.backends.base import LLMResponse +from agentigrid.exago_engine.agent_loop import AgentLoopController +from agentigrid.exago_engine.journal import SearchJournal +from agentigrid.exago_engine.regenerate_report import ( + build_session_stub, + load_journal_export, + regenerate_from_journal, +) +from agentigrid.exago_engine.session_io import load_session, save_session +from agentigrid.token_usage import TokenLedger +from tests.agent_loop.test_agent_loop import ( + BASE_CASE, + MockBackend, + _has_base_case, + _has_sample, + _make_config, + _make_llm_response, + _make_sim_result, + _objectives_response, + _sample_stdout, +) + + +def _text_response(text: str, prompt: int, completion: int) -> LLMResponse: + return LLMResponse(raw_text=text, json_data=None, json_error=None, model="test-model", + backend="test", prompt_tokens=prompt, completion_tokens=completion) + + +def _classification_response() -> LLMResponse: + return _text_response( + "Analysis.\n```json\n{\"goal_type\": \"cost_minimization\", \"best_iteration\": 0, " + "\"best_iteration_rationale\": \"base\"}\n```", 700, 300, + ) + + +def _run_session(tmp_path: Path) -> tuple: + cfg = _make_config(tmp_path, max_iterations=10) + responses = [ + _objectives_response(), # objective extraction (50+20) + _objectives_response(), # steering directive (50+20) + _make_llm_response({ # iter 1 (100+50) + "action": "analyze", "reasoning": "curious", + "query": "tell me a short story about the sky", # no pattern: LLM fallback + }), + _text_response("Once upon a time.", 30, 10), # analyze fallback (30+10) + _make_llm_response({ # iter 2 (100+50) + "action": "modify", "reasoning": "+10%", "mode": "fresh", + "description": "Scale all loads +10%", + "commands": [{"action": "scale_all_loads", "factor": 1.1}], + }), + _make_llm_response({ # iter 3 (100+50) + "action": "complete", "reasoning": "done", + "findings": {"summary": "Done"}, + }), + _classification_response(), # post-search analysis (700+300) + ] + backend = MockBackend(responses) + with patch("agentigrid.exago_engine.agent_loop.create_backend", return_value=backend), \ + patch("agentigrid.exago_engine.agent_loop.SimulationExecutor") as mock_exec_cls: + mock_executor = MagicMock() + mock_executor.run.return_value = _make_sim_result(stdout=_sample_stdout()) + mock_exec_cls.return_value = mock_executor + controller = AgentLoopController(cfg) + controller.inject_steering("Also keep generation cost low") + session = controller.run(BASE_CASE, "Find max feasible load increase") + return cfg, backend, controller, session + + +@pytest.mark.skipif(not _has_base_case, reason="Base case .m file not found") +@pytest.mark.skipif(not _has_sample, reason="sample_opflow_output.txt not found") +class TestExagoSessionTokens: + + def test_every_call_counted(self, tmp_path: Path): + _, backend, controller, session = _run_session(tmp_path) + ledger = session.token_ledger + + assert ledger is controller.token_ledger + assert len(backend.calls) == 7 + assert len(ledger) == len(backend.calls) + by_phase = ledger.by_phase + assert {p: a["calls"] for p, a in by_phase.items()} == { + "objective_extraction": 1, "search": 4, "analysis": 1, "steering": 1, + } + assert by_phase["search"]["total_tokens"] == 3 * 150 + 40 + assert by_phase["analysis"]["total_tokens"] == 1000 + assert ledger.total_prompt == 50 + 50 + 3 * 100 + 30 + 700 + assert ledger.total_completion == 20 + 20 + 3 * 50 + 10 + 300 + assert ledger.total == sum( + (c["prompt_tokens"] or 0) + (c["completion_tokens"] or 0) for c in ledger.calls + ) + # Session totals now cover every phase + assert session.total_prompt_tokens == ledger.total_prompt + assert session.total_completion_tokens == ledger.total_completion + # The analyze fallback is recorded with the iteration it ran in + search_iters = [c["iteration"] for c in ledger.calls if c["phase"] == "search"] + assert search_iters == [1, 1, 2, 3] + + def test_char_lengths_match_backend_calls(self, tmp_path: Path): + _, backend, _, session = _run_session(tmp_path) + ledger = session.token_ledger + assert [c["prompt_chars"] for c in ledger.calls] == [ + len(system) + len(user) for system, user in backend.calls + ] + assert all(c["completion_chars"] is not None for c in ledger.calls) + assert ledger.calls[3]["completion_chars"] == len("Once upon a time.") # analyze fallback + assert ledger.total_prompt_chars == sum(c["prompt_chars"] for c in ledger.calls) + assert ledger.total_completion_chars == sum(c["completion_chars"] for c in ledger.calls) + assert ledger.to_dict()["prompt_chars"] == ledger.total_prompt_chars + + def test_journal_has_token_usage(self, tmp_path: Path): + cfg, _, controller, _ = _run_session(tmp_path) + exports = sorted(cfg.output.workdir.glob("journal_*.json")) + assert exports, "journal was not exported" + data = json.loads(exports[-1].read_text()) + assert data["token_usage"] == controller.token_ledger.to_dict() + assert data["token_usage"]["calls"] == 7 + # and it loads back onto the journal + assert load_journal_export(exports[-1]).token_usage == data["token_usage"] + + def test_classification_prompt_gets_tokens_before_analysis(self, tmp_path: Path): + _, backend, _, _ = _run_session(tmp_path) + before = 50 + 20 + 50 + 20 + 3 * 150 + 40 + _, classification_user_prompt = backend.calls[-1] + assert f"Tokens used before this analysis: ~{before:,}" in classification_user_prompt + + def test_summary_line_printed(self, tmp_path: Path, capsys): + _run_session(tmp_path) + out = capsys.readouterr().out + assert ("Tokens: 1,630 total (1,130 prompt + 500 completion) in 7 LLM calls; " + "objective extraction 70, search 490, analysis 1,000, steering 70") in out + + +def test_old_journal_without_token_usage_loads(tmp_path: Path): + journal = SearchJournal() + path = tmp_path / "old.json" + journal.export_json(path) + raw = json.loads(path.read_text()) + assert "token_usage" not in raw + loaded = load_journal_export(path) + assert loaded.token_usage is None + assert len(TokenLedger.from_dict(loaded.token_usage)) == 0 + + +def _save(tmp_path: Path, token_usage: dict | None) -> Path: + return save_session( + save_dir=tmp_path / "save", goal="g", application="opflow", + base_case_path=tmp_path / "case.m", config_path=None, journal=SearchJournal(), + steering_history=[], active_steering_directives=[], current_network=None, + total_prompt_tokens=300, total_completion_tokens=100, last_iteration=0, + token_usage=token_usage, + ) + + +def test_session_checkpoint_round_trip(tmp_path: Path): + ledger = TokenLedger() + ledger.record("search", _text_response("", 300, 100), 1) + loaded = load_session(_save(tmp_path, ledger.to_dict())) + assert TokenLedger.from_dict(loaded["token_usage"]).calls == ledger.calls + + +def test_old_session_checkpoint_gives_legacy_totals(tmp_path: Path): + save_dir = _save(tmp_path, None) + raw = json.loads((save_dir / "session.json").read_text()) + raw["format_version"] = "1.1" + raw.pop("token_usage") + (save_dir / "session.json").write_text(json.dumps(raw)) + loaded = load_session(save_dir) + assert loaded["token_usage"] is None + ledger = TokenLedger.from_legacy_totals(loaded["total_prompt_tokens"], + loaded["total_completion_tokens"]) + assert ledger.total == 400 and ledger.by_phase["search"]["calls"] == 1 + + +def test_regenerate_records_report_call(tmp_path: Path): + from tests.journal.test_regenerate_report import _cfg + + journal = SearchJournal() + session_ledger = TokenLedger() + session_ledger.record("search", _text_response("", 1000, 200), 1) + journal.token_usage = session_ledger.to_dict() + export = tmp_path / "journal.json" + journal.export_json(export) + + backend = MagicMock() + backend.complete.return_value = _text_response("no json here", 400, 100) + captured = {} + + class _SpyGenerator: + def generate(self, session, **kwargs): + captured["session"] = session + return b"%PDF-1.4\n%%EOF" + + with patch("agentigrid.exago_engine.regenerate_report._import_report_generator", + return_value=_SpyGenerator): + regenerate_from_journal(journal_json_path=export, config=_cfg(tmp_path), + backend=backend, out_dir=tmp_path / "out", goal="g", + generate_analysis=True) + + session = captured["session"] + assert backend.complete.call_count == 1 + assert {p: a["calls"] for p, a in session.token_ledger.by_phase.items()} == { + "search": 1, "report": 1, + } + assert session.total_prompt_tokens + session.total_completion_tokens == 1700 + + +def test_session_stub_takes_journal_usage(tmp_path: Path): + from tests.journal.test_regenerate_report import _cfg + + journal = SearchJournal() + ledger = TokenLedger() + ledger.record("analysis", _text_response("", 10, 5)) + journal.token_usage = ledger.to_dict() + session = build_session_stub(journal, _cfg(tmp_path), goal="g") + assert session.total_prompt_tokens == 10 and session.total_completion_tokens == 5 + assert len(session.token_ledger) == 1 + + +def _stub_session(tmp_path: Path): + from tests.journal.test_regenerate_report import _cfg + + journal = SearchJournal() + ledger = TokenLedger() + ledger.record("search", _text_response("", 1000, 200), 1) + journal.token_usage = ledger.to_dict() + return build_session_stub(journal, _cfg(tmp_path), goal="g") + + +def test_gui_summary_analysis_uses_session_ledger(tmp_path: Path): + from launcher.session_manager import SessionManager + + session = _stub_session(tmp_path) + backend = MockBackend([_classification_response()]) + with patch("launcher.session_manager.create_backend", return_value=backend): + SessionManager().get_summary_analysis(session) + + assert len(backend.calls) == 1 + assert {p: a["calls"] for p, a in session.token_ledger.by_phase.items()} == { + "search": 1, "analysis": 1, + } + assert session.total_prompt_tokens == 1700 + assert session.total_completion_tokens == 500 + assert session.journal.token_usage == session.token_ledger.to_dict() + + +def test_report_token_line_shows_breakdown(tmp_path: Path): + report_generator = pytest.importorskip("launcher.report_generator") + + session = _stub_session(tmp_path) + session.token_ledger.record("analysis", _text_response("", 300, 100)) + session.total_prompt_tokens = session.token_ledger.total_prompt + session.total_completion_tokens = session.token_ledger.total_completion + assert report_generator._token_usage_line(session) == ( + "Token usage: 1,600 in 2 LLM calls (search 1,200, analysis 400)" + ) + session.total_prompt_tokens = session.total_completion_tokens = 0 + assert report_generator._token_usage_line(session) == "Token usage: N/A" diff --git a/agentigrid/tests/e2e/test_e2e.py b/agentigrid/tests/e2e/test_e2e.py index 05ae5f1f..2244ff43 100644 --- a/agentigrid/tests/e2e/test_e2e.py +++ b/agentigrid/tests/e2e/test_e2e.py @@ -113,7 +113,7 @@ def _make_config( data=DataConfig(data_dir=DATA_DIR), llm=LLMConfig( backend="anthropic" if _has_anthropic_key else "openai", - model="claude-sonnet-4-20250514" if _has_anthropic_key else "gpt-4o-mini", + model="claude-sonnet-4-6" if _has_anthropic_key else "gpt-4o-mini", api_key_env="ANTHROPIC_API_KEY" if _has_anthropic_key else "OPENAI_API_KEY", openai_base_url=None, ollama_host="http://localhost:11434", diff --git a/agentigrid/tests/gridkit/test_from_exago.py b/agentigrid/tests/gridkit/test_from_exago.py index 74427fb5..e835230e 100644 --- a/agentigrid/tests/gridkit/test_from_exago.py +++ b/agentigrid/tests/gridkit/test_from_exago.py @@ -2,6 +2,7 @@ from __future__ import annotations +import copy import json import math from pathlib import Path @@ -13,6 +14,7 @@ from agentigrid.gridkit_engine.from_exago import ( SteadyStateError, apply_steady_state, + exago_shifts_to_degrees, find_latest_steady_state, goal_wants_steady_state, ) @@ -178,6 +180,98 @@ def test_shunts_and_branches_from_exago(pair): assert any("64-82" in n for n in notes) +@needs_cases +def test_exago_phase_shifts_read_as_radians(pair): + # ExaGO's saved solution writes shift angles in radians; GridKit wants radians. + case, net = pair + br = net.branches[0] + br.angle = math.radians(-12.0) + notes = exago_shifts_to_degrees(net) + assert br.angle == pytest.approx(-12.0) and "radians to degrees" in notes[0] + out, _ = apply_steady_state(case, net) + phases = [d["params"]["phase"] for d in _devices(out, "Branch") if d["params"]["phase"]] + assert phases == [pytest.approx(math.radians(-12.0))] + + +def _with_static_gen(case, net, name): + # A 5 MW renewable at bus 2, as GridKit's 10k case writes it; ExaGO nets it from Pd. + case["header"]["case_name"] = name + case["devices"].append({"class": "LoadZIP", "id": "load_2_static_gen_1", "ports": {"bus": 2}, + "params": {"Pnom": -0.05, "Qnom": -0.01, "alphaI": 1.0, "alphaP": 0.0}}) + b = next(b for b in net.buses if b.bus_i == 2) + b.Pd -= 5.0 + b.Qd -= 1.0 + return b + + +@needs_cases +def test_static_gen_kept_for_10k_case(pair): + case, net = pair + _with_static_gen(case, net, "ACTIVSg10k") + out, notes = apply_steady_state(case, net) + at2 = {d["id"]: d["params"] for d in _devices(out, "LoadZIP") if d["ports"]["bus"] == 2} + assert at2["load_2_static_gen_1"] == {"Pnom": -0.05, "Qnom": -0.01, "alphaI": 1.0, "alphaP": 0.0} + assert at2["load_2"]["Pnom"] == pytest.approx(7.39 / 100) and at2["load_2"]["alphaI"] == 0.0 + assert any("static_gen" in n for n in notes) + + +@needs_cases +def test_static_gen_rule_off_for_other_cases(pair): + case, net = pair + b = _with_static_gen(case, net, "ACTIVSg200") + out, notes = apply_steady_state(case, net) + at2 = [d for d in _devices(out, "LoadZIP") if d["ports"]["bus"] == 2] + assert [d["id"] for d in at2] == ["load_2"] or [d["id"] for d in at2] == ["load_2_static_gen_1"] + assert at2[0]["params"]["Pnom"] == pytest.approx(b.Pd / 100) + assert not any("static_gen" in n for n in notes) + + +@needs_cases +def test_machines_pair_in_file_order(pair): + # ACTIVSg10k lists a bus's units as 1, 10, 2, ...; pairing must not re-sort them. + case, net = pair + gen = next(g for g in net.generators if g.bus == 135 and g.status) + m1 = next(d for d in _devices(case, "Genrou") if d["ports"]["bus"] == 135) + m10 = copy.deepcopy(m1) + m10["id"] = m1["id"].replace("_1_", "_10_", 1) + case["devices"].insert(case["devices"].index(m1) + 1, m10) + second = copy.deepcopy(gen) + second.Pg = 12.34 + net.generators.insert(net.generators.index(gen) + 1, second) + out, _ = apply_steady_state(case, net) + p0 = {d["id"]: d["params"]["p0"] for d in _devices(out, "Genrou") if d["ports"]["bus"] == 135} + assert p0[m1["id"]] == pytest.approx(gen.Pg / 100) + assert p0[m10["id"]] == pytest.approx(12.34 / 100) + + +@needs_cases +@pytest.mark.parametrize("offset,snapped", [(0.005, True), (-0.005, True), (0.05, False)]) +def test_dispatch_at_limit_snapped(pair, offset, snapped): + case, net = pair + gen = next(g for g in net.generators if g.bus == 135 and g.status) + gen.Pg = gen.Pmax + offset + out, notes = apply_steady_state(case, net) + m = next(d for d in _devices(out, "Genrou") if d["ports"]["bus"] == 135) + assert m["params"]["p0"] == pytest.approx((gen.Pmax if snapped else gen.Pg) / 100) + assert any("set to the limit" in n for n in notes) == snapped + + +@needs_cases +def test_dispatch_at_valve_limit_set_inside(pair): + # Tgov1 rejects a start outside [Pvmin, Pvmax] with a 4-epsilon tolerance. + case, net = pair + gov = next(d for d in _devices(case, "Tgov1")) + m = next(d for d in case["devices"] if d["class"] in MACHINE_CLASSES + and d["ports"].get("pmech") == gov["ports"]["pmech"]) + gen = next(g for g in net.generators if g.bus == m["ports"]["bus"] and g.status) + gen.Pg = gov["params"]["Pvmax"] * gov["params"]["Trate"] + 0.004 + out, notes = apply_steady_state(case, net) + p0 = next(d for d in out["devices"] if d["id"] == m["id"])["params"]["p0"] + assert p0 * 100 / gov["params"]["Trate"] <= gov["params"]["Pvmax"] + assert p0 * 100 == pytest.approx(gov["params"]["Pvmax"] * gov["params"]["Trate"]) + assert any("valve limit" in n for n in notes) + + @needs_cases def test_generator_without_model_is_an_error(pair): case, net = pair diff --git a/agentigrid/tests/gridkit/test_token_accounting.py b/agentigrid/tests/gridkit/test_token_accounting.py new file mode 100644 index 00000000..278c93e0 --- /dev/null +++ b/agentigrid/tests/gridkit/test_token_accounting.py @@ -0,0 +1,60 @@ +"""Every LLM call of a GridKit session is counted (all under ``search``).""" + +from __future__ import annotations + +import json + +from agentigrid.backends.base import LLMResponse +from tests.gridkit.test_agent_loop import IEEE39, ScriptedBackend, _loop, pytestmark # noqa: F401 + + +def _ledger_matches_backend(session, backend) -> None: + ledger = session.token_ledger + assert len(ledger) == len(backend.prompts) + assert set(ledger.by_phase) == {"search"} + assert session.total_prompt_tokens == ledger.total_prompt + assert session.total_completion_tokens == ledger.total_completion + assert ledger.total == sum( + (c["prompt_tokens"] or 0) + (c["completion_tokens"] or 0) for c in ledger.calls + ) + + +def test_parse_failures_counted(tmp_path): + loop, backend = _loop(tmp_path, ["a", "b", "c"]) + session = loop.run(IEEE39, "anything") + assert session.termination_reason == "parse_failures" + _ledger_matches_backend(session, backend) + assert session.token_ledger.total == 3 * 15 + assert [c["prompt_chars"] for c in session.token_ledger.calls] == [ + len(system) + len(user) for system, user in backend.prompts + ] + assert [c["completion_chars"] for c in session.token_ledger.calls] == [1, 1, 1] + assert [c["iteration"] for c in session.token_ledger.calls] == [1, 2, 3] + + +def test_journal_block_and_missing_usage(tmp_path): + loop, backend = _loop(tmp_path, []) + + replies = iter(["not json", {"action": "complete", "summary": "VERDICT: ok", "details": ""}]) + + def complete(system_prompt, user_prompt, temperature=None): + backend.prompts.append((system_prompt, user_prompt)) + reply = next(replies) + is_json = isinstance(reply, dict) + return LLMResponse(raw_text=json.dumps(reply) if is_json else reply, + json_data=reply if is_json else None, + json_error=None if is_json else "not JSON", + model="scripted", backend="scripted", + prompt_tokens=None if not is_json else 20, # usage missing once + completion_tokens=None if not is_json else 8) + + backend.complete = complete + session = loop.run(IEEE39, "anything") + assert session.termination_reason == "completed" + _ledger_matches_backend(session, backend) + assert session.token_ledger.calls_without_usage == 1 + + journal = json.loads(session.journal_path.read_text()) + assert journal["token_usage"] == session.token_ledger.to_dict() + assert journal["token_usage"]["calls"] == 2 + assert journal["token_usage"]["total_tokens"] == 28 diff --git a/agentigrid/tests/journal/test_iteration_count.py b/agentigrid/tests/journal/test_iteration_count.py new file mode 100644 index 00000000..a78b2161 --- /dev/null +++ b/agentigrid/tests/journal/test_iteration_count.py @@ -0,0 +1,51 @@ +"""Iteration count shown to users: agent iterations plus the base case.""" + +from __future__ import annotations + +from agentigrid.exago_engine.journal import SearchJournal, format_iteration_count + +from tests.journal.test_journal import _make_entry + + +def _journal(iterations, discarded=()): + j = SearchJournal() + for i in iterations: + j.add_entry(_make_entry(i)) + j.discarded_actions = [{"iteration": i, "kind": "rejected"} for i in discarded] + return j + + +def test_base_case_is_not_counted_as_an_agent_iteration(): + stats = _journal([0, 1, 2]).summary_stats() + assert stats["total_iterations"] == 3 # unchanged: journal entries + assert stats["llm_iterations"] == 2 + assert stats["has_base_case"] is True + assert format_iteration_count(stats) == "2 + base case" + + +def test_rejected_iterations_still_count(): + # max 4: iteration 3 was rejected (no journal entry, kept in discarded_actions) + stats = _journal([0, 1, 2, 4], discarded=[3]).summary_stats() + assert stats["total_iterations"] == 4 + assert stats["llm_iterations"] == 4 + assert format_iteration_count(stats) == "4 + base case" + + +def test_several_entries_for_one_iteration_count_once(): + stats = _journal([0, 1, 1, 2]).summary_stats() + assert stats["llm_iterations"] == 2 + + +def test_no_base_case(): + stats = _journal([1, 2]).summary_stats() + assert format_iteration_count(stats) == "2" + + +def test_old_stats_without_new_keys_fall_back(): + assert format_iteration_count({"total_iterations": 5}) == "5" + + +def test_empty_journal(): + stats = SearchJournal().summary_stats() + assert stats["llm_iterations"] == 0 and stats["has_base_case"] is False + assert format_iteration_count(stats) == "0" diff --git a/agentigrid/tests/llm/test_backends.py b/agentigrid/tests/llm/test_backends.py index 2111ecfe..97ed40b4 100644 --- a/agentigrid/tests/llm/test_backends.py +++ b/agentigrid/tests/llm/test_backends.py @@ -233,7 +233,7 @@ class TestAnthropicLive: def test_simple_prompt(self): cfg = _dummy_config( backend="anthropic", - model="claude-sonnet-4-20250514", + model="claude-sonnet-4-6", api_key_env="ANTHROPIC_API_KEY", ) backend = create_backend(cfg) diff --git a/agentigrid/tests/llm/test_prompt_cache.py b/agentigrid/tests/llm/test_prompt_cache.py new file mode 100644 index 00000000..1fe46760 --- /dev/null +++ b/agentigrid/tests/llm/test_prompt_cache.py @@ -0,0 +1,227 @@ +"""Prompt caching (Anthropic) and its token accounting. + +Caching must be billing-only: the model sees the same system text either way, +``prompt_tokens`` keeps meaning TOTAL input tokens, and the cache breakdown is +reported separately. The session's TokenLedger records the breakdown, API +errors and max_tokens cut-offs per call. +""" + +from __future__ import annotations + +import os +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from agentigrid.backends import LLMResponse +from agentigrid.config import LLMConfig, load_config +from agentigrid.token_usage import TokenLedger + +try: + import anthropic as _anthropic # noqa: F401 + _has_anthropic = True +except ImportError: + _has_anthropic = False + + +def _cfg(**over) -> LLMConfig: + base = dict( + backend="anthropic", model="claude-sonnet-4-6", api_key_env="TEST_API_KEY", + openai_base_url=None, ollama_host="http://localhost:11434", + ollama_cloud_host=None, temperature=0.3, max_tokens=256, + ) + base.update(over) + return LLMConfig(**base) + + +def _fake_message(input_tokens, creation, read, output_tokens=7, text='{"action":"complete"}'): + usage = SimpleNamespace( + input_tokens=input_tokens, + cache_creation_input_tokens=creation, + cache_read_input_tokens=read, + output_tokens=output_tokens, + ) + return SimpleNamespace(content=[SimpleNamespace(type="text", text=text)], usage=usage) + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + +def test_prompt_cache_defaults_on(): + cfg = load_config(None) + assert cfg.llm.prompt_cache is True + + +def test_prompt_cache_can_be_disabled_by_override(): + cfg = load_config(None, cli_overrides={"llm.prompt_cache": False}) + assert cfg.llm.prompt_cache is False + + +# --------------------------------------------------------------------------- +# Anthropic backend request shape + token accounting +# --------------------------------------------------------------------------- + +@pytest.mark.skipif(not _has_anthropic, reason="anthropic package not installed") +@patch.dict(os.environ, {"TEST_API_KEY": "sk-test"}) +def test_system_prompt_sent_as_cached_block(): + from agentigrid.backends.anthropic_backend import AnthropicBackend + + backend = AnthropicBackend(_cfg(prompt_cache=True)) + with patch.object(backend._client.messages, "create", + return_value=_fake_message(50, 4000, 0)) as create: + backend.complete("SYSTEM TEXT", "user text") + system = create.call_args.kwargs["system"] + assert system == [{"type": "text", "text": "SYSTEM TEXT", + "cache_control": {"type": "ephemeral"}}] + # the user turn is untouched + assert create.call_args.kwargs["messages"] == [{"role": "user", "content": "user text"}] + + +@pytest.mark.skipif(not _has_anthropic, reason="anthropic package not installed") +@patch.dict(os.environ, {"TEST_API_KEY": "sk-test"}) +def test_cache_off_sends_plain_string(): + from agentigrid.backends.anthropic_backend import AnthropicBackend + + backend = AnthropicBackend(_cfg(prompt_cache=False)) + with patch.object(backend._client.messages, "create", + return_value=_fake_message(4050, None, None)) as create: + resp = backend.complete("SYSTEM TEXT", "user text") + assert create.call_args.kwargs["system"] == "SYSTEM TEXT" + assert resp.prompt_tokens == 4050 + assert resp.cache_creation_tokens is None and resp.cache_read_tokens is None + + +@pytest.mark.skipif(not _has_anthropic, reason="anthropic package not installed") +@patch.dict(os.environ, {"TEST_API_KEY": "sk-test"}) +def test_prompt_tokens_is_total_input_with_breakdown(): + from agentigrid.backends.anthropic_backend import AnthropicBackend + + backend = AnthropicBackend(_cfg(prompt_cache=True)) + # first call writes the cache, second reads it + with patch.object(backend._client.messages, "create", + side_effect=[_fake_message(50, 4000, 0), _fake_message(60, 0, 4000)]): + first = backend.complete("S", "u1") + second = backend.complete("S", "u2") + assert (first.prompt_tokens, first.cache_creation_tokens, first.cache_read_tokens) == (4050, 4000, 0) + assert (second.prompt_tokens, second.cache_creation_tokens, second.cache_read_tokens) == (4060, 0, 4000) + assert first.json_data == {"action": "complete"} + + +# --------------------------------------------------------------------------- +# TokenLedger: cache breakdown, API errors, cut-offs +# --------------------------------------------------------------------------- + +def _resp(pt, ct, cw=None, cr=None, **kw): + return LLMResponse(raw_text="x", json_data=None, json_error=None, model="m", + backend="stub", prompt_tokens=pt, completion_tokens=ct, + cache_creation_tokens=cw, cache_read_tokens=cr, **kw) + + +def test_ledger_records_cache_breakdown_per_call_and_in_totals(): + ledger = TokenLedger() + ledger.record("search", _resp(4050, 10, 4000, 0), 1) + ledger.record("search", _resp(4060, 12, 0, 4000), 2) + ledger.record("analysis", _resp(None, None)) + d = ledger.to_dict() + assert (d["calls"], d["prompt_tokens"], d["completion_tokens"]) == (3, 8110, 22) + assert (d["cache_creation_tokens"], d["cache_read_tokens"]) == (4000, 4000) + assert (d["api_errors"], d["truncated"], d["calls_without_usage"]) == (0, 0, 1) + assert d["per_call"][1]["cache_read_tokens"] == 4000 + assert "4,000 prompt tokens read from cache" in ledger.summary_text() + + +def test_ledger_round_trip_keeps_new_fields_and_reads_older_records(): + ledger = TokenLedger() + ledger.record("search", _resp(4050, 10, 4000, 0, stop_reason="max_tokens"), 1) + again = TokenLedger.from_dict(ledger.to_dict()) + assert again.to_dict() == ledger.to_dict() + # a per-call record written before these fields existed + old = TokenLedger.from_dict({"per_call": [{"phase": "search", "prompt_tokens": 5, + "completion_tokens": 1}]}) + assert (old.total_cache_read, old.api_errors, old.truncated) == (0, 0, 0) + + +def test_journal_exports_token_usage_and_llm_settings(tmp_path): + import json + from agentigrid.exago_engine.journal import SearchJournal + + ledger = TokenLedger() + ledger.record("search", _resp(100, 5, 80, 0), 1) + j = SearchJournal() + j.token_usage = ledger.to_dict() + j.llm_settings = {"backend": "anthropic", "model": "m", "max_tokens": 256, + "prompt_cache": True, "temperature_sent": True} + out = tmp_path / "j.json" + j.export_json(out) + data = json.loads(out.read_text()) + assert data["token_usage"]["cache_creation_tokens"] == 80 + assert data["llm_settings"]["prompt_cache"] is True + assert "llm_usage" not in data + + +def test_temperature_rejected_retries_without_it_and_remembers(): + from agentigrid.backends.anthropic_backend import AnthropicBackend + backend = AnthropicBackend(_cfg(model="claude-sonnet-5")) + calls = [] + + def create(**kw): + calls.append(kw) + if "temperature" in kw: + raise RuntimeError("Error code: 400 - `temperature` is deprecated for this model.") + return _fake_message(10, 0, 0) + + with patch.object(backend._client.messages, "create", side_effect=create): + r1 = backend.complete("sys", "u1") + r2 = backend.complete("sys", "u2") + assert not r1.api_error and r1.json_data == {"action": "complete"} + assert ["temperature" in c for c in calls] == [True, False, False] # retried once, then omitted + assert backend.temperature_sent is False and r2.prompt_tokens == 10 + + +@pytest.mark.skipif(not _has_anthropic, reason="anthropic not installed") +@patch.dict(os.environ, {"TEST_API_KEY": "sk-test"}) +def test_other_api_errors_are_flagged_and_counted(): + from agentigrid.backends.anthropic_backend import AnthropicBackend + backend = AnthropicBackend(_cfg()) + ledger = TokenLedger() + with patch.object(backend._client.messages, "create", side_effect=RuntimeError("401 invalid x-api-key")): + r = backend.complete("sys", "u") + ledger.record("search", r, 1) + assert r.api_error and r.json_data is None + t = ledger.to_dict() + assert t["calls"] == 1 and t["api_errors"] == 1 and t["prompt_tokens"] == 0 + with patch.object(backend._client.messages, "create", return_value=_fake_message(5, 0, 0)): + backend.complete("sys", "u") + assert backend.temperature_sent is True + + +# --------------------------------------------------------------------------- +# Output-token cap: stop_reason telemetry and AGENTIGRID_MAX_TOKENS +# --------------------------------------------------------------------------- + +@pytest.mark.skipif(not _has_anthropic, reason="anthropic not installed") +@patch.dict(os.environ, {"TEST_API_KEY": "sk-test"}) +def test_max_tokens_stop_is_reported_and_counted(): + from agentigrid.backends.anthropic_backend import AnthropicBackend + backend = AnthropicBackend(_cfg()) + ledger = TokenLedger() + msg = _fake_message(10, 0, 0, output_tokens=256, text="") + msg.content = [SimpleNamespace(type="thinking", thinking="...")] + msg.stop_reason = "max_tokens" + with patch.object(backend._client.messages, "create", return_value=msg): + r = backend.complete("sys", "u") + ledger.record("search", r, 1) + assert r.stop_reason == "max_tokens" and r.content_types == ["thinking"] + assert r.raw_text == "" and r.json_data is None and not r.api_error + assert ledger.to_dict()["truncated"] == 1 + + +def test_max_tokens_env_override(monkeypatch): + monkeypatch.setenv("AGENTIGRID_MAX_TOKENS", "16000") + assert load_config(None).llm.max_tokens == 16000 + monkeypatch.setenv("AGENTIGRID_MAX_TOKENS", "lots") + assert load_config(None).llm.max_tokens == 4096 # ignored, default kept + monkeypatch.delenv("AGENTIGRID_MAX_TOKENS") + assert load_config(None).llm.max_tokens == 4096 diff --git a/agentigrid/tests/llm/test_token_ledger.py b/agentigrid/tests/llm/test_token_ledger.py new file mode 100644 index 00000000..8ff5c6b2 --- /dev/null +++ b/agentigrid/tests/llm/test_token_ledger.py @@ -0,0 +1,137 @@ +"""Unit tests for agentigrid.token_usage.TokenLedger.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from agentigrid.backends.base import LLMResponse +from agentigrid.token_usage import PHASES, TokenLedger + + +def _resp(prompt, completion, model="m", backend="b") -> LLMResponse: + return LLMResponse(raw_text="", json_data=None, json_error=None, model=model, + backend=backend, prompt_tokens=prompt, completion_tokens=completion) + + +def _ledger() -> TokenLedger: + ledger = TokenLedger() + ledger.record("objective_extraction", _resp(50, 20), 0) + ledger.record("search", _resp(100, 40), 1) + ledger.record("search", _resp(120, 60), 2) + ledger.record("analysis", _resp(300, 200)) + return ledger + + +def test_totals_and_phases(): + ledger = _ledger() + assert len(ledger) == 4 + assert ledger.total_prompt == 570 + assert ledger.total_completion == 320 + assert ledger.total == 890 + assert list(ledger.by_phase) == ["objective_extraction", "search", "analysis"] + assert ledger.by_phase["search"] == { + "calls": 2, "prompt_tokens": 220, "completion_tokens": 100, + "total_tokens": 320, "calls_without_usage": 0, + "prompt_chars": 0, "completion_chars": 0, + } + assert sum(p["total_tokens"] for p in ledger.by_phase.values()) == ledger.total + + +def test_none_usage_kept_per_call_and_counted(): + ledger = TokenLedger() + ledger.record("search", _resp(None, None), 1) # e.g. API error + ledger.record("search", _resp(10, None), 2) # partial usage + ledger.record("search", _resp(10, 5), 3) + assert ledger.calls[0]["prompt_tokens"] is None + assert ledger.calls[1]["completion_tokens"] is None + assert ledger.total_prompt == 20 and ledger.total_completion == 5 + assert ledger.calls_without_usage == 2 + assert ledger.by_phase["search"]["calls_without_usage"] == 2 + assert "2 calls without usage data" in ledger.summary_text() + + +def test_partial_response_objects_never_fail(): + ledger = TokenLedger() + ledger.record("analysis", SimpleNamespace(raw_text="x")) + ledger.record("analysis", MagicMock()) + assert len(ledger) == 2 + assert ledger.calls_without_usage == 2 + json.dumps(ledger.to_dict()) # no mock objects leak into the export + + +def test_unknown_phase_rejected(): + with pytest.raises(ValueError): + TokenLedger().record("planning", _resp(1, 1)) + + +def test_to_dict_round_trip(): + ledger = _ledger() + ledger.record("search", _resp(None, None), 3) + data = json.loads(json.dumps(ledger.to_dict())) + assert data["calls"] == 5 + assert data["total_tokens"] == 890 + assert data["calls_without_usage"] == 1 + assert data["by_phase"]["search"]["calls"] == 3 + back = TokenLedger.from_dict(data) + assert back.calls == ledger.calls + assert back.to_dict() == data + + +def test_from_dict_empty_and_legacy(): + assert len(TokenLedger.from_dict(None)) == 0 + assert len(TokenLedger.from_dict({})) == 0 + legacy = TokenLedger.from_legacy_totals(1000, 500) + assert legacy.total == 1500 + assert legacy.by_phase["search"]["calls"] == 1 + assert legacy.calls[0]["carried_over"] is True + assert len(TokenLedger.from_legacy_totals(0, 0)) == 0 + # carried_over survives a round trip + assert TokenLedger.from_dict(legacy.to_dict()).calls[0]["carried_over"] is True + + +def test_summary_text_format(): + text = _ledger().summary_text() + assert text == ( + "Tokens: 890 total (570 prompt + 320 completion) in 4 LLM calls; " + "objective extraction 70, search 320, analysis 500" + ) + assert TokenLedger().summary_text() == "Tokens: 0 total (0 prompt + 0 completion) in 0 LLM calls" + + +def test_phase_names(): + assert PHASES == ("objective_extraction", "search", "analysis", "steering", "report", "other") + + +def test_char_totals_equal_sum_over_calls(): + ledger = TokenLedger() + r1 = _resp(10, 5); r1.raw_text = "abcd" + r2 = _resp(20, 6); r2.raw_text = "xy" + ledger.record("search", r1, 1, prompt_chars=120) + ledger.record("search", r2, 2, prompt_chars=80) + ledger.record("analysis", _resp(30, 7), prompt_chars=None) # call site without a length + ledger.record("other", SimpleNamespace()) # no raw_text at all + + assert [c["prompt_chars"] for c in ledger.calls] == [120, 80, None, None] + assert [c["completion_chars"] for c in ledger.calls] == [4, 2, 0, None] + assert ledger.total_prompt_chars == sum(c["prompt_chars"] or 0 for c in ledger.calls) == 200 + assert ledger.total_completion_chars == sum(c["completion_chars"] or 0 for c in ledger.calls) == 6 + assert ledger.by_phase["search"]["prompt_chars"] == 200 + assert ledger.by_phase["search"]["completion_chars"] == 6 + assert sum(a["prompt_chars"] for a in ledger.by_phase.values()) == ledger.total_prompt_chars + data = ledger.to_dict() + assert data["prompt_chars"] == 200 and data["completion_chars"] == 6 + assert TokenLedger.from_dict(data).calls == ledger.calls + + +def test_old_dict_without_char_fields_loads(): + old = {"per_call": [{"phase": "search", "iteration": 1, "model": "m", "backend": "b", + "prompt_tokens": 10, "completion_tokens": 5}]} + ledger = TokenLedger.from_dict(old) + assert ledger.calls[0]["prompt_chars"] is None + assert ledger.calls[0]["completion_chars"] is None + assert ledger.total == 15 and ledger.total_prompt_chars == 0 + assert ledger.to_dict()["by_phase"]["search"]["prompt_chars"] == 0