strix-agent 0.4.0__py3-none-any.whl

strix/llm/llm.py

@@ -0,0 +1,465 @@
+import logging
+import os
+from dataclasses import dataclass
+from enum import Enum
+from fnmatch import fnmatch
+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
+
+api_base = (
+    os.getenv("LLM_API_BASE")
+    or os.getenv("OPENAI_API_BASE")
+    or os.getenv("LITELLM_BASE_URL")
+    or os.getenv("OLLAMA_API_BASE")
+)
+if api_base:
+    litellm.api_base = api_base
+
+
+class LLMRequestFailedError(Exception):
+    def __init__(self, message: str, details: str | None = None):
+        super().__init__(message)
+        self.message = message
+        self.details = details
+
+
+SUPPORTS_STOP_WORDS_FALSE_PATTERNS: list[str] = [
+    "o1*",
+    "grok-4-0709",
+    "grok-code-fast-1",
+    "deepseek-r1-0528*",
+]
+
+REASONING_EFFORT_PATTERNS: list[str] = [
+    "o1-2024-12-17",
+    "o1",
+    "o3",
+    "o3-2025-04-16",
+    "o3-mini-2025-01-31",
+    "o3-mini",
+    "o4-mini",
+    "o4-mini-2025-04-16",
+    "gemini-2.5-flash",
+    "gemini-2.5-pro",
+    "gpt-5*",
+    "deepseek-r1-0528*",
+    "claude-sonnet-4-5*",
+    "claude-haiku-4-5*",
+]
+
+
+def normalize_model_name(model: str) -> str:
+    raw = (model or "").strip().lower()
+    if "/" in raw:
+        name = raw.split("/")[-1]
+        if ":" in name:
+            name = name.split(":", 1)[0]
+    else:
+        name = raw
+    if name.endswith("-gguf"):
+        name = name[: -len("-gguf")]
+    return name
+
+
+def model_matches(model: str, patterns: list[str]) -> bool:
+    raw = (model or "").strip().lower()
+    name = normalize_model_name(model)
+    for pat in patterns:
+        pat_l = pat.lower()
+        if "/" in pat_l:
+            if fnmatch(raw, pat_l):
+                return True
+        elif fnmatch(name, pat_l):
+            return True
+    return False
+
+
+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, agent_id: str | None = None
+    ):
+        self.config = config
+        self.agent_name = agent_name
+        self.agent_id = agent_id
+        self._total_stats = RequestStats()
+        self._last_request_stats = RequestStats()
+
+        self.memory_compressor = MemoryCompressor(
+            model_name=self.config.model_name,
+            timeout=self.config.timeout,
+        )
+
+        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 set_agent_identity(self, agent_name: str | None, agent_id: str | None) -> None:
+        if agent_name:
+            self.agent_name = agent_name
+        if agent_id:
+            self.agent_id = agent_id
+
+    def _build_identity_message(self) -> dict[str, Any] | None:
+        if not (self.agent_name and str(self.agent_name).strip()):
+            return None
+        identity_name = self.agent_name
+        identity_id = self.agent_id
+        content = (
+            "\n\n"
+            "<agent_identity>\n"
+            "<meta>Internal metadata: do not echo or reference; "
+            "not part of history or tool calls.</meta>\n"
+            "<note>You are now assuming the role of this agent. "
+            "Act strictly as this agent and maintain self-identity for this step. "
+            "Now go answer the next needed step!</note>\n"
+            f"<agent_name>{identity_name}</agent_name>\n"
+            f"<agent_id>{identity_id}</agent_id>\n"
+            "</agent_identity>\n\n"
+        )
+        return {"role": "user", "content": content}
+
+    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(  # noqa: PLR0912, PLR0915
+        self,
+        conversation_history: list[dict[str, Any]],
+        scan_id: str | None = None,
+        step_number: int = 1,
+    ) -> LLMResponse:
+        messages = [{"role": "system", "content": self.system_prompt}]
+
+        identity_message = self._build_identity_message()
+        if identity_message:
+            messages.append(identity_message)
+
+        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 litellm.RateLimitError as e:
+            raise LLMRequestFailedError("LLM request failed: Rate limit exceeded", str(e)) from e
+        except litellm.AuthenticationError as e:
+            raise LLMRequestFailedError("LLM request failed: Invalid API key", str(e)) from e
+        except litellm.NotFoundError as e:
+            raise LLMRequestFailedError("LLM request failed: Model not found", str(e)) from e
+        except litellm.ContextWindowExceededError as e:
+            raise LLMRequestFailedError("LLM request failed: Context too long", str(e)) from e
+        except litellm.ContentPolicyViolationError as e:
+            raise LLMRequestFailedError(
+                "LLM request failed: Content policy violation", str(e)
+            ) from e
+        except litellm.ServiceUnavailableError as e:
+            raise LLMRequestFailedError("LLM request failed: Service unavailable", str(e)) from e
+        except litellm.Timeout as e:
+            raise LLMRequestFailedError("LLM request failed: Request timed out", str(e)) from e
+        except litellm.UnprocessableEntityError as e:
+            raise LLMRequestFailedError("LLM request failed: Unprocessable entity", str(e)) from e
+        except litellm.InternalServerError as e:
+            raise LLMRequestFailedError("LLM request failed: Internal server error", str(e)) from e
+        except litellm.APIConnectionError as e:
+            raise LLMRequestFailedError("LLM request failed: Connection error", str(e)) from e
+        except litellm.UnsupportedParamsError as e:
+            raise LLMRequestFailedError("LLM request failed: Unsupported parameters", str(e)) from e
+        except litellm.BudgetExceededError as e:
+            raise LLMRequestFailedError("LLM request failed: Budget exceeded", str(e)) from e
+        except litellm.APIResponseValidationError as e:
+            raise LLMRequestFailedError(
+                "LLM request failed: Response validation error", str(e)
+            ) from e
+        except litellm.JSONSchemaValidationError as e:
+            raise LLMRequestFailedError(
+                "LLM request failed: JSON schema validation error", str(e)
+            ) from e
+        except litellm.InvalidRequestError as e:
+            raise LLMRequestFailedError("LLM request failed: Invalid request", str(e)) from e
+        except litellm.BadRequestError as e:
+            raise LLMRequestFailedError("LLM request failed: Bad request", str(e)) from e
+        except litellm.APIError as e:
+            raise LLMRequestFailedError("LLM request failed: API error", str(e)) from e
+        except litellm.OpenAIError as e:
+            raise LLMRequestFailedError("LLM request failed: OpenAI error", str(e)) from e
+        except Exception as e:
+            raise LLMRequestFailedError(f"LLM request failed: {type(e).__name__}", str(e)) from e
+
+    @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),
+        }
+
+    def _should_include_stop_param(self) -> bool:
+        if not self.config.model_name:
+            return True
+
+        return not model_matches(self.config.model_name, SUPPORTS_STOP_WORDS_FALSE_PATTERNS)
+
+    def _should_include_reasoning_effort(self) -> bool:
+        if not self.config.model_name:
+            return False
+
+        return model_matches(self.config.model_name, REASONING_EFFORT_PATTERNS)
+
+    async def _make_request(
+        self,
+        messages: list[dict[str, Any]],
+    ) -> ModelResponse:
+        completion_args: dict[str, Any] = {
+            "model": self.config.model_name,
+            "messages": messages,
+            "timeout": self.config.timeout,
+        }
+
+        if self._should_include_stop_param():
+            completion_args["stop"] = ["</function>"]
+
+        if self._should_include_reasoning_effort():
+            completion_args["reasoning_effort"] = "high"
+
+        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 Exception as e:  # noqa: BLE001
+                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 Exception as e:  # noqa: BLE001
+            logger.warning(f"Failed to update usage stats: {e}")

strix/llm/memory_compressor.py

@@ -0,0 +1,212 @@
+import logging
+import os
+from typing import Any
+
+import litellm
+
+
+logger = logging.getLogger(__name__)
+
+
+MAX_TOTAL_TOKENS = 100_000
+MIN_RECENT_MESSAGES = 15
+
+SUMMARY_PROMPT_TEMPLATE = """You are an agent performing context
+condensation for a security agent. Your job is to compress scan data while preserving
+ALL operationally critical information for continuing the security assessment.
+
+CRITICAL ELEMENTS TO PRESERVE:
+- Discovered vulnerabilities and potential attack vectors
+- Scan results and tool outputs (compressed but maintaining key findings)
+- Access credentials, tokens, or authentication details found
+- System architecture insights and potential weak points
+- Progress made in the assessment
+- Failed attempts and dead ends (to avoid duplication)
+- Any decisions made about the testing approach
+
+COMPRESSION GUIDELINES:
+- Preserve exact technical details (URLs, paths, parameters, payloads)
+- Summarize verbose tool outputs while keeping critical findings
+- Maintain version numbers, specific technologies identified
+- Keep exact error messages that might indicate vulnerabilities
+- Compress repetitive or similar findings into consolidated form
+
+Remember: Another security agent will use this summary to continue the assessment.
+They must be able to pick up exactly where you left off without losing any
+operational advantage or context needed to find vulnerabilities.
+
+CONVERSATION SEGMENT TO SUMMARIZE:
+{conversation}
+
+Provide a technically precise summary that preserves all operational security context while
+keeping the summary concise and to the point."""
+
+
+def _count_tokens(text: str, model: str) -> int:
+    try:
+        count = litellm.token_counter(model=model, text=text)
+        return int(count)
+    except Exception:
+        logger.exception("Failed to count tokens")
+        return len(text) // 4  # Rough estimate
+
+
+def _get_message_tokens(msg: dict[str, Any], model: str) -> int:
+    content = msg.get("content", "")
+    if isinstance(content, str):
+        return _count_tokens(content, model)
+    if isinstance(content, list):
+        return sum(
+            _count_tokens(item.get("text", ""), model)
+            for item in content
+            if isinstance(item, dict) and item.get("type") == "text"
+        )
+    return 0
+
+
+def _extract_message_text(msg: dict[str, Any]) -> str:
+    content = msg.get("content", "")
+    if isinstance(content, str):
+        return content
+
+    if isinstance(content, list):
+        parts = []
+        for item in content:
+            if isinstance(item, dict):
+                if item.get("type") == "text":
+                    parts.append(item.get("text", ""))
+                elif item.get("type") == "image_url":
+                    parts.append("[IMAGE]")
+        return " ".join(parts)
+
+    return str(content)
+
+
+def _summarize_messages(
+    messages: list[dict[str, Any]],
+    model: str,
+    timeout: int = 600,
+) -> dict[str, Any]:
+    if not messages:
+        empty_summary = "<context_summary message_count='0'>{text}</context_summary>"
+        return {
+            "role": "assistant",
+            "content": empty_summary.format(text="No messages to summarize"),
+        }
+
+    formatted = []
+    for msg in messages:
+        role = msg.get("role", "unknown")
+        text = _extract_message_text(msg)
+        formatted.append(f"{role}: {text}")
+
+    conversation = "\n".join(formatted)
+    prompt = SUMMARY_PROMPT_TEMPLATE.format(conversation=conversation)
+
+    try:
+        completion_args = {
+            "model": model,
+            "messages": [{"role": "user", "content": prompt}],
+            "timeout": timeout,
+        }
+
+        response = litellm.completion(**completion_args)
+        summary = response.choices[0].message.content or ""
+        if not summary.strip():
+            return messages[0]
+        summary_msg = "<context_summary message_count='{count}'>{text}</context_summary>"
+        return {
+            "role": "assistant",
+            "content": summary_msg.format(count=len(messages), text=summary),
+        }
+    except Exception:
+        logger.exception("Failed to summarize messages")
+        return messages[0]
+
+
+def _handle_images(messages: list[dict[str, Any]], max_images: int) -> None:
+    image_count = 0
+    for msg in reversed(messages):
+        content = msg.get("content", [])
+        if isinstance(content, list):
+            for item in content:
+                if isinstance(item, dict) and item.get("type") == "image_url":
+                    if image_count >= max_images:
+                        item.update(
+                            {
+                                "type": "text",
+                                "text": "[Previously attached image removed to preserve context]",
+                            }
+                        )
+                    else:
+                        image_count += 1
+
+
+class MemoryCompressor:
+    def __init__(
+        self,
+        max_images: int = 3,
+        model_name: str | None = None,
+        timeout: int = 600,
+    ):
+        self.max_images = max_images
+        self.model_name = model_name or os.getenv("STRIX_LLM", "openai/gpt-5")
+        self.timeout = timeout
+
+        if not self.model_name:
+            raise ValueError("STRIX_LLM environment variable must be set and not empty")
+
+    def compress_history(
+        self,
+        messages: list[dict[str, Any]],
+    ) -> list[dict[str, Any]]:
+        """Compress conversation history to stay within token limits.
+
+        Strategy:
+        1. Handle image limits first
+        2. Keep all system messages
+        3. Keep minimum recent messages
+        4. Summarize older messages when total tokens exceed limit
+
+        The compression preserves:
+        - All system messages unchanged
+        - Most recent messages intact
+        - Critical security context in summaries
+        - Recent images for visual context
+        - Technical details and findings
+        """
+        if not messages:
+            return messages
+
+        _handle_images(messages, self.max_images)
+
+        system_msgs = []
+        regular_msgs = []
+        for msg in messages:
+            if msg.get("role") == "system":
+                system_msgs.append(msg)
+            else:
+                regular_msgs.append(msg)
+
+        recent_msgs = regular_msgs[-MIN_RECENT_MESSAGES:]
+        old_msgs = regular_msgs[:-MIN_RECENT_MESSAGES]
+
+        # Type assertion since we ensure model_name is not None in __init__
+        model_name: str = self.model_name  # type: ignore[assignment]
+
+        total_tokens = sum(
+            _get_message_tokens(msg, model_name) for msg in system_msgs + regular_msgs
+        )
+
+        if total_tokens <= MAX_TOTAL_TOKENS * 0.9:
+            return messages
+
+        compressed = []
+        chunk_size = 10
+        for i in range(0, len(old_msgs), chunk_size):
+            chunk = old_msgs[i : i + chunk_size]
+            summary = _summarize_messages(chunk, model_name, self.timeout)
+            if summary:
+                compressed.append(summary)
+
+        return system_msgs + compressed + recent_msgs