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

vanna/integrations/postgres/sql_runner.py

@@ -0,0 +1,112 @@
+"""PostgreSQL implementation of SqlRunner interface."""
+
+from typing import Optional
+import pandas as pd
+
+from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
+from vanna.core.tool import ToolContext
+
+
+class PostgresRunner(SqlRunner):
+    """PostgreSQL implementation of the SqlRunner interface."""
+
+    def __init__(
+        self,
+        connection_string: Optional[str] = None,
+        host: Optional[str] = None,
+        port: Optional[int] = 5432,
+        database: Optional[str] = None,
+        user: Optional[str] = None,
+        password: Optional[str] = None,
+        **kwargs,
+    ):
+        """Initialize with PostgreSQL connection parameters.
+
+        You can either provide a connection_string OR individual parameters (host, database, etc.).
+        If connection_string is provided, it takes precedence.
+
+        Args:
+            connection_string: PostgreSQL connection string (e.g., "postgresql://user:password@host:port/database")
+            host: Database host address
+            port: Database port (default: 5432)
+            database: Database name
+            user: Database user
+            password: Database password
+            **kwargs: Additional psycopg2 connection parameters (sslmode, connect_timeout, etc.)
+        """
+        try:
+            import psycopg2
+            import psycopg2.extras
+
+            self.psycopg2 = psycopg2
+        except Exception as e:
+            raise ImportError(
+                "psycopg2 package is required. Install with: pip install 'vanna[postgres]'"
+            ) from e
+
+        if connection_string:
+            self.connection_string = connection_string
+            self.connection_params = None
+        elif host and database and user:
+            self.connection_string = None
+            self.connection_params = {
+                "host": host,
+                "port": port,
+                "database": database,
+                "user": user,
+                "password": password,
+                **kwargs,
+            }
+        else:
+            raise ValueError(
+                "Either provide connection_string OR (host, database, and user) parameters"
+            )
+
+    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
+        """Execute SQL query against PostgreSQL database and return results as DataFrame.
+
+        Args:
+            args: SQL query arguments
+            context: Tool execution context
+
+        Returns:
+            DataFrame with query results
+
+        Raises:
+            psycopg2.Error: If query execution fails
+        """
+        # Connect to the database using either connection string or parameters
+        if self.connection_string:
+            conn = self.psycopg2.connect(self.connection_string)
+        else:
+            conn = self.psycopg2.connect(**self.connection_params)
+
+        cursor = conn.cursor(cursor_factory=self.psycopg2.extras.RealDictCursor)
+
+        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/premium/agent_memory/__init__.py

@@ -0,0 +1,7 @@
+"""
+Cloud-based agent memory implementations.
+"""
+
+from .premium import CloudAgentMemory
+
+__all__ = ["CloudAgentMemory"]

vanna/integrations/premium/agent_memory/premium.py

@@ -0,0 +1,186 @@
+"""
+Cloud-based implementation of AgentMemory.
+
+This implementation uses Vanna's premium cloud service for storing and searching
+tool usage patterns with advanced similarity search and analytics.
+"""
+
+import json
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+import httpx
+
+from vanna.capabilities.agent_memory import (
+    AgentMemory,
+    TextMemory,
+    TextMemorySearchResult,
+    ToolMemory,
+    ToolMemorySearchResult,
+)
+from vanna.core.tool import ToolContext
+
+
+class CloudAgentMemory(AgentMemory):
+    """Cloud-based implementation of AgentMemory."""
+
+    def __init__(
+        self,
+        api_base_url: str = "https://api.vanna.ai",
+        api_key: Optional[str] = None,
+        organization_id: Optional[str] = None,
+    ):
+        self.api_base_url = api_base_url.rstrip("/")
+        self.api_key = api_key
+        self.organization_id = organization_id
+        self._client = httpx.AsyncClient(base_url=self.api_base_url, timeout=30.0)
+
+    def _get_headers(self) -> Dict[str, str]:
+        """Get request headers with authentication."""
+        headers = {"Content-Type": "application/json"}
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        if self.organization_id:
+            headers["X-Organization-ID"] = self.organization_id
+        return headers
+
+    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 to premium cloud storage."""
+        import uuid
+
+        payload = {
+            "id": str(uuid.uuid4()),
+            "question": question,
+            "tool_name": tool_name,
+            "args": args,
+            "success": success,
+            "metadata": metadata or {},
+            "timestamp": datetime.now().isoformat(),
+        }
+
+        response = await self._client.post(
+            "/memory/tool-usage", json=payload, headers=self._get_headers()
+        )
+        response.raise_for_status()
+
+    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 in premium cloud storage."""
+        params = {
+            "question": question,
+            "limit": limit,
+            "similarity_threshold": similarity_threshold,
+        }
+        if tool_name_filter:
+            params["tool_name_filter"] = tool_name_filter
+
+        response = await self._client.get(
+            "/memory/search-similar", params=params, headers=self._get_headers()
+        )
+        response.raise_for_status()
+
+        data = response.json()
+        results = []
+
+        for item in data.get("results", []):
+            memory = ToolMemory(**item["memory"])
+            result = ToolMemorySearchResult(
+                memory=memory,
+                similarity_score=item["similarity_score"],
+                rank=item["rank"],
+            )
+            results.append(result)
+
+        return results
+
+    async def get_recent_memories(
+        self, context: ToolContext, limit: int = 10
+    ) -> List[ToolMemory]:
+        """Get recently added memories from premium cloud storage."""
+        params = {"limit": limit}
+
+        response = await self._client.get(
+            "/memory/recent", params=params, headers=self._get_headers()
+        )
+        response.raise_for_status()
+
+        data = response.json()
+        memories = []
+
+        for item in data.get("memories", []):
+            memory = ToolMemory(**item)
+            memories.append(memory)
+
+        return memories
+
+    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
+        """Delete a memory by its ID from premium cloud storage."""
+        response = await self._client.delete(
+            f"/memory/{memory_id}", headers=self._get_headers()
+        )
+
+        if response.status_code == 404:
+            return False
+
+        response.raise_for_status()
+        return True
+
+    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
+        """Cloud implementation does not yet support text memories."""
+        raise NotImplementedError("CloudAgentMemory does not support text memories.")
+
+    async def search_text_memories(
+        self,
+        query: str,
+        context: ToolContext,
+        *,
+        limit: int = 10,
+        similarity_threshold: float = 0.7,
+    ) -> List[TextMemorySearchResult]:
+        """Cloud implementation does not yet support text memories."""
+        return []
+
+    async def get_recent_text_memories(
+        self, context: ToolContext, limit: int = 10
+    ) -> List[TextMemory]:
+        """Cloud implementation does not yet support text memories."""
+        return []
+
+    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
+        """Cloud implementation does not yet support text memories."""
+        return False
+
+    async def clear_memories(
+        self,
+        context: ToolContext,
+        tool_name: Optional[str] = None,
+        before_date: Optional[str] = None,
+    ) -> int:
+        """Clear stored memories from premium cloud storage."""
+        payload = {}
+        if tool_name:
+            payload["tool_name"] = tool_name
+        if before_date:
+            payload["before_date"] = before_date
+
+        response = await self._client.delete(
+            "/memory/clear", json=payload, headers=self._get_headers()
+        )
+        response.raise_for_status()
+
+        data = response.json()
+        return data.get("deleted_count", 0)

vanna/integrations/presto/__init__.py

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

vanna/integrations/presto/sql_runner.py

@@ -0,0 +1,107 @@
+"""Presto implementation of SqlRunner interface."""
+
+from typing import Optional
+import pandas as pd
+
+from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
+from vanna.core.tool import ToolContext
+
+
+class PrestoRunner(SqlRunner):
+    """Presto implementation of the SqlRunner interface."""
+
+    def __init__(
+        self,
+        host: str,
+        catalog: str = "hive",
+        schema: str = "default",
+        user: Optional[str] = None,
+        password: Optional[str] = None,
+        port: int = 443,
+        combined_pem_path: Optional[str] = None,
+        protocol: str = "https",
+        requests_kwargs: Optional[dict] = None,
+        **kwargs,
+    ):
+        """Initialize with Presto connection parameters.
+
+        Args:
+            host: The host address of the Presto database
+            catalog: The catalog to use in the Presto environment (default: 'hive')
+            schema: The schema to use in the Presto environment (default: 'default')
+            user: The username for authentication
+            password: The password for authentication
+            port: The port number for the Presto connection (default: 443)
+            combined_pem_path: The path to the combined pem file for SSL connection
+            protocol: The protocol to use for the connection (default: 'https')
+            requests_kwargs: Additional keyword arguments for requests
+            **kwargs: Additional pyhive connection parameters
+        """
+        try:
+            from pyhive import presto
+
+            self.presto = presto
+        except ImportError as e:
+            raise ImportError(
+                "pyhive package is required. Install with: pip install pyhive"
+            ) from e
+
+        self.host = host
+        self.catalog = catalog
+        self.schema = schema
+        self.user = user
+        self.password = password
+        self.port = port
+        self.protocol = protocol
+        self.kwargs = kwargs
+
+        # Set up requests_kwargs for SSL if combined_pem_path is provided
+        if requests_kwargs is None and combined_pem_path is not None:
+            self.requests_kwargs = {"verify": combined_pem_path}
+        else:
+            self.requests_kwargs = requests_kwargs
+
+    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
+        """Execute SQL query against Presto database and return results as DataFrame.
+
+        Args:
+            args: SQL query arguments
+            context: Tool execution context
+
+        Returns:
+            DataFrame with query results
+
+        Raises:
+            presto.Error: If query execution fails
+        """
+        # Connect to the database
+        conn = self.presto.Connection(
+            host=self.host,
+            username=self.user,
+            password=self.password,
+            catalog=self.catalog,
+            schema=self.schema,
+            port=self.port,
+            protocol=self.protocol,
+            requests_kwargs=self.requests_kwargs,
+            **self.kwargs,
+        )
+
+        try:
+            # Strip and remove trailing semicolons (Presto doesn't like them)
+            sql = args.sql.rstrip()
+            if sql.endswith(";"):
+                sql = sql[:-1]
+
+            cursor = conn.cursor()
+            cursor.execute(sql)
+            results = cursor.fetchall()
+
+            # Create a pandas dataframe from the results
+            df = pd.DataFrame(results, columns=[desc[0] for desc in cursor.description])
+
+            cursor.close()
+            return df
+
+        finally:
+            conn.close()

vanna/integrations/qdrant/__init__.py

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