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

vanna/integrations/azuresearch/agent_memory.py

@@ -0,0 +1,413 @@
+"""
+Azure AI Search implementation of AgentMemory.
+
+This implementation uses Azure Cognitive Search 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 azure.search.documents import SearchClient
+    from azure.search.documents.indexes import SearchIndexClient
+    from azure.search.documents.indexes.models import (
+        SearchIndex,
+        SearchField,
+        SearchFieldDataType,
+        VectorSearch,
+        VectorSearchAlgorithmConfiguration,
+    )
+    from azure.core.credentials import AzureKeyCredential
+
+    AZURE_SEARCH_AVAILABLE = True
+except ImportError:
+    AZURE_SEARCH_AVAILABLE = False
+
+from vanna.capabilities.agent_memory import (
+    AgentMemory,
+    TextMemory,
+    TextMemorySearchResult,
+    ToolMemory,
+    ToolMemorySearchResult,
+)
+from vanna.core.tool import ToolContext
+
+
+class AzureAISearchAgentMemory(AgentMemory):
+    """Azure AI Search-based implementation of AgentMemory."""
+
+    def __init__(
+        self,
+        endpoint: str,
+        api_key: str,
+        index_name: str = "tool-memories",
+        dimension: int = 384,
+    ):
+        if not AZURE_SEARCH_AVAILABLE:
+            raise ImportError(
+                "Azure Search is required for AzureAISearchAgentMemory. "
+                "Install with: pip install azure-search-documents"
+            )
+
+        self.endpoint = endpoint
+        self.api_key = api_key
+        self.index_name = index_name
+        self.dimension = dimension
+        self._credential = AzureKeyCredential(api_key)
+        self._search_client = None
+        self._index_client = None
+        self._executor = ThreadPoolExecutor(max_workers=2)
+
+    def _get_index_client(self):
+        """Get or create index client."""
+        if self._index_client is None:
+            self._index_client = SearchIndexClient(
+                endpoint=self.endpoint, credential=self._credential
+            )
+            self._ensure_index_exists()
+        return self._index_client
+
+    def _get_search_client(self):
+        """Get or create search client."""
+        if self._search_client is None:
+            self._get_index_client()  # Ensure index exists
+            self._search_client = SearchClient(
+                endpoint=self.endpoint,
+                index_name=self.index_name,
+                credential=self._credential,
+            )
+        return self._search_client
+
+    def _ensure_index_exists(self):
+        """Create index if it doesn't exist."""
+        try:
+            self._index_client.get_index(self.index_name)
+        except Exception:
+            # Create index with vector search configuration
+            fields = [
+                SearchField(
+                    name="memory_id", type=SearchFieldDataType.String, key=True
+                ),
+                SearchField(
+                    name="question", type=SearchFieldDataType.String, searchable=True
+                ),
+                SearchField(
+                    name="tool_name", type=SearchFieldDataType.String, filterable=True
+                ),
+                SearchField(name="args_json", type=SearchFieldDataType.String),
+                SearchField(
+                    name="timestamp",
+                    type=SearchFieldDataType.String,
+                    sortable=True,
+                    filterable=True,
+                ),
+                SearchField(
+                    name="success", type=SearchFieldDataType.Boolean, filterable=True
+                ),
+                SearchField(name="metadata_json", type=SearchFieldDataType.String),
+                SearchField(
+                    name="embedding",
+                    type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
+                    searchable=True,
+                    vector_search_dimensions=self.dimension,
+                    vector_search_configuration="vector-config",
+                ),
+            ]
+
+            vector_search = VectorSearch(
+                algorithm_configurations=[
+                    VectorSearchAlgorithmConfiguration(name="vector-config")
+                ]
+            )
+
+            index = SearchIndex(
+                name=self.index_name, fields=fields, vector_search=vector_search
+            )
+
+            self._index_client.create_index(index)
+
+    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_search_client()
+
+            memory_id = str(uuid.uuid4())
+            timestamp = datetime.now().isoformat()
+            embedding = self._create_embedding(question)
+
+            document = {
+                "memory_id": memory_id,
+                "question": question,
+                "tool_name": tool_name,
+                "args_json": json.dumps(args),
+                "timestamp": timestamp,
+                "success": success,
+                "metadata_json": json.dumps(metadata or {}),
+                "embedding": embedding,
+            }
+
+            client.upload_documents(documents=[document])
+
+        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_search_client()
+
+            embedding = self._create_embedding(question)
+
+            # Build filter
+            filter_expr = "success eq true"
+            if tool_name_filter:
+                filter_expr += f" and tool_name eq '{tool_name_filter}'"
+
+            results = client.search(
+                search_text=None, vector=embedding, top_k=limit, filter=filter_expr
+            )
+
+            search_results = []
+            for i, doc in enumerate(results):
+                # Azure returns similarity score in @search.score
+                similarity_score = doc.get("@search.score", 0)
+
+                if similarity_score >= similarity_threshold:
+                    args = json.loads(doc.get("args_json", "{}"))
+                    metadata_dict = json.loads(doc.get("metadata_json", "{}"))
+
+                    memory = ToolMemory(
+                        memory_id=doc["memory_id"],
+                        question=doc["question"],
+                        tool_name=doc["tool_name"],
+                        args=args,
+                        timestamp=doc.get("timestamp"),
+                        success=doc.get("success", True),
+                        metadata=metadata_dict,
+                    )
+
+                    search_results.append(
+                        ToolMemorySearchResult(
+                            memory=memory, similarity_score=similarity_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_search_client()
+
+            results = client.search(
+                search_text="*", top=limit, order_by=["timestamp desc"]
+            )
+
+            memories = []
+            for doc in results:
+                args = json.loads(doc.get("args_json", "{}"))
+                metadata_dict = json.loads(doc.get("metadata_json", "{}"))
+
+                memory = ToolMemory(
+                    memory_id=doc["memory_id"],
+                    question=doc["question"],
+                    tool_name=doc["tool_name"],
+                    args=args,
+                    timestamp=doc.get("timestamp"),
+                    success=doc.get("success", True),
+                    metadata=metadata_dict,
+                )
+                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."""
+
+        def _delete():
+            client = self._get_search_client()
+
+            try:
+                client.delete_documents(documents=[{"memory_id": memory_id}])
+                return True
+            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_search_client()
+
+            memory_id = str(uuid.uuid4())
+            timestamp = datetime.now().isoformat()
+            embedding = self._create_embedding(content)
+
+            document = {
+                "memory_id": memory_id,
+                "content": content,
+                "timestamp": timestamp,
+                "embedding": embedding,
+            }
+
+            client.upload_documents(documents=[document])
+
+            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_search_client()
+
+            embedding = self._create_embedding(query)
+
+            results = client.search(search_text=None, vector=embedding, top_k=limit)
+
+            search_results = []
+            for i, doc in enumerate(results):
+                similarity_score = doc.get("@search.score", 0)
+
+                if similarity_score >= similarity_threshold:
+                    memory = TextMemory(
+                        memory_id=doc["memory_id"],
+                        content=doc.get("content", ""),
+                        timestamp=doc.get("timestamp"),
+                    )
+
+                    search_results.append(
+                        TextMemorySearchResult(
+                            memory=memory, similarity_score=similarity_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_search_client()
+
+            results = client.search(
+                search_text="*", top=limit, order_by=["timestamp desc"]
+            )
+
+            memories = []
+            for doc in results:
+                # Skip if this is a tool memory (has tool_name field)
+                if "tool_name" in doc:
+                    continue
+
+                memory = TextMemory(
+                    memory_id=doc["memory_id"],
+                    content=doc.get("content", ""),
+                    timestamp=doc.get("timestamp"),
+                )
+                memories.append(memory)
+
+            return memories[:limit]
+
+        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_search_client()
+
+            try:
+                client.delete_documents(documents=[{"memory_id": memory_id}])
+                return True
+            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_search_client()
+
+            # Build filter
+            filter_parts = []
+            if tool_name:
+                filter_parts.append(f"tool_name eq '{tool_name}'")
+            if before_date:
+                filter_parts.append(f"timestamp lt '{before_date}'")
+
+            filter_expr = " and ".join(filter_parts) if filter_parts else None
+
+            # Search for documents to delete
+            results = client.search(
+                search_text="*", filter=filter_expr, select=["memory_id"]
+            )
+
+            docs_to_delete = [{"memory_id": doc["memory_id"]} for doc in results]
+
+            if docs_to_delete:
+                client.delete_documents(documents=docs_to_delete)
+
+            return len(docs_to_delete)
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)

vanna/integrations/bigquery/__init__.py

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

vanna/integrations/bigquery/sql_runner.py

@@ -0,0 +1,81 @@
+"""BigQuery 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 BigQueryRunner(SqlRunner):
+    """BigQuery implementation of the SqlRunner interface."""
+
+    def __init__(self, project_id: str, cred_file_path: Optional[str] = None, **kwargs):
+        """Initialize with BigQuery connection parameters.
+
+        Args:
+            project_id: Google Cloud Project ID
+            cred_file_path: Path to Google Cloud credentials JSON file (optional)
+            **kwargs: Additional google.cloud.bigquery.Client parameters
+        """
+        try:
+            from google.cloud import bigquery
+            from google.oauth2 import service_account
+
+            self.bigquery = bigquery
+            self.service_account = service_account
+        except ImportError as e:
+            raise ImportError(
+                "google-cloud-bigquery package is required. "
+                "Install with: pip install 'vanna[bigquery]'"
+            ) from e
+
+        self.project_id = project_id
+        self.cred_file_path = cred_file_path
+        self.kwargs = kwargs
+        self._client = None
+
+    def _get_client(self):
+        """Get or create BigQuery client."""
+        if self._client is not None:
+            return self._client
+
+        if self.cred_file_path:
+            import json
+
+            with open(self.cred_file_path, "r") as f:
+                credentials = (
+                    self.service_account.Credentials.from_service_account_info(
+                        json.loads(f.read()),
+                        scopes=["https://www.googleapis.com/auth/cloud-platform"],
+                    )
+                )
+            self._client = self.bigquery.Client(
+                project=self.project_id, credentials=credentials, **self.kwargs
+            )
+        else:
+            # Use default credentials
+            self._client = self.bigquery.Client(project=self.project_id, **self.kwargs)
+
+        return self._client
+
+    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
+        """Execute SQL query against BigQuery database and return results as DataFrame.
+
+        Args:
+            args: SQL query arguments
+            context: Tool execution context
+
+        Returns:
+            DataFrame with query results
+
+        Raises:
+            google.api_core.exceptions.GoogleAPIError: If query execution fails
+        """
+        client = self._get_client()
+
+        # Execute the query
+        job = client.query(args.sql)
+        df = job.result().to_dataframe()
+
+        return df

vanna/integrations/chromadb/__init__.py

@@ -0,0 +1,104 @@
+"""
+ChromaDB integration for Vanna Agents.
+"""
+
+from .agent_memory import ChromaAgentMemory
+
+
+def get_device() -> str:
+    """Detect the best available device for embeddings.
+
+    This function checks for GPU availability and returns the appropriate device string
+    for use with embedding models. It prioritizes hardware acceleration when available.
+
+    Returns:
+        str: Device string - 'cuda' if NVIDIA GPU available, 'mps' if Apple Silicon,
+             'cpu' otherwise.
+
+    Examples:
+        >>> device = get_device()
+        >>> print(f"Using device: {device}")
+        Using device: cuda
+
+        # Use with ChromaDB SentenceTransformer embeddings
+        >>> from chromadb.utils import embedding_functions
+        >>> ef = embedding_functions.SentenceTransformerEmbeddingFunction(
+        ...     model_name="sentence-transformers/all-MiniLM-L6-v2",
+        ...     device=get_device()
+        ... )
+        >>> memory = ChromaAgentMemory(embedding_function=ef)
+    """
+    try:
+        import torch
+
+        # Check for CUDA (NVIDIA GPUs)
+        if torch.cuda.is_available():
+            return "cuda"
+
+        # Check for MPS (Apple Silicon GPUs)
+        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            return "mps"
+
+    except ImportError:
+        # PyTorch not installed, fall back to CPU
+        pass
+
+    return "cpu"
+
+
+def create_sentence_transformer_embedding_function(
+    model_name: str = "sentence-transformers/all-MiniLM-L6-v2", device: str = None
+):
+    """Create a SentenceTransformer embedding function with automatic device detection.
+
+    This convenience function creates a ChromaDB-compatible SentenceTransformer embedding
+    function with intelligent device selection. If no device is specified, it automatically
+    detects and uses the best available hardware (CUDA, MPS, or CPU).
+
+    Note: This requires the 'sentence-transformers' package to be installed.
+    Install with: pip install sentence-transformers
+
+    Args:
+        model_name: The name of the sentence-transformer model to use.
+                   Defaults to "sentence-transformers/all-MiniLM-L6-v2".
+        device: Optional device string ('cuda', 'mps', or 'cpu'). If None,
+               automatically detects the best available device.
+
+    Returns:
+        A ChromaDB SentenceTransformer embedding function configured for the
+        specified/detected device.
+
+    Examples:
+        # Automatic device detection (uses CUDA/MPS if available)
+        >>> from vanna.integrations.chromadb import ChromaAgentMemory, create_sentence_transformer_embedding_function
+        >>> ef = create_sentence_transformer_embedding_function()
+        >>> memory = ChromaAgentMemory(embedding_function=ef)
+
+        # Explicitly use CUDA
+        >>> ef_cuda = create_sentence_transformer_embedding_function(device="cuda")
+        >>> memory = ChromaAgentMemory(embedding_function=ef_cuda)
+
+        # Use a different model
+        >>> ef_large = create_sentence_transformer_embedding_function(
+        ...     model_name="sentence-transformers/all-mpnet-base-v2"
+        ... )
+        >>> memory = ChromaAgentMemory(embedding_function=ef_large)
+    """
+    try:
+        from chromadb.utils import embedding_functions
+    except ImportError:
+        raise ImportError("ChromaDB is required. Install with: pip install chromadb")
+
+    if device is None:
+        device = get_device()
+
+    return embedding_functions.SentenceTransformerEmbeddingFunction(
+        model_name=model_name, device=device
+    )
+
+
+__all__ = [
+    "ChromaAgentMemory",
+    "get_device",
+    "create_sentence_transformer_embedding_function",
+]