strix-agent 0.1.1__py3-none-any.whl

strix/cli/tracer.py

@@ -0,0 +1,308 @@
+import logging
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any, Optional
+from uuid import uuid4
+
+
+logger = logging.getLogger(__name__)
+
+_global_tracer: Optional["Tracer"] = None
+
+
+def get_global_tracer() -> Optional["Tracer"]:
+    return _global_tracer
+
+
+def set_global_tracer(tracer: "Tracer") -> None:
+    global _global_tracer  # noqa: PLW0603
+    _global_tracer = tracer
+
+
+class Tracer:
+    def __init__(self, run_name: str | None = None):
+        self.run_name = run_name
+        self.run_id = run_name or f"run-{uuid4().hex[:8]}"
+        self.start_time = datetime.now(UTC).isoformat()
+        self.end_time: str | None = None
+
+        self.agents: dict[str, dict[str, Any]] = {}
+        self.tool_executions: dict[int, dict[str, Any]] = {}
+        self.chat_messages: list[dict[str, Any]] = []
+
+        self.vulnerability_reports: list[dict[str, Any]] = []
+        self.final_scan_result: str | None = None
+
+        self.scan_results: dict[str, Any] | None = None
+        self.scan_config: dict[str, Any] | None = None
+        self.run_metadata: dict[str, Any] = {
+            "run_id": self.run_id,
+            "run_name": self.run_name,
+            "start_time": self.start_time,
+            "end_time": None,
+            "target": None,
+            "scan_type": None,
+            "status": "running",
+        }
+        self._run_dir: Path | None = None
+        self._next_execution_id = 1
+        self._next_message_id = 1
+
+    def set_run_name(self, run_name: str) -> None:
+        self.run_name = run_name
+        self.run_id = run_name
+
+    def get_run_dir(self) -> Path:
+        if self._run_dir is None:
+            workspace_root = Path(__file__).parent.parent.parent
+            runs_dir = workspace_root / "agent_runs"
+            runs_dir.mkdir(exist_ok=True)
+
+            run_dir_name = self.run_name if self.run_name else self.run_id
+            self._run_dir = runs_dir / run_dir_name
+            self._run_dir.mkdir(exist_ok=True)
+
+        return self._run_dir
+
+    def add_vulnerability_report(
+        self,
+        title: str,
+        content: str,
+        severity: str,
+    ) -> str:
+        report_id = f"vuln-{len(self.vulnerability_reports) + 1:04d}"
+
+        report = {
+            "id": report_id,
+            "title": title.strip(),
+            "content": content.strip(),
+            "severity": severity.lower().strip(),
+            "timestamp": datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S UTC"),
+        }
+
+        self.vulnerability_reports.append(report)
+        logger.info(f"Added vulnerability report: {report_id} - {title}")
+        return report_id
+
+    def set_final_scan_result(
+        self,
+        content: str,
+        success: bool = True,
+    ) -> None:
+        self.final_scan_result = content.strip()
+
+        self.scan_results = {
+            "scan_completed": True,
+            "content": content,
+            "success": success,
+        }
+
+        logger.info(f"Set final scan result: success={success}")
+
+    def log_agent_creation(
+        self, agent_id: str, name: str, task: str, parent_id: str | None = None
+    ) -> None:
+        agent_data: dict[str, Any] = {
+            "id": agent_id,
+            "name": name,
+            "task": task,
+            "status": "running",
+            "parent_id": parent_id,
+            "created_at": datetime.now(UTC).isoformat(),
+            "updated_at": datetime.now(UTC).isoformat(),
+            "tool_executions": [],
+        }
+
+        self.agents[agent_id] = agent_data
+
+    def log_chat_message(
+        self,
+        content: str,
+        role: str,
+        agent_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> int:
+        message_id = self._next_message_id
+        self._next_message_id += 1
+
+        message_data = {
+            "message_id": message_id,
+            "content": content,
+            "role": role,
+            "agent_id": agent_id,
+            "timestamp": datetime.now(UTC).isoformat(),
+            "metadata": metadata or {},
+        }
+
+        self.chat_messages.append(message_data)
+        return message_id
+
+    def log_tool_execution_start(self, agent_id: str, tool_name: str, args: dict[str, Any]) -> int:
+        execution_id = self._next_execution_id
+        self._next_execution_id += 1
+
+        now = datetime.now(UTC).isoformat()
+        execution_data = {
+            "execution_id": execution_id,
+            "agent_id": agent_id,
+            "tool_name": tool_name,
+            "args": args,
+            "status": "running",
+            "result": None,
+            "timestamp": now,
+            "started_at": now,
+            "completed_at": None,
+        }
+
+        self.tool_executions[execution_id] = execution_data
+
+        if agent_id in self.agents:
+            self.agents[agent_id]["tool_executions"].append(execution_id)
+
+        return execution_id
+
+    def update_tool_execution(
+        self, execution_id: int, status: str, result: Any | None = None
+    ) -> None:
+        if execution_id in self.tool_executions:
+            self.tool_executions[execution_id]["status"] = status
+            self.tool_executions[execution_id]["result"] = result
+            self.tool_executions[execution_id]["completed_at"] = datetime.now(UTC).isoformat()
+
+    def update_agent_status(self, agent_id: str, status: str) -> None:
+        if agent_id in self.agents:
+            self.agents[agent_id]["status"] = status
+            self.agents[agent_id]["updated_at"] = datetime.now(UTC).isoformat()
+
+    def set_scan_config(self, config: dict[str, Any]) -> None:
+        self.scan_config = config
+        self.run_metadata.update(
+            {
+                "target": config.get("target", {}),
+                "scan_type": config.get("scan_type", "general"),
+                "user_instructions": config.get("user_instructions", ""),
+                "max_iterations": config.get("max_iterations", 200),
+            }
+        )
+
+    def save_run_data(self) -> None:
+        try:
+            run_dir = self.get_run_dir()
+            self.end_time = datetime.now(UTC).isoformat()
+
+            if self.final_scan_result:
+                scan_report_file = run_dir / "scan_report.md"
+                with scan_report_file.open("w", encoding="utf-8") as f:
+                    f.write("# Security Scan Report\n\n")
+                    f.write(
+                        f"**Generated:** {datetime.now(UTC).strftime('%Y-%m-%d %H:%M:%S UTC')}\n\n"
+                    )
+                    f.write(f"{self.final_scan_result}\n")
+                logger.info(f"Saved final scan report to: {scan_report_file}")
+
+            if self.vulnerability_reports:
+                vuln_dir = run_dir / "vulnerabilities"
+                vuln_dir.mkdir(exist_ok=True)
+
+                severity_order = {"critical": 0, "high": 1, "medium": 2, "low": 3, "info": 4}
+                sorted_reports = sorted(
+                    self.vulnerability_reports,
+                    key=lambda x: (severity_order.get(x["severity"], 5), x["timestamp"]),
+                )
+
+                for report in sorted_reports:
+                    vuln_file = vuln_dir / f"{report['id']}.md"
+                    with vuln_file.open("w", encoding="utf-8") as f:
+                        f.write(f"# {report['title']}\n\n")
+                        f.write(f"**ID:** {report['id']}\n")
+                        f.write(f"**Severity:** {report['severity'].upper()}\n")
+                        f.write(f"**Found:** {report['timestamp']}\n\n")
+                        f.write("## Description\n\n")
+                        f.write(f"{report['content']}\n")
+
+                vuln_csv_file = run_dir / "vulnerabilities.csv"
+                with vuln_csv_file.open("w", encoding="utf-8", newline="") as f:
+                    import csv
+
+                    fieldnames = ["id", "title", "severity", "timestamp", "file"]
+                    writer = csv.DictWriter(f, fieldnames=fieldnames)
+                    writer.writeheader()
+
+                    for report in sorted_reports:
+                        writer.writerow(
+                            {
+                                "id": report["id"],
+                                "title": report["title"],
+                                "severity": report["severity"].upper(),
+                                "timestamp": report["timestamp"],
+                                "file": f"vulnerabilities/{report['id']}.md",
+                            }
+                        )
+
+                logger.info(
+                    f"Saved {len(self.vulnerability_reports)} vulnerability reports to: {vuln_dir}"
+                )
+                logger.info(f"Saved vulnerability index to: {vuln_csv_file}")
+
+            logger.info(f"📊 Essential scan data saved to: {run_dir}")
+
+        except (OSError, RuntimeError):
+            logger.exception("Failed to save scan data")
+
+    def _calculate_duration(self) -> float:
+        try:
+            start = datetime.fromisoformat(self.start_time.replace("Z", "+00:00"))
+            if self.end_time:
+                end = datetime.fromisoformat(self.end_time.replace("Z", "+00:00"))
+                return (end - start).total_seconds()
+        except (ValueError, TypeError):
+            pass
+        return 0.0
+
+    def get_agent_tools(self, agent_id: str) -> list[dict[str, Any]]:
+        return [
+            exec_data
+            for exec_data in self.tool_executions.values()
+            if exec_data.get("agent_id") == agent_id
+        ]
+
+    def get_real_tool_count(self) -> int:
+        return sum(
+            1
+            for exec_data in self.tool_executions.values()
+            if exec_data.get("tool_name") not in ["scan_start_info", "subagent_start_info"]
+        )
+
+    def get_total_llm_stats(self) -> dict[str, Any]:
+        from strix.tools.agents_graph.agents_graph_actions import _agent_instances
+
+        total_stats = {
+            "input_tokens": 0,
+            "output_tokens": 0,
+            "cached_tokens": 0,
+            "cache_creation_tokens": 0,
+            "cost": 0.0,
+            "requests": 0,
+            "failed_requests": 0,
+        }
+
+        for agent_instance in _agent_instances.values():
+            if hasattr(agent_instance, "llm") and hasattr(agent_instance.llm, "_total_stats"):
+                agent_stats = agent_instance.llm._total_stats
+                total_stats["input_tokens"] += agent_stats.input_tokens
+                total_stats["output_tokens"] += agent_stats.output_tokens
+                total_stats["cached_tokens"] += agent_stats.cached_tokens
+                total_stats["cache_creation_tokens"] += agent_stats.cache_creation_tokens
+                total_stats["cost"] += agent_stats.cost
+                total_stats["requests"] += agent_stats.requests
+                total_stats["failed_requests"] += agent_stats.failed_requests
+
+        total_stats["cost"] = round(total_stats["cost"], 4)
+
+        return {
+            "total": total_stats,
+            "total_tokens": total_stats["input_tokens"] + total_stats["output_tokens"],
+        }
+
+    def cleanup(self) -> None:
+        self.save_run_data()

strix/llm/__init__.py

@@ -0,0 +1,14 @@
+import litellm
+
+from .config import LLMConfig
+from .llm import LLM
+
+
+__all__ = [
+    "LLM",
+    "LLMConfig",
+]
+
+litellm.drop_params = True
+litellm.suppress_debug_info = True
+litellm.set_verbose = False

strix/llm/config.py

@@ -0,0 +1,19 @@
+import os
+
+
+class LLMConfig:
+    def __init__(
+        self,
+        model_name: str | None = None,
+        temperature: float = 0,
+        enable_prompt_caching: bool = True,
+        prompt_modules: list[str] | None = None,
+    ):
+        self.model_name = model_name or os.getenv("STRIX_LLM", "anthropic/claude-sonnet-4-20250514")
+
+        if not self.model_name:
+            raise ValueError("STRIX_LLM environment variable must be set and not empty")
+
+        self.temperature = max(0.0, min(1.0, temperature))
+        self.enable_prompt_caching = enable_prompt_caching
+        self.prompt_modules = prompt_modules or []

strix/llm/llm.py

@@ -0,0 +1,310 @@
+import logging
+import os
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+import litellm
+from jinja2 import (
+    Environment,
+    FileSystemLoader,
+    select_autoescape,
+)
+from litellm import ModelResponse, completion_cost
+from litellm.utils import supports_prompt_caching
+
+from strix.llm.config import LLMConfig
+from strix.llm.memory_compressor import MemoryCompressor
+from strix.llm.request_queue import get_global_queue
+from strix.llm.utils import _truncate_to_first_function, parse_tool_invocations
+from strix.prompts import load_prompt_modules
+from strix.tools import get_tools_prompt
+
+
+logger = logging.getLogger(__name__)
+
+api_key = os.getenv("LLM_API_KEY")
+if api_key:
+    litellm.api_key = api_key
+
+
+class StepRole(str, Enum):
+    AGENT = "agent"
+    USER = "user"
+    SYSTEM = "system"
+
+
+@dataclass
+class LLMResponse:
+    content: str
+    tool_invocations: list[dict[str, Any]] | None = None
+    scan_id: str | None = None
+    step_number: int = 1
+    role: StepRole = StepRole.AGENT
+
+
+@dataclass
+class RequestStats:
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cached_tokens: int = 0
+    cache_creation_tokens: int = 0
+    cost: float = 0.0
+    requests: int = 0
+    failed_requests: int = 0
+
+    def to_dict(self) -> dict[str, int | float]:
+        return {
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "cached_tokens": self.cached_tokens,
+            "cache_creation_tokens": self.cache_creation_tokens,
+            "cost": round(self.cost, 4),
+            "requests": self.requests,
+            "failed_requests": self.failed_requests,
+        }
+
+
+class LLM:
+    def __init__(self, config: LLMConfig, agent_name: str | None = None):
+        self.config = config
+        self.agent_name = agent_name
+        self._total_stats = RequestStats()
+        self._last_request_stats = RequestStats()
+
+        self.memory_compressor = MemoryCompressor()
+
+        if agent_name:
+            prompt_dir = Path(__file__).parent.parent / "agents" / agent_name
+            prompts_dir = Path(__file__).parent.parent / "prompts"
+
+            loader = FileSystemLoader([prompt_dir, prompts_dir])
+            self.jinja_env = Environment(
+                loader=loader,
+                autoescape=select_autoescape(enabled_extensions=(), default_for_string=False),
+            )
+
+            try:
+                prompt_module_content = load_prompt_modules(
+                    self.config.prompt_modules or [], self.jinja_env
+                )
+
+                def get_module(name: str) -> str:
+                    return prompt_module_content.get(name, "")
+
+                self.jinja_env.globals["get_module"] = get_module
+
+                self.system_prompt = self.jinja_env.get_template("system_prompt.jinja").render(
+                    get_tools_prompt=get_tools_prompt,
+                    loaded_module_names=list(prompt_module_content.keys()),
+                    **prompt_module_content,
+                )
+            except (FileNotFoundError, OSError, ValueError) as e:
+                logger.warning(f"Failed to load system prompt for {agent_name}: {e}")
+                self.system_prompt = "You are a helpful AI assistant."
+        else:
+            self.system_prompt = "You are a helpful AI assistant."
+
+    def _add_cache_control_to_content(
+        self, content: str | list[dict[str, Any]]
+    ) -> str | list[dict[str, Any]]:
+        if isinstance(content, str):
+            return [{"type": "text", "text": content, "cache_control": {"type": "ephemeral"}}]
+        if isinstance(content, list) and content:
+            last_item = content[-1]
+            if isinstance(last_item, dict) and last_item.get("type") == "text":
+                return content[:-1] + [{**last_item, "cache_control": {"type": "ephemeral"}}]
+        return content
+
+    def _is_anthropic_model(self) -> bool:
+        if not self.config.model_name:
+            return False
+        model_lower = self.config.model_name.lower()
+        return any(provider in model_lower for provider in ["anthropic/", "claude"])
+
+    def _calculate_cache_interval(self, total_messages: int) -> int:
+        if total_messages <= 1:
+            return 10
+
+        max_cached_messages = 3
+        non_system_messages = total_messages - 1
+
+        interval = 10
+        while non_system_messages // interval > max_cached_messages:
+            interval += 10
+
+        return interval
+
+    def _prepare_cached_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        if (
+            not self.config.enable_prompt_caching
+            or not supports_prompt_caching(self.config.model_name)
+            or not messages
+        ):
+            return messages
+
+        if not self._is_anthropic_model():
+            return messages
+
+        cached_messages = list(messages)
+
+        if cached_messages and cached_messages[0].get("role") == "system":
+            system_message = cached_messages[0].copy()
+            system_message["content"] = self._add_cache_control_to_content(
+                system_message["content"]
+            )
+            cached_messages[0] = system_message
+
+        total_messages = len(cached_messages)
+        if total_messages > 1:
+            interval = self._calculate_cache_interval(total_messages)
+
+            cached_count = 0
+            for i in range(interval, total_messages, interval):
+                if cached_count >= 3:
+                    break
+
+                if i < len(cached_messages):
+                    message = cached_messages[i].copy()
+                    message["content"] = self._add_cache_control_to_content(message["content"])
+                    cached_messages[i] = message
+                    cached_count += 1
+
+        return cached_messages
+
+    async def generate(
+        self,
+        conversation_history: list[dict[str, Any]],
+        scan_id: str | None = None,
+        step_number: int = 1,
+    ) -> LLMResponse:
+        messages = [{"role": "system", "content": self.system_prompt}]
+
+        compressed_history = list(self.memory_compressor.compress_history(conversation_history))
+
+        conversation_history.clear()
+        conversation_history.extend(compressed_history)
+        messages.extend(compressed_history)
+
+        cached_messages = self._prepare_cached_messages(messages)
+
+        try:
+            response = await self._make_request(cached_messages)
+            self._update_usage_stats(response)
+
+            content = ""
+            if (
+                response.choices
+                and hasattr(response.choices[0], "message")
+                and response.choices[0].message
+            ):
+                content = getattr(response.choices[0].message, "content", "") or ""
+
+            content = _truncate_to_first_function(content)
+
+            if "</function>" in content:
+                function_end_index = content.find("</function>") + len("</function>")
+                content = content[:function_end_index]
+
+            tool_invocations = parse_tool_invocations(content)
+
+            return LLMResponse(
+                scan_id=scan_id,
+                step_number=step_number,
+                role=StepRole.AGENT,
+                content=content,
+                tool_invocations=tool_invocations if tool_invocations else None,
+            )
+
+        except (ValueError, TypeError, RuntimeError):
+            logger.exception("Error in LLM generation")
+            return LLMResponse(
+                scan_id=scan_id,
+                step_number=step_number,
+                role=StepRole.AGENT,
+                content="An error occurred while generating the response",
+                tool_invocations=None,
+            )
+
+    @property
+    def usage_stats(self) -> dict[str, dict[str, int | float]]:
+        return {
+            "total": self._total_stats.to_dict(),
+            "last_request": self._last_request_stats.to_dict(),
+        }
+
+    def get_cache_config(self) -> dict[str, bool]:
+        return {
+            "enabled": self.config.enable_prompt_caching,
+            "supported": supports_prompt_caching(self.config.model_name),
+        }
+
+    async def _make_request(
+        self,
+        messages: list[dict[str, Any]],
+    ) -> ModelResponse:
+        completion_args = {
+            "model": self.config.model_name,
+            "messages": messages,
+            "temperature": self.config.temperature,
+            "stop": ["</function>"],
+        }
+
+        queue = get_global_queue()
+        response = await queue.make_request(completion_args)
+
+        self._total_stats.requests += 1
+        self._last_request_stats = RequestStats(requests=1)
+
+        return response
+
+    def _update_usage_stats(self, response: ModelResponse) -> None:
+        try:
+            if hasattr(response, "usage") and response.usage:
+                input_tokens = getattr(response.usage, "prompt_tokens", 0)
+                output_tokens = getattr(response.usage, "completion_tokens", 0)
+
+                cached_tokens = 0
+                cache_creation_tokens = 0
+
+                if hasattr(response.usage, "prompt_tokens_details"):
+                    prompt_details = response.usage.prompt_tokens_details
+                    if hasattr(prompt_details, "cached_tokens"):
+                        cached_tokens = prompt_details.cached_tokens or 0
+
+                if hasattr(response.usage, "cache_creation_input_tokens"):
+                    cache_creation_tokens = response.usage.cache_creation_input_tokens or 0
+
+            else:
+                input_tokens = 0
+                output_tokens = 0
+                cached_tokens = 0
+                cache_creation_tokens = 0
+
+            try:
+                cost = completion_cost(response) or 0.0
+            except (ValueError, TypeError, RuntimeError) as e:
+                logger.warning(f"Failed to calculate cost: {e}")
+                cost = 0.0
+
+            self._total_stats.input_tokens += input_tokens
+            self._total_stats.output_tokens += output_tokens
+            self._total_stats.cached_tokens += cached_tokens
+            self._total_stats.cache_creation_tokens += cache_creation_tokens
+            self._total_stats.cost += cost
+
+            self._last_request_stats.input_tokens = input_tokens
+            self._last_request_stats.output_tokens = output_tokens
+            self._last_request_stats.cached_tokens = cached_tokens
+            self._last_request_stats.cache_creation_tokens = cache_creation_tokens
+            self._last_request_stats.cost = cost
+
+            if cached_tokens > 0:
+                logger.info(f"Cache hit: {cached_tokens} cached tokens, {input_tokens} new tokens")
+            if cache_creation_tokens > 0:
+                logger.info(f"Cache creation: {cache_creation_tokens} tokens written to cache")
+
+            logger.info(f"Usage stats: {self.usage_stats}")
+        except (AttributeError, TypeError, ValueError) as e:
+            logger.warning(f"Failed to update usage stats: {e}")