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

vanna/integrations/local/agent_memory/in_memory.py

@@ -0,0 +1,285 @@
+"""
+Demo in-memory implementation of AgentMemory.
+
+This implementation provides a zero-dependency, minimal storage solution that
+keeps all memories in RAM. It uses simple similarity algorithms (Jaccard and
+difflib) instead of vector embeddings. Perfect for demos and testing.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import difflib
+import time
+import uuid
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from vanna.capabilities.agent_memory import (
+    AgentMemory,
+    TextMemory,
+    TextMemorySearchResult,
+    ToolMemory,
+    ToolMemorySearchResult,
+)
+from vanna.core.tool import ToolContext
+
+
+class DemoAgentMemory(AgentMemory):
+    """
+    Minimal, dependency-free in-memory storage for demos and testing.
+    - O(n) search over an in-memory list
+    - Simple similarity: max(Jaccard(token sets), difflib ratio)
+    - Optional FIFO eviction via max_items
+    - Async-safe with an asyncio.Lock
+    """
+
+    def __init__(self, *, max_items: int = 10_000):
+        """
+        Initialize the in-memory storage.
+
+        Args:
+            max_items: Maximum number of memories to keep. Oldest memories are
+                      evicted when this limit is reached (FIFO).
+        """
+        self._memories: List[ToolMemory] = []
+        self._text_memories: List[TextMemory] = []
+        self._lock = asyncio.Lock()
+        self._max_items = max_items
+
+    @staticmethod
+    def _now_iso() -> str:
+        """Get current timestamp in ISO format."""
+        return datetime.now().isoformat()
+
+    @staticmethod
+    def _normalize(text: str) -> str:
+        """Normalize text by lowercasing and collapsing whitespace."""
+        return " ".join(text.lower().split())
+
+    @staticmethod
+    def _tokenize(text: str) -> set[str]:
+        """Simple tokenizer that splits on whitespace."""
+        return set(text.lower().split())
+
+    @classmethod
+    def _similarity(cls, a: str, b: str) -> float:
+        """
+        Calculate similarity between two strings using multiple methods.
+
+        Returns the maximum of Jaccard similarity and difflib ratio.
+        """
+        a_norm, b_norm = cls._normalize(a), cls._normalize(b)
+
+        # Jaccard over whitespace tokens
+        ta, tb = cls._tokenize(a_norm), cls._tokenize(b_norm)
+        if not ta and not tb:
+            jaccard = 1.0
+        elif not ta or not tb:
+            jaccard = 0.0
+        else:
+            jaccard = len(ta & tb) / max(1, len(ta | tb))
+
+        # difflib ratio
+        ratio = difflib.SequenceMatcher(None, a_norm, b_norm).ratio()
+
+        # Take the better of the two cheap measures
+        return max(jaccard, ratio)
+
+    async def save_tool_usage(
+        self,
+        question: str,
+        tool_name: str,
+        args: Dict[str, Any],
+        context: ToolContext,
+        success: bool = True,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Save a tool usage pattern for future reference."""
+        tm = ToolMemory(
+            memory_id=str(uuid.uuid4()),
+            question=question,
+            tool_name=tool_name,
+            args=args,
+            timestamp=self._now_iso(),
+            success=success,
+            metadata=metadata or {},
+        )
+        async with self._lock:
+            self._memories.append(tm)
+            # Optional FIFO eviction
+            if len(self._memories) > self._max_items:
+                overflow = len(self._memories) - self._max_items
+                del self._memories[:overflow]
+
+    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
+        """Store a text memory in RAM."""
+        tm = TextMemory(
+            memory_id=str(uuid.uuid4()), content=content, timestamp=self._now_iso()
+        )
+        async with self._lock:
+            self._text_memories.append(tm)
+            if len(self._text_memories) > self._max_items:
+                overflow = len(self._text_memories) - self._max_items
+                del self._text_memories[:overflow]
+        return tm
+
+    async def search_similar_usage(
+        self,
+        question: str,
+        context: ToolContext,
+        *,
+        limit: int = 10,
+        similarity_threshold: float = 0.7,
+        tool_name_filter: Optional[str] = None,
+    ) -> List[ToolMemorySearchResult]:
+        """Search for similar tool usage patterns based on a question."""
+        q = self._normalize(question)
+
+        async with self._lock:
+            # Filter candidates by tool name and success status
+            candidates = [
+                m
+                for m in self._memories
+                if m.success
+                and (tool_name_filter is None or m.tool_name == tool_name_filter)
+            ]
+
+            # Score each candidate by question similarity
+            results: List[tuple[ToolMemory, float]] = []
+            for m in candidates:
+                score = self._similarity(q, m.question)
+                results.append((m, min(score, 1.0)))
+
+            # Filter by threshold and sort by score
+            results = [(m, s) for (m, s) in results if s >= similarity_threshold]
+            results.sort(key=lambda x: x[1], reverse=True)
+
+            # Build ranked response
+            out: List[ToolMemorySearchResult] = []
+            for idx, (m, s) in enumerate(results[:limit], start=1):
+                out.append(
+                    ToolMemorySearchResult(memory=m, similarity_score=s, rank=idx)
+                )
+            return out
+
+    async def search_text_memories(
+        self,
+        query: str,
+        context: ToolContext,
+        *,
+        limit: int = 10,
+        similarity_threshold: float = 0.7,
+    ) -> List[TextMemorySearchResult]:
+        """Search free-form text memories using the demo similarity metric."""
+        normalized_query = self._normalize(query)
+
+        async with self._lock:
+            scored: List[tuple[TextMemory, float]] = []
+            for memory in self._text_memories:
+                score = self._similarity(normalized_query, memory.content)
+                scored.append((memory, min(score, 1.0)))
+
+            scored = [
+                (memory, score)
+                for memory, score in scored
+                if score >= similarity_threshold
+            ]
+            scored.sort(key=lambda item: item[1], reverse=True)
+
+            results: List[TextMemorySearchResult] = []
+            for idx, (memory, score) in enumerate(scored[:limit], start=1):
+                results.append(
+                    TextMemorySearchResult(
+                        memory=memory, similarity_score=score, rank=idx
+                    )
+                )
+            return results
+
+    async def get_recent_memories(
+        self, context: ToolContext, limit: int = 10
+    ) -> List[ToolMemory]:
+        """Get recently added memories. Returns most recent memories first."""
+        async with self._lock:
+            # Return memories in reverse order (most recent first)
+            return list(reversed(self._memories[-limit:]))
+
+    async def get_recent_text_memories(
+        self, context: ToolContext, limit: int = 10
+    ) -> List[TextMemory]:
+        """Return recently added text memories."""
+        async with self._lock:
+            return list(reversed(self._text_memories[-limit:]))
+
+    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
+        """Delete a stored text memory by ID."""
+        async with self._lock:
+            for index, memory in enumerate(self._text_memories):
+                if memory.memory_id == memory_id:
+                    del self._text_memories[index]
+                    return True
+            return False
+
+    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
+        """Delete a memory by its ID. Returns True if deleted, False if not found."""
+        async with self._lock:
+            for i, m in enumerate(self._memories):
+                if m.memory_id == memory_id:
+                    del self._memories[i]
+                    return True
+            return False
+
+    async def clear_memories(
+        self,
+        context: ToolContext,
+        tool_name: Optional[str] = None,
+        before_date: Optional[str] = None,
+    ) -> int:
+        """Clear stored memories. Returns number of memories deleted."""
+        async with self._lock:
+            original_tool_count = len(self._memories)
+            original_text_count = len(self._text_memories)
+
+            # Filter memories to keep
+            kept_memories = []
+            for m in self._memories:
+                should_delete = True
+
+                # Check tool name filter
+                if tool_name and m.tool_name != tool_name:
+                    should_delete = False
+
+                # Check date filter
+                if should_delete and before_date and m.timestamp:
+                    if m.timestamp >= before_date:
+                        should_delete = False
+
+                # If no filters specified, delete all
+                if tool_name is None and before_date is None:
+                    should_delete = True
+
+                # Keep if should not delete
+                if not should_delete:
+                    kept_memories.append(m)
+
+            self._memories = kept_memories
+            deleted_tool_count = original_tool_count - len(self._memories)
+
+            # Apply filters to text memories (tool filter ignored)
+            kept_text_memories = []
+            for memory in self._text_memories:
+                should_delete = (
+                    tool_name is None
+                )  # only delete text when not targeting a tool
+
+                if before_date and memory.timestamp:
+                    if memory.timestamp >= before_date:
+                        should_delete = False
+
+                if not should_delete:
+                    kept_text_memories.append(memory)
+
+            self._text_memories = kept_text_memories
+            deleted_text_count = original_text_count - len(self._text_memories)
+
+            return deleted_tool_count + deleted_text_count

vanna/integrations/local/audit.py

@@ -0,0 +1,59 @@
+"""
+Local audit logger implementation using Python logging.
+
+This module provides a simple audit logger that writes events using
+the standard Python logging module, useful for development and testing.
+"""
+
+import json
+import logging
+from typing import Optional
+
+from vanna.core.audit import AuditEvent, AuditLogger
+
+logger = logging.getLogger(__name__)
+
+
+class LoggingAuditLogger(AuditLogger):
+    """Audit logger that writes events to Python logger as structured JSON.
+
+    This implementation uses logger.info() to emit audit events as JSON,
+    making them easy to parse and route to log aggregation systems.
+
+    Example:
+        audit_logger = LoggingAuditLogger()
+        agent = Agent(
+            llm_service=...,
+            audit_logger=audit_logger
+        )
+    """
+
+    def __init__(self, log_level: int = logging.INFO):
+        """Initialize the logging audit logger.
+
+        Args:
+            log_level: Log level to use for audit events (default: INFO)
+        """
+        self.log_level = log_level
+
+    async def log_event(self, event: AuditEvent) -> None:
+        """Log an audit event as structured JSON.
+
+        Args:
+            event: The audit event to log
+        """
+        try:
+            # Convert event to dict for JSON serialization
+            event_dict = event.model_dump(mode="json", exclude_none=True)
+
+            # Format as single-line JSON for easy parsing
+            event_json = json.dumps(event_dict, separators=(",", ":"))
+
+            # Log with structured prefix for easy filtering
+            logger.log(
+                self.log_level,
+                f"[AUDIT] {event.event_type.value} | {event_json}",
+            )
+        except Exception as e:
+            # Don't fail the operation if audit logging fails
+            logger.error(f"Failed to log audit event: {e}", exc_info=True)

vanna/integrations/local/file_system.py

@@ -0,0 +1,242 @@
+"""
+Local file system implementation.
+
+This module provides a local file system implementation with per-user isolation.
+"""
+
+import asyncio
+import hashlib
+from pathlib import Path
+from typing import List, Optional
+
+from vanna.capabilities.file_system import CommandResult, FileSearchMatch, FileSystem
+from vanna.core.tool import ToolContext
+
+MAX_SEARCH_FILE_BYTES = 1_000_000
+
+
+class LocalFileSystem(FileSystem):
+    """Local file system implementation with per-user isolation."""
+
+    def __init__(self, working_directory: str = "."):
+        """Initialize with a working directory.
+
+        Args:
+            working_directory: Base directory where user-specific folders will be created
+        """
+        self.working_directory = Path(working_directory)
+
+    def _get_user_directory(self, context: ToolContext) -> Path:
+        """Get the user-specific directory by hashing the user ID.
+
+        Args:
+            context: Tool context containing user information
+
+        Returns:
+            Path to the user-specific directory
+        """
+        # Hash the user ID to create a directory name
+        user_hash = hashlib.sha256(context.user.id.encode()).hexdigest()[:16]
+        user_dir = self.working_directory / user_hash
+
+        # Create the directory if it doesn't exist
+        user_dir.mkdir(parents=True, exist_ok=True)
+
+        return user_dir
+
+    def _resolve_path(self, path: str, context: ToolContext) -> Path:
+        """Resolve a path relative to the user's directory.
+
+        Args:
+            path: Path relative to user directory
+            context: Tool context containing user information
+
+        Returns:
+            Absolute path within user's directory
+        """
+        user_dir = self._get_user_directory(context)
+        resolved = user_dir / path
+
+        # Ensure the path is within the user's directory (prevent directory traversal)
+        try:
+            resolved.resolve().relative_to(user_dir.resolve())
+        except ValueError:
+            raise PermissionError(
+                f"Access denied: path '{path}' is outside user directory"
+            )
+
+        return resolved
+
+    async def list_files(self, directory: str, context: ToolContext) -> List[str]:
+        """List files in a directory within the user's isolated space."""
+        directory_path = self._resolve_path(directory, context)
+
+        if not directory_path.exists():
+            raise FileNotFoundError(f"Directory '{directory}' does not exist")
+
+        if not directory_path.is_dir():
+            raise NotADirectoryError(f"'{directory}' is not a directory")
+
+        files = []
+        for item in directory_path.iterdir():
+            if item.is_file():
+                files.append(item.name)
+
+        return sorted(files)
+
+    async def read_file(self, filename: str, context: ToolContext) -> str:
+        """Read the contents of a file within the user's isolated space."""
+        file_path = self._resolve_path(filename, context)
+
+        if not file_path.exists():
+            raise FileNotFoundError(f"File '{filename}' does not exist")
+
+        if not file_path.is_file():
+            raise IsADirectoryError(f"'{filename}' is a directory, not a file")
+
+        return file_path.read_text(encoding="utf-8")
+
+    async def write_file(
+        self, filename: str, content: str, context: ToolContext, overwrite: bool = False
+    ) -> None:
+        """Write content to a file within the user's isolated space."""
+        file_path = self._resolve_path(filename, context)
+
+        # Create parent directories if they don't exist
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        if file_path.exists() and not overwrite:
+            raise FileExistsError(
+                f"File '{filename}' already exists. Use overwrite=True to replace it."
+            )
+
+        file_path.write_text(content, encoding="utf-8")
+
+    async def exists(self, path: str, context: ToolContext) -> bool:
+        """Check if a file or directory exists within the user's isolated space."""
+        try:
+            resolved_path = self._resolve_path(path, context)
+            return resolved_path.exists()
+        except PermissionError:
+            return False
+
+    async def is_directory(self, path: str, context: ToolContext) -> bool:
+        """Check if a path is a directory within the user's isolated space."""
+        try:
+            resolved_path = self._resolve_path(path, context)
+            return resolved_path.exists() and resolved_path.is_dir()
+        except PermissionError:
+            return False
+
+    async def search_files(
+        self,
+        query: str,
+        context: ToolContext,
+        *,
+        max_results: int = 20,
+        include_content: bool = False,
+    ) -> List[FileSearchMatch]:
+        """Search for files within the user's isolated space."""
+
+        trimmed_query = query.strip()
+        if not trimmed_query:
+            raise ValueError("Search query must not be empty")
+
+        user_dir = self._get_user_directory(context)
+        matches: List[FileSearchMatch] = []
+        query_lower = trimmed_query.lower()
+
+        for path in user_dir.rglob("*"):
+            if len(matches) >= max_results:
+                break
+
+            if not path.is_file():
+                continue
+
+            relative_path = path.relative_to(user_dir).as_posix()
+            include_entry = False
+            snippet: Optional[str] = None
+
+            if query_lower in path.name.lower():
+                include_entry = True
+                snippet = "[filename match]"
+
+            content: Optional[str] = None
+            if include_content:
+                try:
+                    size = path.stat().st_size
+                except OSError:
+                    if include_entry:
+                        matches.append(
+                            FileSearchMatch(path=relative_path, snippet=snippet)
+                        )
+                    continue
+
+                if size <= MAX_SEARCH_FILE_BYTES:
+                    try:
+                        content = path.read_text(encoding="utf-8")
+                    except (UnicodeDecodeError, OSError):
+                        content = None
+                elif not include_entry:
+                    # Skip oversized files if they do not match by name
+                    continue
+
+            if include_content and content is not None:
+                if query_lower in content.lower():
+                    # Create snippet
+                    lowered = content.lower()
+                    index = lowered.find(query_lower)
+                    if index != -1:
+                        context_window = 60
+                        start = max(0, index - context_window)
+                        end = min(len(content), index + len(query) + context_window)
+                        snippet = content[start:end].replace("\n", " ").strip()
+                        if start > 0:
+                            snippet = f"…{snippet}"
+                        if end < len(content):
+                            snippet = f"{snippet}…"
+                    include_entry = True
+                elif not include_entry:
+                    continue
+
+            if include_entry:
+                matches.append(FileSearchMatch(path=relative_path, snippet=snippet))
+
+        return matches
+
+    async def run_bash(
+        self,
+        command: str,
+        context: ToolContext,
+        *,
+        timeout: Optional[float] = None,
+    ) -> CommandResult:
+        """Execute a bash command within the user's isolated space."""
+
+        if not command.strip():
+            raise ValueError("Command must not be empty")
+
+        user_dir = self._get_user_directory(context)
+
+        process = await asyncio.create_subprocess_shell(
+            command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=str(user_dir),
+        )
+
+        try:
+            stdout_bytes, stderr_bytes = await asyncio.wait_for(
+                process.communicate(), timeout=timeout
+            )
+        except asyncio.TimeoutError as exc:
+            process.kill()
+            await process.wait()
+            raise TimeoutError(f"Command timed out after {timeout} seconds") from exc
+
+        stdout = stdout_bytes.decode("utf-8", errors="replace")
+        stderr = stderr_bytes.decode("utf-8", errors="replace")
+
+        return CommandResult(
+            stdout=stdout, stderr=stderr, returncode=process.returncode or 0
+        )