diff --git a/.env.example b/.env.example index 6a5e9b7..5b5beff 100644 --- a/.env.example +++ b/.env.example @@ -15,11 +15,6 @@ TAVILY_API_KEY= # Other providers: azure/, bedrock/, groq/, ollama/, together_ai/ (see litellm docs) PENTESTAGENT_MODEL=gpt-5 -# Provider selection: -# Note: The app determines provider from `PENTESTAGENT_MODEL` prefix -# (e.g., `ollama/...`, `gpt-5`, `claude-...`, `gemini/...`). No separate -# `LLM_PROVIDER` variable is used. - # Ollama base URL (set this when using an `ollama/...` model) # Example: http://127.0.0.1:11434 or http://192.168.0.165:11434 OLLAMA_BASE_URL=http://127.0.0.1:11434 @@ -32,7 +27,7 @@ OLLAMA_BASE_URL=http://127.0.0.1:11434 PENTESTAGENT_EMBEDDINGS=local # Settings -PENTESTAGENT_DEBUG=true +PENTESTAGENT_DEBUG=false # Optional: manually declare model/context and daily token budgeting # Useful when provider metadata isn't available or you want to enforce local limits. diff --git a/pentestagent/interface/cli.py b/pentestagent/interface/cli.py index 9f360fd..bbaf25e 100644 --- a/pentestagent/interface/cli.py +++ b/pentestagent/interface/cli.py @@ -1,5 +1,6 @@ """Non-interactive CLI mode for PentestAgent.""" +import ast import asyncio import time from datetime import datetime @@ -9,9 +10,6 @@ from rich.console import Console from rich.markdown import Markdown from rich.panel import Panel from rich.text import Text -import json -import ast -import re from ..config.constants import AGENT_MAX_ITERATIONS, ORCHESTRATOR_MAX_ITERATIONS @@ -206,7 +204,9 @@ async def run_cli( except Exception: parsed_for_summary = None - if isinstance(parsed_for_summary, dict) and isinstance(parsed_for_summary.get("summary"), str): + if isinstance(parsed_for_summary, dict) and isinstance( + parsed_for_summary.get("summary"), str + ): console.print( Panel( Markdown(parsed_for_summary.get("summary")), diff --git a/pentestagent/interface/tui.py b/pentestagent/interface/tui.py index 152c1f6..c839fee 100644 --- a/pentestagent/interface/tui.py +++ b/pentestagent/interface/tui.py @@ -426,13 +426,15 @@ class MemoryDiagnostics(Static): # Use a consistent bar width for all bars and align labels bar_width = 28 labels = ["Tokens:", "Messages:", "Retention:"] - label_width = max(len(l) for l in labels) + label_width = max(len(label_text) for label_text in labels) # Tokens line ratio = current_tokens / max(1, budget) bar = self._bar(ratio, width=bar_width) label = "Tokens:".ljust(label_width) - text.append(f"{label} [{bar}] {current_tokens:,} / {budget:,}\n", style="#9a9a9a") + text.append( + f"{label} [{bar}] {current_tokens:,} / {budget:,}\n", style="#9a9a9a" + ) # Messages line (scale messages to an expected max window) expected_msgs_max = max(1, recent_keep * 6) @@ -445,7 +447,9 @@ class MemoryDiagnostics(Static): k_ratio = min(1.0, recent_keep / max(1, recent_keep)) keep_bar = self._bar(k_ratio, width=bar_width) label = "Retention:".ljust(label_width) - text.append(f"{label} [{keep_bar}] keeping last {recent_keep}\n", style="#9a9a9a") + text.append( + f"{label} [{keep_bar}] keeping last {recent_keep}\n", style="#9a9a9a" + ) # Summary status summary_state = "active" if has_summary else "inactive" @@ -453,7 +457,9 @@ class MemoryDiagnostics(Static): text.append(f"Summary: {emoji} {summary_state}\n", style="#9a9a9a") # Summarized / threshold - text.append(f"Summarized: {summarized_count} / {thresh:,}\n", style="#9a9a9a") + text.append( + f"Summarized: {summarized_count} / {thresh:,}\n", style="#9a9a9a" + ) text.append(f"Threshold: {thresh:,}\n", style="#9a9a9a") except Exception as e: @@ -501,7 +507,10 @@ class TokenDiagnostics(Static): text.append("Token Usage Diagnostics\n", style="bold #d4d4d4") if not token_tracker: - text.append("Token tracker not available (tools/token_tracker).\n", style="#9a9a9a") + text.append( + "Token tracker not available (tools/token_tracker).\n", + style="#9a9a9a", + ) return text stats = token_tracker.get_stats_sync() @@ -569,7 +578,10 @@ class TokenDiagnostics(Static): text.append("Environment configuration errors:\n", style="#ef4444") for e in env_errors: text.append(f" - {e}\n", style="#9a9a9a") - text.append("\nSet environment variables correctly to compute costs.\n", style="#9a9a9a") + text.append( + "\nSet environment variables correctly to compute costs.\n", + style="#9a9a9a", + ) return text # Compute costs @@ -580,7 +592,10 @@ class TokenDiagnostics(Static): else: # Require per-direction costs to be present to compute if input_cost_per_m is None or output_cost_per_m is None: - text.append("Cost vars missing. Set COST_PER_MILLION or both INPUT_COST_PER_MILLION and OUTPUT_COST_PER_MILLION.\n", style="#9a9a9a") + text.append( + "Cost vars missing. Set COST_PER_MILLION or both INPUT_COST_PER_MILLION and OUTPUT_COST_PER_MILLION.\n", + style="#9a9a9a", + ) # Still show numeric token stats below input_cost = output_cost = None else: @@ -618,7 +633,7 @@ class TokenDiagnostics(Static): "Current date:", "Reset occurred:", ] - label_width = max(len(l) for l in labels) + label_width = max(len(label_text) for label_text in labels) # Last command tokens label = "Last command:".ljust(label_width) @@ -1137,6 +1152,7 @@ class PentestAgentTUI(App): # Mount a new diagnostics panel with a unique ID and scroll into view try: import uuid + panel_id = f"memory-diagnostics-{uuid.uuid4().hex}" except Exception: panel_id = None @@ -1158,6 +1174,7 @@ class PentestAgentTUI(App): # Mount a new diagnostics panel with a unique ID and scroll into view try: import uuid + panel_id = f"token-diagnostics-{uuid.uuid4().hex}" except Exception: panel_id = None @@ -1275,10 +1292,14 @@ class PentestAgentTUI(App): scroll = self.query_one("#chat-scroll", ScrollableContainer) updated = False for child in scroll.children: - if isinstance(child, SystemMessage) and "PentestAgent ready" in getattr(child, "message_content", ""): + if isinstance(child, SystemMessage) and "PentestAgent ready" in getattr( + child, "message_content", "" + ): # Append Target line if not already present if "Target:" not in child.message_content: - child.message_content = child.message_content + f"\n Target: {target}" + child.message_content = ( + child.message_content + f"\n Target: {target}" + ) try: child.refresh() except Exception: diff --git a/pentestagent/tools/token_tracker.py b/pentestagent/tools/token_tracker.py index 3cb32a8..848ffd2 100644 --- a/pentestagent/tools/token_tracker.py +++ b/pentestagent/tools/token_tracker.py @@ -9,7 +9,7 @@ import json import threading from datetime import date from pathlib import Path -from typing import Dict, Any +from typing import Any, Dict # Persistent storage (loot root) _data_file: Path = Path("loot/token_usage.json")