vanna 0.7.9__py3-none-any.whl → 2.0.0__py3-none-any.whl

vanna/core/agent/config.py

@@ -0,0 +1,123 @@
+"""
+Agent configuration.
+
+This module contains configuration models that control agent behavior.
+"""
+
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+from .._compat import StrEnum
+
+if TYPE_CHECKING:
+    from ..user import User
+
+
+class UiFeature(StrEnum):
+    UI_FEATURE_SHOW_TOOL_NAMES = "tool_names"
+    UI_FEATURE_SHOW_TOOL_ARGUMENTS = "tool_arguments"
+    UI_FEATURE_SHOW_TOOL_ERROR = "tool_error"
+    UI_FEATURE_SHOW_TOOL_INVOCATION_MESSAGE_IN_CHAT = "tool_invocation_message_in_chat"
+    UI_FEATURE_SHOW_MEMORY_DETAILED_RESULTS = "memory_detailed_results"
+
+
+# Optional: you can also define defaults if you want a shared baseline
+DEFAULT_UI_FEATURES: Dict[str, List[str]] = {
+    UiFeature.UI_FEATURE_SHOW_TOOL_NAMES: ["admin", "user"],
+    UiFeature.UI_FEATURE_SHOW_TOOL_ARGUMENTS: ["admin"],
+    UiFeature.UI_FEATURE_SHOW_TOOL_ERROR: ["admin"],
+    UiFeature.UI_FEATURE_SHOW_TOOL_INVOCATION_MESSAGE_IN_CHAT: ["admin"],
+    UiFeature.UI_FEATURE_SHOW_MEMORY_DETAILED_RESULTS: ["admin"],
+}
+
+
+class UiFeatures(BaseModel):
+    """UI features with group-based access control using the same pattern as tools.
+
+    Each field specifies which groups can access that UI feature.
+    Empty list means the feature is accessible to all users.
+    Uses the same intersection logic as tool access control.
+    """
+
+    # Custom features for extensibility
+    feature_group_access: Dict[str, List[str]] = Field(
+        default_factory=lambda: DEFAULT_UI_FEATURES.copy(),
+        description="Which groups can access UI features",
+    )
+
+    def can_user_access_feature(self, feature_name: str, user: "User") -> bool:
+        """Check if user can access a UI feature using same logic as tools.
+
+        Args:
+            feature_name: Name of the UI feature to check
+            user: User object with group_memberships
+
+        Returns:
+            True if user has access, False otherwise
+        """
+        # Then try custom features
+        if feature_name in self.feature_group_access:
+            allowed_groups = self.feature_group_access[feature_name]
+        else:
+            # Feature doesn't exist, deny access
+            return False
+
+        # Empty list means all users can access (same as tools)
+        if not allowed_groups:
+            return True
+
+        # Same intersection logic as tool access control
+        user_groups = set(user.group_memberships)
+        feature_groups = set(allowed_groups)
+        return bool(user_groups & feature_groups)
+
+    def register_feature(self, name: str, access_groups: List[str]) -> None:
+        """Register a custom UI feature with group access control.
+
+        Args:
+            name: Name of the custom feature
+            access_groups: List of groups that can access this feature
+        """
+        self.feature_group_access[name] = access_groups
+
+
+class AuditConfig(BaseModel):
+    """Configuration for audit logging."""
+
+    enabled: bool = Field(default=True, description="Enable audit logging")
+    log_tool_access_checks: bool = Field(
+        default=True, description="Log tool access permission checks"
+    )
+    log_tool_invocations: bool = Field(
+        default=True, description="Log tool invocations with parameters"
+    )
+    log_tool_results: bool = Field(
+        default=True, description="Log tool execution results"
+    )
+    log_ui_feature_checks: bool = Field(
+        default=False, description="Log UI feature access checks (can be noisy)"
+    )
+    log_ai_responses: bool = Field(
+        default=True, description="Log AI-generated responses"
+    )
+    include_full_ai_responses: bool = Field(
+        default=False,
+        description="Include full AI response text in logs (privacy concern)",
+    )
+    sanitize_tool_parameters: bool = Field(
+        default=True, description="Sanitize sensitive parameters (passwords, tokens)"
+    )
+
+
+class AgentConfig(BaseModel):
+    """Configuration for agent behavior."""
+
+    max_tool_iterations: int = Field(default=10, gt=0)
+    stream_responses: bool = Field(default=True)
+    auto_save_conversations: bool = Field(default=True)
+    include_thinking_indicators: bool = Field(default=True)
+    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
+    max_tokens: Optional[int] = Field(default=None, gt=0)
+    ui_features: UiFeatures = Field(default_factory=UiFeatures)
+    audit_config: AuditConfig = Field(default_factory=AuditConfig)

vanna/core/audit/__init__.py

@@ -0,0 +1,28 @@
+"""
+Audit logging for the Vanna Agents framework.
+
+This module provides interfaces and models for audit logging, enabling
+tracking of user actions, tool invocations, and access control decisions.
+"""
+
+from .base import AuditLogger
+from .models import (
+    AiResponseEvent,
+    AuditEvent,
+    AuditEventType,
+    ToolAccessCheckEvent,
+    ToolInvocationEvent,
+    ToolResultEvent,
+    UiFeatureAccessCheckEvent,
+)
+
+__all__ = [
+    "AuditLogger",
+    "AuditEvent",
+    "AuditEventType",
+    "ToolAccessCheckEvent",
+    "ToolInvocationEvent",
+    "ToolResultEvent",
+    "UiFeatureAccessCheckEvent",
+    "AiResponseEvent",
+]

vanna/core/audit/base.py

@@ -0,0 +1,299 @@
+"""
+Base audit logger interface.
+
+Audit loggers enable tracking user actions, tool invocations, and access control
+decisions for security, compliance, and debugging.
+"""
+
+import hashlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from .models import (
+    AiResponseEvent,
+    AuditEvent,
+    ToolAccessCheckEvent,
+    ToolInvocationEvent,
+    ToolResultEvent,
+    UiFeatureAccessCheckEvent,
+)
+
+if TYPE_CHECKING:
+    from ..tool.models import ToolCall, ToolContext, ToolResult
+    from ..user.models import User
+
+
+class AuditLogger(ABC):
+    """Abstract base class for audit logging implementations.
+
+    Implementations can:
+    - Write to files (JSON, CSV, etc.)
+    - Send to databases (Postgres, MongoDB, etc.)
+    - Stream to cloud services (CloudWatch, Datadog, etc.)
+    - Send to SIEM systems (Splunk, Elastic, etc.)
+
+    Example:
+        class PostgresAuditLogger(AuditLogger):
+            async def log_event(self, event: AuditEvent) -> None:
+                await self.db.execute(
+                    "INSERT INTO audit_log (...) VALUES (...)",
+                    event.model_dump()
+                )
+
+        agent = Agent(
+            llm_service=...,
+            audit_logger=PostgresAuditLogger(db_pool)
+        )
+    """
+
+    @abstractmethod
+    async def log_event(self, event: AuditEvent) -> None:
+        """Log a single audit event.
+
+        Args:
+            event: The audit event to log
+
+        Raises:
+            Exception: If logging fails critically
+        """
+        pass
+
+    async def log_tool_access_check(
+        self,
+        user: "User",
+        tool_name: str,
+        access_granted: bool,
+        required_groups: List[str],
+        context: "ToolContext",
+        reason: Optional[str] = None,
+    ) -> None:
+        """Convenience method for logging tool access checks.
+
+        Args:
+            user: User attempting to access the tool
+            tool_name: Name of the tool being accessed
+            access_granted: Whether access was granted
+            required_groups: Groups required to access the tool
+            context: Tool execution context
+            reason: Optional reason for denial
+        """
+        event = ToolAccessCheckEvent(
+            user_id=user.id,
+            username=user.username,
+            user_email=user.email,
+            user_groups=user.group_memberships,
+            conversation_id=context.conversation_id,
+            request_id=context.request_id,
+            tool_name=tool_name,
+            access_granted=access_granted,
+            required_groups=required_groups,
+            reason=reason,
+        )
+        await self.log_event(event)
+
+    async def log_tool_invocation(
+        self,
+        user: "User",
+        tool_call: "ToolCall",
+        ui_features: List[str],
+        context: "ToolContext",
+        sanitize_parameters: bool = True,
+    ) -> None:
+        """Convenience method for logging tool invocations.
+
+        Args:
+            user: User invoking the tool
+            tool_call: Tool call information
+            ui_features: List of UI features available to the user
+            context: Tool execution context
+            sanitize_parameters: Whether to sanitize sensitive parameters
+        """
+        parameters = tool_call.arguments.copy()
+        sanitized = False
+
+        if sanitize_parameters:
+            parameters, sanitized = self._sanitize_parameters(parameters)
+
+        event = ToolInvocationEvent(
+            user_id=user.id,
+            username=user.username,
+            user_email=user.email,
+            user_groups=user.group_memberships,
+            conversation_id=context.conversation_id,
+            request_id=context.request_id,
+            tool_call_id=tool_call.id,
+            tool_name=tool_call.name,
+            parameters=parameters,
+            parameters_sanitized=sanitized,
+            ui_features_available=ui_features,
+        )
+        await self.log_event(event)
+
+    async def log_tool_result(
+        self,
+        user: "User",
+        tool_call: "ToolCall",
+        result: "ToolResult",
+        context: "ToolContext",
+    ) -> None:
+        """Convenience method for logging tool results.
+
+        Args:
+            user: User who invoked the tool
+            tool_call: Tool call information
+            result: Tool execution result
+            context: Tool execution context
+        """
+        event = ToolResultEvent(
+            user_id=user.id,
+            username=user.username,
+            user_email=user.email,
+            user_groups=user.group_memberships,
+            conversation_id=context.conversation_id,
+            request_id=context.request_id,
+            tool_call_id=tool_call.id,
+            tool_name=tool_call.name,
+            success=result.success,
+            error=result.error,
+            execution_time_ms=result.metadata.get("execution_time_ms", 0.0),
+            result_size_bytes=(
+                len(result.result_for_llm.encode("utf-8"))
+                if result.result_for_llm
+                else 0
+            ),
+            ui_component_type=(
+                result.ui_component.__class__.__name__ if result.ui_component else None
+            ),
+        )
+        await self.log_event(event)
+
+    async def log_ui_feature_access(
+        self,
+        user: "User",
+        feature_name: str,
+        access_granted: bool,
+        required_groups: List[str],
+        conversation_id: str,
+        request_id: str,
+    ) -> None:
+        """Convenience method for logging UI feature access checks.
+
+        Args:
+            user: User attempting to access the feature
+            feature_name: Name of the UI feature
+            access_granted: Whether access was granted
+            required_groups: Groups required to access the feature
+            conversation_id: Conversation identifier
+            request_id: Request identifier
+        """
+        event = UiFeatureAccessCheckEvent(
+            user_id=user.id,
+            username=user.username,
+            user_email=user.email,
+            user_groups=user.group_memberships,
+            conversation_id=conversation_id,
+            request_id=request_id,
+            feature_name=feature_name,
+            access_granted=access_granted,
+            required_groups=required_groups,
+        )
+        await self.log_event(event)
+
+    async def log_ai_response(
+        self,
+        user: "User",
+        conversation_id: str,
+        request_id: str,
+        response_text: str,
+        tool_calls: List["ToolCall"],
+        model_info: Optional[Dict[str, Any]] = None,
+        include_full_text: bool = False,
+    ) -> None:
+        """Convenience method for logging AI responses.
+
+        Args:
+            user: User receiving the response
+            conversation_id: Conversation identifier
+            request_id: Request identifier
+            response_text: The AI-generated response text
+            tool_calls: List of tool calls in the response
+            model_info: Optional model configuration info
+            include_full_text: Whether to include full response text
+        """
+        response_hash = hashlib.sha256(response_text.encode("utf-8")).hexdigest()
+
+        event = AiResponseEvent(
+            user_id=user.id,
+            username=user.username,
+            user_email=user.email,
+            user_groups=user.group_memberships,
+            conversation_id=conversation_id,
+            request_id=request_id,
+            response_length_chars=len(response_text),
+            response_text=response_text if include_full_text else None,
+            response_hash=response_hash,
+            model_name=model_info.get("model") if model_info else None,
+            temperature=model_info.get("temperature") if model_info else None,
+            tool_calls_count=len(tool_calls),
+            tool_names=[tc.name for tc in tool_calls],
+        )
+        await self.log_event(event)
+
+    async def query_events(
+        self,
+        filters: Optional[Dict[str, Any]] = None,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+        limit: int = 100,
+    ) -> List[AuditEvent]:
+        """Query audit events (optional, for implementations that support it).
+
+        Args:
+            filters: Filter criteria (user_id, event_type, etc.)
+            start_time: Filter events after this time
+            end_time: Filter events before this time
+            limit: Maximum number of events to return
+
+        Returns:
+            List of matching audit events
+
+        Raises:
+            NotImplementedError: If query not supported by implementation
+        """
+        raise NotImplementedError("Query not supported by this implementation")
+
+    def _sanitize_parameters(
+        self, parameters: Dict[str, Any]
+    ) -> tuple[Dict[str, Any], bool]:
+        """Sanitize sensitive data from parameters.
+
+        Args:
+            parameters: Raw parameters dict
+
+        Returns:
+            Tuple of (sanitized_parameters, was_sanitized)
+        """
+        sanitized = parameters.copy()
+        was_sanitized = False
+
+        # Common sensitive field patterns
+        sensitive_patterns = [
+            "password",
+            "secret",
+            "token",
+            "api_key",
+            "apikey",
+            "credential",
+            "auth",
+            "private_key",
+            "access_key",
+        ]
+
+        for key in list(sanitized.keys()):
+            key_lower = key.lower()
+            if any(pattern in key_lower for pattern in sensitive_patterns):
+                sanitized[key] = "[REDACTED]"
+                was_sanitized = True
+
+        return sanitized, was_sanitized

vanna/core/audit/models.py

@@ -0,0 +1,131 @@
+"""
+Audit event models.
+
+This module contains data models for audit logging events.
+"""
+
+import uuid
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+from .._compat import StrEnum
+
+
+class AuditEventType(StrEnum):
+    """Types of audit events."""
+
+    # Access control events
+    TOOL_ACCESS_CHECK = "tool_access_check"
+    UI_FEATURE_ACCESS_CHECK = "ui_feature_access_check"
+
+    # Tool execution events
+    TOOL_INVOCATION = "tool_invocation"
+    TOOL_RESULT = "tool_result"
+
+    # Conversation events
+    MESSAGE_RECEIVED = "message_received"
+    AI_RESPONSE_GENERATED = "ai_response_generated"
+    CONVERSATION_CREATED = "conversation_created"
+
+    # Security events
+    ACCESS_DENIED = "access_denied"
+    AUTHENTICATION_ATTEMPT = "authentication_attempt"
+
+
+class AuditEvent(BaseModel):
+    """Base audit event with common fields."""
+
+    event_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    event_type: AuditEventType
+    timestamp: datetime = Field(default_factory=datetime.utcnow)
+
+    # User context
+    user_id: str
+    username: Optional[str] = None
+    user_email: Optional[str] = None
+    user_groups: List[str] = Field(default_factory=list)
+
+    # Request context
+    conversation_id: str
+    request_id: str
+    remote_addr: Optional[str] = None
+
+    # Event-specific data
+    details: Dict[str, Any] = Field(default_factory=dict)
+
+    # Privacy/redaction markers
+    contains_pii: bool = False
+    redacted_fields: List[str] = Field(default_factory=list)
+
+
+class ToolAccessCheckEvent(AuditEvent):
+    """Audit event for tool access permission checks."""
+
+    event_type: AuditEventType = AuditEventType.TOOL_ACCESS_CHECK
+    tool_name: str
+    access_granted: bool
+    required_groups: List[str] = Field(default_factory=list)
+    reason: Optional[str] = None
+
+
+class ToolInvocationEvent(AuditEvent):
+    """Audit event for actual tool executions."""
+
+    event_type: AuditEventType = AuditEventType.TOOL_INVOCATION
+    tool_call_id: str
+    tool_name: str
+
+    # Parameters with sanitization support
+    parameters: Dict[str, Any] = Field(default_factory=dict)
+    parameters_sanitized: bool = False
+
+    # UI context at invocation time
+    ui_features_available: List[str] = Field(default_factory=list)
+
+
+class ToolResultEvent(AuditEvent):
+    """Audit event for tool execution results."""
+
+    event_type: AuditEventType = AuditEventType.TOOL_RESULT
+    tool_call_id: str
+    tool_name: str
+    success: bool
+    error: Optional[str] = None
+    execution_time_ms: float = 0.0
+
+    # Result metadata (without full content for size)
+    result_size_bytes: Optional[int] = None
+    ui_component_type: Optional[str] = None
+
+
+class UiFeatureAccessCheckEvent(AuditEvent):
+    """Audit event for UI feature access checks."""
+
+    event_type: AuditEventType = AuditEventType.UI_FEATURE_ACCESS_CHECK
+    feature_name: str
+    access_granted: bool
+    required_groups: List[str] = Field(default_factory=list)
+
+
+class AiResponseEvent(AuditEvent):
+    """Audit event for AI-generated responses."""
+
+    event_type: AuditEventType = AuditEventType.AI_RESPONSE_GENERATED
+
+    # Response metadata
+    response_length_chars: int
+    response_length_tokens: Optional[int] = None
+
+    # Full text (optional, configurable)
+    response_text: Optional[str] = None
+    response_hash: str  # SHA256 for integrity verification
+
+    # Model info
+    model_name: Optional[str] = None
+    temperature: Optional[float] = None
+
+    # Tool calls in response
+    tool_calls_count: int = 0
+    tool_names: List[str] = Field(default_factory=list)