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

vanna/integrations/qdrant/agent_memory.py

@@ -0,0 +1,461 @@
+"""
+Qdrant vector database implementation of AgentMemory.
+
+This implementation uses Qdrant for vector storage of tool usage patterns.
+"""
+
+import json
+import uuid
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+
+try:
+    from qdrant_client import QdrantClient
+    from qdrant_client.models import (
+        Distance,
+        VectorParams,
+        PointStruct,
+        Filter,
+        FieldCondition,
+        MatchValue,
+    )
+
+    QDRANT_AVAILABLE = True
+except ImportError:
+    QDRANT_AVAILABLE = False
+
+from vanna.capabilities.agent_memory import (
+    AgentMemory,
+    TextMemory,
+    TextMemorySearchResult,
+    ToolMemory,
+    ToolMemorySearchResult,
+)
+from vanna.core.tool import ToolContext
+
+
+class QdrantAgentMemory(AgentMemory):
+    """Qdrant-based implementation of AgentMemory."""
+
+    def __init__(
+        self,
+        collection_name: str = "tool_memories",
+        url: Optional[str] = None,
+        path: Optional[str] = None,
+        api_key: Optional[str] = None,
+        dimension: int = 384,
+    ):
+        if not QDRANT_AVAILABLE:
+            raise ImportError(
+                "Qdrant is required for QdrantAgentMemory. Install with: pip install qdrant-client"
+            )
+
+        self.collection_name = collection_name
+        self.url = url
+        self.path = path
+        self.api_key = api_key
+        self.dimension = dimension
+        self._client = None
+        self._executor = ThreadPoolExecutor(max_workers=2)
+
+    def _get_client(self):
+        """Get or create Qdrant client."""
+        if self._client is None:
+            if self.url:
+                self._client = QdrantClient(url=self.url, api_key=self.api_key)
+            else:
+                self._client = QdrantClient(path=self.path or ":memory:")
+
+            # Create collection if it doesn't exist
+            collections = self._client.get_collections().collections
+            if not any(c.name == self.collection_name for c in collections):
+                self._client.create_collection(
+                    collection_name=self.collection_name,
+                    vectors_config=VectorParams(
+                        size=self.dimension, distance=Distance.COSINE
+                    ),
+                )
+        return self._client
+
+    def _create_embedding(self, text: str) -> List[float]:
+        """Create a simple embedding from text (placeholder)."""
+        import hashlib
+
+        hash_val = int(hashlib.md5(text.encode()).hexdigest(), 16)
+        return [(hash_val >> i) % 100 / 100.0 for i in range(self.dimension)]
+
+    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."""
+
+        def _save():
+            client = self._get_client()
+
+            memory_id = str(uuid.uuid4())
+            timestamp = datetime.now().isoformat()
+            embedding = self._create_embedding(question)
+
+            payload = {
+                "question": question,
+                "tool_name": tool_name,
+                "args": args,
+                "timestamp": timestamp,
+                "success": success,
+                "metadata": metadata or {},
+            }
+
+            point = PointStruct(id=memory_id, vector=embedding, payload=payload)
+
+            client.upsert(collection_name=self.collection_name, points=[point])
+
+        await asyncio.get_event_loop().run_in_executor(self._executor, _save)
+
+    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."""
+
+        def _search():
+            client = self._get_client()
+
+            embedding = self._create_embedding(question)
+
+            # Build filter
+            query_filter = None
+            conditions = [FieldCondition(key="success", match=MatchValue(value=True))]
+            if tool_name_filter:
+                conditions.append(
+                    FieldCondition(
+                        key="tool_name", match=MatchValue(value=tool_name_filter)
+                    )
+                )
+
+            if conditions:
+                query_filter = Filter(must=conditions)
+
+            # Use query_points for newer qdrant-client (1.8.0+) or search for older versions
+            if hasattr(client, "query_points"):
+                results = client.query_points(
+                    collection_name=self.collection_name,
+                    query=embedding,
+                    limit=limit,
+                    query_filter=query_filter,
+                    score_threshold=similarity_threshold,
+                ).points
+            else:
+                # Fallback to search method for older qdrant-client versions
+                results = client.search(
+                    collection_name=self.collection_name,
+                    query_vector=embedding,
+                    limit=limit,
+                    query_filter=query_filter,
+                    score_threshold=similarity_threshold,
+                )
+
+            search_results = []
+            for i, hit in enumerate(results):
+                payload = hit.payload
+
+                memory = ToolMemory(
+                    memory_id=str(hit.id),
+                    question=payload["question"],
+                    tool_name=payload["tool_name"],
+                    args=payload["args"],
+                    timestamp=payload.get("timestamp"),
+                    success=payload.get("success", True),
+                    metadata=payload.get("metadata", {}),
+                )
+
+                search_results.append(
+                    ToolMemorySearchResult(
+                        memory=memory, similarity_score=hit.score, rank=i + 1
+                    )
+                )
+
+            return search_results
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)
+
+    async def get_recent_memories(
+        self, context: ToolContext, limit: int = 10
+    ) -> List[ToolMemory]:
+        """Get recently added memories."""
+
+        def _get_recent():
+            client = self._get_client()
+
+            # Scroll through all points and sort by timestamp
+            points, _ = client.scroll(
+                collection_name=self.collection_name,
+                limit=1000,  # Get more than we need to sort
+                with_payload=True,
+                with_vectors=False,
+            )
+
+            # Sort by timestamp
+            sorted_points = sorted(
+                points, key=lambda p: p.payload.get("timestamp", ""), reverse=True
+            )
+
+            memories = []
+            for point in sorted_points[:limit]:
+                payload = point.payload
+                memory = ToolMemory(
+                    memory_id=str(point.id),
+                    question=payload["question"],
+                    tool_name=payload["tool_name"],
+                    args=payload["args"],
+                    timestamp=payload.get("timestamp"),
+                    success=payload.get("success", True),
+                    metadata=payload.get("metadata", {}),
+                )
+                memories.append(memory)
+
+            return memories
+
+        return await asyncio.get_event_loop().run_in_executor(
+            self._executor, _get_recent
+        )
+
+    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."""
+
+        def _delete():
+            client = self._get_client()
+
+            try:
+                # Check if the point exists before attempting to delete
+                points = client.retrieve(
+                    collection_name=self.collection_name,
+                    ids=[memory_id],
+                    with_payload=False,
+                    with_vectors=False,
+                )
+
+                if points and len(points) > 0:
+                    client.delete(
+                        collection_name=self.collection_name,
+                        points_selector=[memory_id],
+                    )
+                    return True
+                return False
+            except Exception:
+                return False
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)
+
+    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
+        """Save a text memory."""
+
+        def _save():
+            client = self._get_client()
+
+            memory_id = str(uuid.uuid4())
+            timestamp = datetime.now().isoformat()
+            embedding = self._create_embedding(content)
+
+            payload = {
+                "content": content,
+                "timestamp": timestamp,
+                "is_text_memory": True,
+            }
+
+            point = PointStruct(id=memory_id, vector=embedding, payload=payload)
+
+            client.upsert(collection_name=self.collection_name, points=[point])
+
+            return TextMemory(memory_id=memory_id, content=content, timestamp=timestamp)
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _save)
+
+    async def search_text_memories(
+        self,
+        query: str,
+        context: ToolContext,
+        *,
+        limit: int = 10,
+        similarity_threshold: float = 0.7,
+    ) -> List[TextMemorySearchResult]:
+        """Search for similar text memories."""
+
+        def _search():
+            client = self._get_client()
+
+            embedding = self._create_embedding(query)
+
+            query_filter = Filter(
+                must=[
+                    FieldCondition(key="is_text_memory", match=MatchValue(value=True))
+                ]
+            )
+
+            # Use query_points for newer qdrant-client (1.8.0+) or search for older versions
+            if hasattr(client, "query_points"):
+                results = client.query_points(
+                    collection_name=self.collection_name,
+                    query=embedding,
+                    limit=limit,
+                    query_filter=query_filter,
+                    score_threshold=similarity_threshold,
+                ).points
+            else:
+                # Fallback to search method for older qdrant-client versions
+                results = client.search(
+                    collection_name=self.collection_name,
+                    query_vector=embedding,
+                    limit=limit,
+                    query_filter=query_filter,
+                    score_threshold=similarity_threshold,
+                )
+
+            search_results = []
+            for i, hit in enumerate(results):
+                payload = hit.payload
+
+                memory = TextMemory(
+                    memory_id=str(hit.id),
+                    content=payload.get("content", ""),
+                    timestamp=payload.get("timestamp"),
+                )
+
+                search_results.append(
+                    TextMemorySearchResult(
+                        memory=memory, similarity_score=hit.score, rank=i + 1
+                    )
+                )
+
+            return search_results
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)
+
+    async def get_recent_text_memories(
+        self, context: ToolContext, limit: int = 10
+    ) -> List[TextMemory]:
+        """Get recently added text memories."""
+
+        def _get_recent():
+            client = self._get_client()
+
+            # Scroll through text memory points and sort by timestamp
+            points, _ = client.scroll(
+                collection_name=self.collection_name,
+                scroll_filter=Filter(
+                    must=[
+                        FieldCondition(
+                            key="is_text_memory", match=MatchValue(value=True)
+                        )
+                    ]
+                ),
+                limit=1000,
+                with_payload=True,
+                with_vectors=False,
+            )
+
+            # Sort by timestamp
+            sorted_points = sorted(
+                points, key=lambda p: p.payload.get("timestamp", ""), reverse=True
+            )
+
+            memories = []
+            for point in sorted_points[:limit]:
+                payload = point.payload
+                memory = TextMemory(
+                    memory_id=str(point.id),
+                    content=payload.get("content", ""),
+                    timestamp=payload.get("timestamp"),
+                )
+                memories.append(memory)
+
+            return memories
+
+        return await asyncio.get_event_loop().run_in_executor(
+            self._executor, _get_recent
+        )
+
+    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
+        """Delete a text memory by its ID."""
+
+        def _delete():
+            client = self._get_client()
+
+            try:
+                # Check if the point exists before attempting to delete
+                points = client.retrieve(
+                    collection_name=self.collection_name,
+                    ids=[memory_id],
+                    with_payload=False,
+                    with_vectors=False,
+                )
+
+                if points and len(points) > 0:
+                    client.delete(
+                        collection_name=self.collection_name,
+                        points_selector=[memory_id],
+                    )
+                    return True
+                return False
+            except Exception:
+                return False
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)
+
+    async def clear_memories(
+        self,
+        context: ToolContext,
+        tool_name: Optional[str] = None,
+        before_date: Optional[str] = None,
+    ) -> int:
+        """Clear stored memories."""
+
+        def _clear():
+            client = self._get_client()
+
+            # Build filter
+            conditions = []
+            if tool_name:
+                conditions.append(
+                    FieldCondition(key="tool_name", match=MatchValue(value=tool_name))
+                )
+            if before_date:
+                conditions.append(
+                    FieldCondition(key="timestamp", match=MatchValue(value=before_date))
+                )
+
+            if conditions or (tool_name is None and before_date is None):
+                # Delete with filter or delete all
+                query_filter = Filter(must=conditions) if conditions else None
+
+                if query_filter:
+                    client.delete(
+                        collection_name=self.collection_name,
+                        points_selector=query_filter,
+                    )
+                else:
+                    # Delete all points
+                    client.delete_collection(collection_name=self.collection_name)
+                    # Recreate empty collection
+                    client.create_collection(
+                        collection_name=self.collection_name,
+                        vectors_config=VectorParams(
+                            size=self.dimension, distance=Distance.COSINE
+                        ),
+                    )
+
+            return 0  # Qdrant doesn't return count
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)

vanna/integrations/snowflake/__init__.py

@@ -0,0 +1,5 @@
+"""Snowflake integration for Vanna."""
+
+from .sql_runner import SnowflakeRunner
+
+__all__ = ["SnowflakeRunner"]

vanna/integrations/snowflake/sql_runner.py

@@ -0,0 +1,147 @@
+"""Snowflake implementation of SqlRunner interface."""
+
+from typing import Optional, Union
+import os
+import pandas as pd
+
+from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
+from vanna.core.tool import ToolContext
+
+
+class SnowflakeRunner(SqlRunner):
+    """Snowflake implementation of the SqlRunner interface."""
+
+    def __init__(
+        self,
+        account: str,
+        username: str,
+        password: Optional[str] = None,
+        database: str = "",
+        role: Optional[str] = None,
+        warehouse: Optional[str] = None,
+        private_key_path: Optional[str] = None,
+        private_key_passphrase: Optional[str] = None,
+        private_key_content: Optional[bytes] = None,
+        **kwargs,
+    ):
+        """Initialize with Snowflake connection parameters.
+
+        Args:
+            account: Snowflake account identifier
+            username: Database user
+            password: Database password (optional if using key-pair auth)
+            database: Database name
+            role: Snowflake role to use (optional)
+            warehouse: Snowflake warehouse to use (optional)
+            private_key_path: Path to private key file for RSA key-pair authentication (optional)
+            private_key_passphrase: Passphrase for encrypted private key (optional)
+            private_key_content: Private key content as bytes (optional, alternative to private_key_path)
+            **kwargs: Additional snowflake.connector connection parameters
+
+        Note:
+            Either password OR private_key_path/private_key_content must be provided.
+            RSA key-pair authentication is recommended for production systems as Snowflake
+            is deprecating user/password authentication.
+        """
+        try:
+            import snowflake.connector
+
+            self.snowflake = snowflake.connector
+        except ImportError as e:
+            raise ImportError(
+                "snowflake-connector-python package is required. "
+                "Install with: pip install 'vanna[snowflake]'"
+            ) from e
+
+        # Validate that at least one authentication method is provided
+        if not password and not private_key_path and not private_key_content:
+            raise ValueError(
+                "Either password or private_key_path/private_key_content must be provided for authentication"
+            )
+
+        # Validate private key path exists if provided
+        if private_key_path and not os.path.isfile(private_key_path):
+            raise FileNotFoundError(f"Private key file not found: {private_key_path}")
+
+        self.account = account
+        self.username = username
+        self.password = password
+        self.database = database
+        self.role = role
+        self.warehouse = warehouse
+        self.private_key_path = private_key_path
+        self.private_key_passphrase = private_key_passphrase
+        self.private_key_content = private_key_content
+        self.kwargs = kwargs
+
+    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
+        """Execute SQL query against Snowflake database and return results as DataFrame.
+
+        Args:
+            args: SQL query arguments
+            context: Tool execution context
+
+        Returns:
+            DataFrame with query results
+
+        Raises:
+            snowflake.connector.Error: If query execution fails
+        """
+        # Build connection parameters
+        conn_params = {
+            "user": self.username,
+            "account": self.account,
+            "client_session_keep_alive": True,
+        }
+
+        # Add database if specified
+        if self.database:
+            conn_params["database"] = self.database
+
+        # Configure authentication method
+        if self.private_key_path or self.private_key_content:
+            # Use RSA key-pair authentication
+            if self.private_key_path:
+                conn_params["private_key_path"] = self.private_key_path
+            else:
+                conn_params["private_key_content"] = self.private_key_content
+
+            # Add passphrase if provided
+            if self.private_key_passphrase:
+                conn_params["private_key_passphrase"] = self.private_key_passphrase
+        else:
+            # Use password authentication (fallback)
+            conn_params["password"] = self.password
+
+        # Add any additional kwargs
+        conn_params.update(self.kwargs)
+
+        # Connect to the database
+        conn = self.snowflake.connect(**conn_params)
+
+        cursor = conn.cursor()
+
+        try:
+            # Set role if specified
+            if self.role:
+                cursor.execute(f"USE ROLE {self.role}")
+
+            # Set warehouse if specified
+            if self.warehouse:
+                cursor.execute(f"USE WAREHOUSE {self.warehouse}")
+
+            # Use the specified database if provided
+            if self.database:
+                cursor.execute(f"USE DATABASE {self.database}")
+
+            # Execute the query
+            cursor.execute(args.sql)
+            results = cursor.fetchall()
+
+            # Create a pandas dataframe from the results
+            df = pd.DataFrame(results, columns=[desc[0] for desc in cursor.description])
+            return df
+
+        finally:
+            cursor.close()
+            conn.close()

vanna/integrations/sqlite/__init__.py

@@ -0,0 +1,9 @@
+"""
+SQLite integration.
+
+This module provides SQLite runner implementation.
+"""
+
+from .sql_runner import SqliteRunner
+
+__all__ = ["SqliteRunner"]

vanna/integrations/sqlite/sql_runner.py

@@ -0,0 +1,65 @@
+"""SQLite implementation of SqlRunner interface."""
+
+import sqlite3
+import pandas as pd
+
+from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
+from vanna.core.tool import ToolContext
+
+
+class SqliteRunner(SqlRunner):
+    """SQLite implementation of the SqlRunner interface."""
+
+    def __init__(self, database_path: str):
+        """Initialize with a SQLite database path.
+
+        Args:
+            database_path: Path to the SQLite database file
+        """
+        self.database_path = database_path
+
+    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
+        """Execute SQL query against SQLite database and return results as DataFrame.
+
+        Args:
+            args: SQL query arguments
+            context: Tool execution context
+
+        Returns:
+            DataFrame with query results
+
+        Raises:
+            sqlite3.Error: If query execution fails
+        """
+        # Connect to the database
+        conn = sqlite3.connect(self.database_path)
+        conn.row_factory = sqlite3.Row  # Enable column access by name
+        cursor = conn.cursor()
+
+        try:
+            # Execute the query
+            cursor.execute(args.sql)
+
+            # Determine if this is a SELECT query or modification query
+            query_type = args.sql.strip().upper().split()[0]
+
+            if query_type == "SELECT":
+                # Fetch results for SELECT queries
+                rows = cursor.fetchall()
+                if not rows:
+                    # Return empty DataFrame
+                    return pd.DataFrame()
+
+                # Convert rows to list of dictionaries
+                results_data = [dict(row) for row in rows]
+                return pd.DataFrame(results_data)
+            else:
+                # For non-SELECT queries (INSERT, UPDATE, DELETE, etc.)
+                conn.commit()
+                rows_affected = cursor.rowcount
+                # Return a DataFrame indicating rows affected
+                return pd.DataFrame({"rows_affected": [rows_affected]})
+
+        finally:
+            cursor.close()
+            conn.close()

vanna/integrations/weaviate/__init__.py

@@ -0,0 +1,7 @@
+"""
+Weaviate integration for Vanna Agents.
+"""
+
+from .agent_memory import WeaviateAgentMemory
+
+__all__ = ["WeaviateAgentMemory"]