aiptx 2.0.7__py3-none-any.whl

aipt_v2/llm/memory.py

@@ -0,0 +1,214 @@
+from __future__ import annotations
+
+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("AIPT_LLM", "openai/gpt-5")
+        self.timeout = timeout
+
+        if not self.model_name:
+            raise ValueError("AIPT_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

aipt_v2/llm/request_queue.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import threading
+import time
+from typing import Any
+
+import litellm
+from litellm import ModelResponse, completion
+from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential
+
+
+logger = logging.getLogger(__name__)
+
+
+def should_retry_exception(exception: Exception) -> bool:
+    status_code = None
+
+    if hasattr(exception, "status_code"):
+        status_code = exception.status_code
+    elif hasattr(exception, "response") and hasattr(exception.response, "status_code"):
+        status_code = exception.response.status_code
+
+    if status_code is not None:
+        return bool(litellm._should_retry(status_code))
+    return True
+
+
+class LLMRequestQueue:
+    def __init__(self, max_concurrent: int = 1, delay_between_requests: float = 4.0):
+        rate_limit_delay = os.getenv("LLM_RATE_LIMIT_DELAY")
+        if rate_limit_delay:
+            delay_between_requests = float(rate_limit_delay)
+
+        rate_limit_concurrent = os.getenv("LLM_RATE_LIMIT_CONCURRENT")
+        if rate_limit_concurrent:
+            max_concurrent = int(rate_limit_concurrent)
+
+        self.max_concurrent = max_concurrent
+        self.delay_between_requests = delay_between_requests
+        self._semaphore = threading.BoundedSemaphore(max_concurrent)
+        self._last_request_time = 0.0
+        self._lock = threading.Lock()
+
+    async def make_request(self, completion_args: dict[str, Any]) -> ModelResponse:
+        try:
+            while not self._semaphore.acquire(timeout=0.2):
+                await asyncio.sleep(0.1)
+
+            with self._lock:
+                now = time.time()
+                time_since_last = now - self._last_request_time
+                sleep_needed = max(0, self.delay_between_requests - time_since_last)
+                self._last_request_time = now + sleep_needed
+
+            if sleep_needed > 0:
+                await asyncio.sleep(sleep_needed)
+
+            return await self._reliable_request(completion_args)
+        finally:
+            self._semaphore.release()
+
+    @retry(  # type: ignore[misc]
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=8, min=8, max=64),
+        retry=retry_if_exception(should_retry_exception),
+        reraise=True,
+    )
+    async def _reliable_request(self, completion_args: dict[str, Any]) -> ModelResponse:
+        response = completion(**completion_args, stream=False)
+        if isinstance(response, ModelResponse):
+            return response
+        self._raise_unexpected_response()
+        raise RuntimeError("Unreachable code")
+
+    def _raise_unexpected_response(self) -> None:
+        raise RuntimeError("Unexpected response type")
+
+
+_global_queue: LLMRequestQueue | None = None
+
+
+def get_global_queue() -> LLMRequestQueue:
+    global _global_queue  # noqa: PLW0603
+    if _global_queue is None:
+        _global_queue = LLMRequestQueue()
+    return _global_queue

aipt_v2/llm/utils.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import html
+import re
+from typing import Any
+
+
+def _truncate_to_first_function(content: str) -> str:
+    if not content:
+        return content
+
+    function_starts = [match.start() for match in re.finditer(r"<function=", content)]
+
+    if len(function_starts) >= 2:
+        second_function_start = function_starts[1]
+
+        return content[:second_function_start].rstrip()
+
+    return content
+
+
+def parse_tool_invocations(content: str) -> list[dict[str, Any]] | None:
+    content = _fix_stopword(content)
+
+    tool_invocations: list[dict[str, Any]] = []
+
+    fn_regex_pattern = r"<function=([^>]+)>\n?(.*?)</function>"
+    fn_param_regex_pattern = r"<parameter=([^>]+)>(.*?)</parameter>"
+
+    fn_matches = re.finditer(fn_regex_pattern, content, re.DOTALL)
+
+    for fn_match in fn_matches:
+        fn_name = fn_match.group(1)
+        fn_body = fn_match.group(2)
+
+        param_matches = re.finditer(fn_param_regex_pattern, fn_body, re.DOTALL)
+
+        args = {}
+        for param_match in param_matches:
+            param_name = param_match.group(1)
+            param_value = param_match.group(2).strip()
+
+            param_value = html.unescape(param_value)
+            args[param_name] = param_value
+
+        tool_invocations.append({"toolName": fn_name, "args": args})
+
+    return tool_invocations if tool_invocations else None
+
+
+def _fix_stopword(content: str) -> str:
+    if "<function=" in content and content.count("<function=") == 1:
+        if content.endswith("</"):
+            content = content.rstrip() + "function>"
+        elif not content.rstrip().endswith("</function>"):
+            content = content + "\n</function>"
+    return content
+
+
+def format_tool_call(tool_name: str, args: dict[str, Any]) -> str:
+    xml_parts = [f"<function={tool_name}>"]
+
+    for key, value in args.items():
+        xml_parts.append(f"<parameter={key}>{value}</parameter>")
+
+    xml_parts.append("</function>")
+
+    return "\n".join(xml_parts)
+
+
+def clean_content(content: str) -> str:
+    if not content:
+        return ""
+
+    content = _fix_stopword(content)
+
+    tool_pattern = r"<function=[^>]+>.*?</function>"
+    cleaned = re.sub(tool_pattern, "", content, flags=re.DOTALL)
+
+    hidden_xml_patterns = [
+        r"<inter_agent_message>.*?</inter_agent_message>",
+        r"<agent_completion_report>.*?</agent_completion_report>",
+    ]
+    for pattern in hidden_xml_patterns:
+        cleaned = re.sub(pattern, "", cleaned, flags=re.DOTALL | re.IGNORECASE)
+
+    cleaned = re.sub(r"\n\s*\n", "\n\n", cleaned)
+
+    return cleaned.strip()