aiptx 2.0.2__py3-none-any.whl

aipt_v2/tools/tool_processing.py

@@ -0,0 +1,108 @@
+"""
+AIPT Tool Processing - Execute tool invocations from LLM responses
+"""
+
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+async def process_tool_invocations(
+    actions: list[dict[str, Any]],
+    conversation_history: list[dict[str, Any]],
+    state: Any,
+) -> bool:
+    """
+    Process tool invocations from LLM response.
+
+    Args:
+        actions: List of tool invocation dicts with 'name' and 'arguments'
+        conversation_history: Mutable conversation history
+        state: Agent state object
+
+    Returns:
+        True if agent should finish, False otherwise
+    """
+    for action in actions:
+        tool_name = action.get("name", "")
+        tool_args = action.get("arguments", {})
+
+        logger.info(f"Executing tool: {tool_name}")
+
+        # Check for finish tools
+        if tool_name in ["finish_scan", "agent_finish"]:
+            result = tool_args.get("result", "Task completed")
+            conversation_history.append({
+                "role": "user",
+                "content": f"Tool {tool_name} executed. Result: {result}",
+            })
+            return True
+
+        # Execute the tool
+        try:
+            result = await _execute_tool(tool_name, tool_args, state)
+            conversation_history.append({
+                "role": "user",
+                "content": f"Tool {tool_name} result:\n{result}",
+            })
+        except Exception as e:
+            error_msg = f"Tool {tool_name} failed: {str(e)}"
+            logger.error(error_msg)
+            conversation_history.append({
+                "role": "user",
+                "content": error_msg,
+            })
+
+    return False
+
+
+async def _execute_tool(name: str, args: dict[str, Any], state: Any) -> str:
+    """Execute a single tool and return result"""
+    # Import tool executors lazily
+    if name == "execute_command":
+        return await _execute_command(args, state)
+    elif name == "browser_navigate":
+        return await _browser_navigate(args, state)
+    elif name == "browser_screenshot":
+        return await _browser_screenshot(args, state)
+    else:
+        return f"Tool '{name}' executed with args: {args}"
+
+
+async def _execute_command(args: dict[str, Any], state: Any) -> str:
+    """Execute a shell command in the sandbox"""
+    import asyncio
+
+    command = args.get("command", "")
+    timeout = args.get("timeout", 60)
+
+    # Use subprocess for now (Docker integration later)
+    try:
+        proc = await asyncio.create_subprocess_shell(
+            command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
+        output = stdout.decode() if stdout else ""
+        errors = stderr.decode() if stderr else ""
+        return output + errors if errors else output
+    except asyncio.TimeoutError:
+        return f"Command timed out after {timeout} seconds"
+    except Exception as e:
+        return f"Command failed: {str(e)}"
+
+
+async def _browser_navigate(args: dict[str, Any], state: Any) -> str:
+    """Navigate browser to URL"""
+    url = args.get("url", "")
+    return f"Navigated to: {url}"
+
+
+async def _browser_screenshot(args: dict[str, Any], state: Any) -> str:
+    """Take browser screenshot"""
+    return "Screenshot taken"
+
+
+__all__ = ["process_tool_invocations"]

aipt_v2/utils/__init__.py

@@ -0,0 +1,17 @@
+"""
+AIPT v2 Utilities Module
+========================
+
+Provides common utilities used across the framework:
+- Structured logging with secret redaction
+- Model management wrappers
+- Searcher utilities
+"""
+
+from .logging import logger, setup_logging, get_logger
+
+__all__ = [
+    "logger",
+    "setup_logging",
+    "get_logger",
+]

aipt_v2/utils/logging.py

@@ -0,0 +1,201 @@
+"""
+Structured Logging Configuration for AIPT v2
+=============================================
+
+Provides:
+- Structured logging via structlog
+- Automatic secret redaction
+- JSON format for production
+- Console format for development
+"""
+
+import logging
+import os
+import sys
+import re
+from typing import Any, Optional
+from functools import lru_cache
+
+try:
+    import structlog
+    STRUCTLOG_AVAILABLE = True
+except ImportError:
+    STRUCTLOG_AVAILABLE = False
+
+
+# Patterns for secret redaction
+SECRET_PATTERNS = [
+    r"api[_-]?key",
+    r"apikey",
+    r"token",
+    r"secret",
+    r"password",
+    r"credential",
+    r"auth",
+    r"bearer",
+    r"sk-[a-zA-Z0-9]+",
+    r"pk-[a-zA-Z0-9]+",
+    r"access[_-]?key",
+    r"private[_-]?key",
+]
+
+SECRET_REGEX = re.compile("|".join(SECRET_PATTERNS), re.IGNORECASE)
+
+
+def _should_redact(key: str) -> bool:
+    """Check if a key should be redacted."""
+    return bool(SECRET_REGEX.search(key))
+
+
+def _redact_value(value: str) -> str:
+    """Redact a sensitive value, keeping first/last chars for debugging."""
+    if len(value) <= 8:
+        return "[REDACTED]"
+    return f"{value[:4]}...{value[-4:]}"
+
+
+def _redact_processor(logger: Any, method_name: str, event_dict: dict) -> dict:
+    """Structlog processor to redact sensitive information."""
+    for key, value in list(event_dict.items()):
+        if isinstance(value, str):
+            if _should_redact(key):
+                event_dict[key] = "[REDACTED]"
+            elif len(value) > 20 and SECRET_REGEX.search(value):
+                event_dict[key] = _redact_value(value)
+    return event_dict
+
+
+def setup_logging(
+    level: str = "INFO",
+    json_format: bool = False,
+    redact_secrets: bool = True,
+) -> Any:
+    """
+    Configure structured logging for the application.
+
+    Args:
+        level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
+        json_format: Use JSON format (for production)
+        redact_secrets: Automatically redact sensitive values
+
+    Returns:
+        Configured logger instance
+    """
+    log_level = getattr(logging, level.upper(), logging.INFO)
+
+    if STRUCTLOG_AVAILABLE:
+        processors = [
+            structlog.stdlib.filter_by_level,
+            structlog.stdlib.add_logger_name,
+            structlog.stdlib.add_log_level,
+            structlog.processors.TimeStamper(fmt="iso"),
+            structlog.processors.StackInfoRenderer(),
+        ]
+
+        if redact_secrets:
+            processors.append(_redact_processor)
+
+        if json_format:
+            processors.append(structlog.processors.JSONRenderer())
+        else:
+            processors.append(structlog.dev.ConsoleRenderer(colors=True))
+
+        structlog.configure(
+            processors=processors,
+            wrapper_class=structlog.stdlib.BoundLogger,
+            context_class=dict,
+            logger_factory=structlog.stdlib.LoggerFactory(),
+            cache_logger_on_first_use=True,
+        )
+
+        logging.basicConfig(
+            format="%(message)s",
+            stream=sys.stdout,
+            level=log_level,
+        )
+
+        return structlog.get_logger()
+    else:
+        # Fallback to standard logging
+        logging.basicConfig(
+            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+            stream=sys.stdout,
+            level=log_level,
+        )
+        return logging.getLogger("aipt_v2")
+
+
+@lru_cache(maxsize=1)
+def get_logger() -> Any:
+    """Get or create the global logger instance."""
+    log_level = os.getenv("AIPT_LOG_LEVEL", "INFO")
+    json_format = os.getenv("AIPT_LOG_FORMAT", "console").lower() == "json"
+    redact = os.getenv("AIPT_REDACT_SECRETS", "true").lower() == "true"
+
+    return setup_logging(
+        level=log_level,
+        json_format=json_format,
+        redact_secrets=redact,
+    )
+
+
+# Global logger instance
+logger = get_logger()
+
+
+class LoggerAdapter:
+    """
+    Adapter for consistent logging interface.
+
+    Provides methods that work whether structlog is available or not.
+    """
+
+    def __init__(self, logger_instance: Any):
+        self._logger = logger_instance
+        self._is_structlog = STRUCTLOG_AVAILABLE
+
+    def _log(self, level: str, msg: str, **kwargs):
+        """Internal log method."""
+        if self._is_structlog:
+            getattr(self._logger, level)(msg, **kwargs)
+        else:
+            extra = " ".join(f"{k}={v}" for k, v in kwargs.items())
+            full_msg = f"{msg} {extra}" if extra else msg
+            getattr(self._logger, level)(full_msg)
+
+    def debug(self, msg: str, **kwargs):
+        self._log("debug", msg, **kwargs)
+
+    def info(self, msg: str, **kwargs):
+        self._log("info", msg, **kwargs)
+
+    def warning(self, msg: str, **kwargs):
+        self._log("warning", msg, **kwargs)
+
+    def error(self, msg: str, exc_info: bool = False, **kwargs):
+        if self._is_structlog:
+            if exc_info:
+                kwargs["exc_info"] = True
+            self._logger.error(msg, **kwargs)
+        else:
+            extra = " ".join(f"{k}={v}" for k, v in kwargs.items())
+            full_msg = f"{msg} {extra}" if extra else msg
+            self._logger.error(full_msg, exc_info=exc_info)
+
+    def critical(self, msg: str, **kwargs):
+        self._log("critical", msg, **kwargs)
+
+    def exception(self, msg: str, **kwargs):
+        """Log exception with traceback."""
+        self.error(msg, exc_info=True, **kwargs)
+
+
+# Export a wrapped logger for consistent interface
+def create_logger(name: Optional[str] = None) -> LoggerAdapter:
+    """Create a named logger instance."""
+    if STRUCTLOG_AVAILABLE:
+        base_logger = structlog.get_logger(name) if name else get_logger()
+    else:
+        base_logger = logging.getLogger(name or "aipt_v2")
+
+    return LoggerAdapter(base_logger)

aipt_v2/utils/model_manager.py

@@ -0,0 +1,187 @@
+"""
+Model Manager for AIPT v2
+=========================
+
+Provides a unified interface for LLM model access.
+This is a compatibility layer for intelligence modules that
+reference utils.model_manager.
+"""
+
+import os
+from typing import Any, Optional, Dict, List
+from dataclasses import dataclass
+
+from aipt_v2.utils.logging import logger
+
+
+@dataclass
+class ModelConfig:
+    """Configuration for model instances."""
+    model_name: str = "gpt-4"
+    temperature: float = 0.7
+    max_tokens: int = 4096
+    timeout: int = 120
+
+
+class ModelWrapper:
+    """
+    Wrapper around litellm for consistent model access.
+
+    Provides both sync and async completion methods.
+    """
+
+    def __init__(self, config: ModelConfig):
+        self.config = config
+        self._litellm = None
+
+    def _get_litellm(self):
+        """Lazy load litellm."""
+        if self._litellm is None:
+            try:
+                import litellm
+                self._litellm = litellm
+            except ImportError:
+                raise ImportError(
+                    "litellm is required for model_manager. "
+                    "Install with: pip install litellm"
+                )
+        return self._litellm
+
+    def complete(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ) -> str:
+        """
+        Synchronous completion.
+
+        Args:
+            prompt: User prompt
+            system_prompt: Optional system prompt
+            **kwargs: Additional arguments passed to litellm
+
+        Returns:
+            Model response text
+        """
+        litellm = self._get_litellm()
+
+        messages: List[Dict[str, str]] = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+
+        try:
+            response = litellm.completion(
+                model=self.config.model_name,
+                messages=messages,
+                temperature=kwargs.get("temperature", self.config.temperature),
+                max_tokens=kwargs.get("max_tokens", self.config.max_tokens),
+                timeout=kwargs.get("timeout", self.config.timeout),
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            logger.error("Model completion failed", model=self.config.model_name, error=str(e))
+            raise
+
+    async def acomplete(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ) -> str:
+        """
+        Asynchronous completion.
+
+        Args:
+            prompt: User prompt
+            system_prompt: Optional system prompt
+            **kwargs: Additional arguments passed to litellm
+
+        Returns:
+            Model response text
+        """
+        litellm = self._get_litellm()
+
+        messages: List[Dict[str, str]] = []
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+
+        try:
+            response = await litellm.acompletion(
+                model=self.config.model_name,
+                messages=messages,
+                temperature=kwargs.get("temperature", self.config.temperature),
+                max_tokens=kwargs.get("max_tokens", self.config.max_tokens),
+                timeout=kwargs.get("timeout", self.config.timeout),
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            logger.error("Async model completion failed", model=self.config.model_name, error=str(e))
+            raise
+
+    def embed(self, text: str) -> List[float]:
+        """
+        Get embedding for text.
+
+        Args:
+            text: Text to embed
+
+        Returns:
+            Embedding vector
+        """
+        litellm = self._get_litellm()
+
+        try:
+            response = litellm.embedding(
+                model="text-embedding-ada-002",
+                input=text,
+            )
+            return response.data[0]["embedding"]
+        except Exception as e:
+            logger.error("Embedding failed", error=str(e))
+            raise
+
+
+# Cache for model instances
+_model_cache: Dict[str, ModelWrapper] = {}
+
+
+def get_model(
+    model_name: Optional[str] = None,
+    **kwargs
+) -> ModelWrapper:
+    """
+    Get or create a model instance.
+
+    Args:
+        model_name: Model identifier (default from env or gpt-4)
+        **kwargs: Additional config options
+
+    Returns:
+        ModelWrapper instance
+    """
+    if model_name is None:
+        model_name = os.getenv("AIPT_LLM_MODEL", "gpt-4")
+
+    cache_key = f"{model_name}:{hash(frozenset(kwargs.items()))}"
+
+    if cache_key not in _model_cache:
+        config = ModelConfig(
+            model_name=model_name,
+            temperature=kwargs.get("temperature", 0.7),
+            max_tokens=kwargs.get("max_tokens", 4096),
+            timeout=kwargs.get("timeout", 120),
+        )
+        _model_cache[cache_key] = ModelWrapper(config)
+        logger.info("Created model instance", model=model_name)
+
+    return _model_cache[cache_key]
+
+
+def clear_model_cache():
+    """Clear the model cache."""
+    global _model_cache
+    _model_cache = {}
+    logger.info("Model cache cleared")