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

vanna/integrations/local/file_system_conversation_store.py

@@ -0,0 +1,255 @@
+"""
+File system conversation store implementation.
+
+This module provides a file-based implementation of the ConversationStore
+interface that persists conversations to disk as a directory structure.
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Dict, List, Optional
+from datetime import datetime
+import time
+
+from vanna.core.storage import ConversationStore, Conversation, Message
+from vanna.core.user import User
+
+
+class FileSystemConversationStore(ConversationStore):
+    """File system-based conversation store.
+
+    Stores conversations as directories with individual message files:
+    conversations/{conversation_id}/
+        metadata.json - conversation metadata (id, user info, timestamps)
+        messages/
+            {timestamp}_{index}.json - individual message files
+    """
+
+    def __init__(self, base_dir: str = "conversations") -> None:
+        """Initialize the file system conversation store.
+
+        Args:
+            base_dir: Base directory for storing conversations
+        """
+        self.base_dir = Path(base_dir)
+        self.base_dir.mkdir(parents=True, exist_ok=True)
+
+    def _get_conversation_dir(self, conversation_id: str) -> Path:
+        """Get the directory path for a conversation."""
+        return self.base_dir / conversation_id
+
+    def _get_metadata_path(self, conversation_id: str) -> Path:
+        """Get the metadata file path for a conversation."""
+        return self._get_conversation_dir(conversation_id) / "metadata.json"
+
+    def _get_messages_dir(self, conversation_id: str) -> Path:
+        """Get the messages directory for a conversation."""
+        return self._get_conversation_dir(conversation_id) / "messages"
+
+    def _save_metadata(self, conversation: Conversation) -> None:
+        """Save conversation metadata to disk."""
+        conv_dir = self._get_conversation_dir(conversation.id)
+        conv_dir.mkdir(parents=True, exist_ok=True)
+
+        metadata = {
+            "id": conversation.id,
+            "user": conversation.user.model_dump(mode="json"),
+            "created_at": conversation.created_at.isoformat(),
+            "updated_at": conversation.updated_at.isoformat(),
+        }
+
+        metadata_path = self._get_metadata_path(conversation.id)
+        with open(metadata_path, "w") as f:
+            json.dump(metadata, f, indent=2)
+
+    def _load_messages(self, conversation_id: str) -> List[Message]:
+        """Load all messages for a conversation."""
+        messages_dir = self._get_messages_dir(conversation_id)
+
+        if not messages_dir.exists():
+            return []
+
+        messages = []
+        # Sort message files by name (timestamp_index ensures correct order)
+        message_files = sorted(messages_dir.glob("*.json"))
+
+        for file_path in message_files:
+            try:
+                with open(file_path, "r") as f:
+                    data = json.load(f)
+                message = Message.model_validate(data)
+                messages.append(message)
+            except (json.JSONDecodeError, ValueError) as e:
+                print(f"Failed to load message from {file_path}: {e}")
+                continue
+
+        return messages
+
+    def _append_message(
+        self, conversation_id: str, message: Message, index: int
+    ) -> None:
+        """Append a message to the conversation."""
+        messages_dir = self._get_messages_dir(conversation_id)
+        messages_dir.mkdir(parents=True, exist_ok=True)
+
+        # Use timestamp + index to ensure unique, ordered filenames
+        timestamp = int(time.time() * 1000000)  # microseconds
+        filename = f"{timestamp}_{index:06d}.json"
+        file_path = messages_dir / filename
+
+        with open(file_path, "w") as f:
+            json.dump(message.model_dump(mode="json"), f, indent=2)
+
+    async def create_conversation(
+        self, conversation_id: str, user: User, initial_message: str
+    ) -> Conversation:
+        """Create a new conversation with the specified ID."""
+        conversation = Conversation(
+            id=conversation_id,
+            user=user,
+            messages=[Message(role="user", content=initial_message)],
+        )
+
+        # Save metadata
+        self._save_metadata(conversation)
+
+        # Save initial message
+        self._append_message(conversation_id, conversation.messages[0], 0)
+
+        return conversation
+
+    async def get_conversation(
+        self, conversation_id: str, user: User
+    ) -> Optional[Conversation]:
+        """Get conversation by ID, scoped to user."""
+        metadata_path = self._get_metadata_path(conversation_id)
+
+        if not metadata_path.exists():
+            return None
+
+        try:
+            # Load metadata
+            with open(metadata_path, "r") as f:
+                metadata = json.load(f)
+
+            # Verify ownership
+            if metadata["user"]["id"] != user.id:
+                return None
+
+            # Load all messages
+            messages = self._load_messages(conversation_id)
+
+            # Reconstruct conversation
+            conversation = Conversation(
+                id=metadata["id"],
+                user=User.model_validate(metadata["user"]),
+                messages=messages,
+                created_at=datetime.fromisoformat(metadata["created_at"]),
+                updated_at=datetime.fromisoformat(metadata["updated_at"]),
+            )
+
+            return conversation
+        except (json.JSONDecodeError, ValueError, KeyError) as e:
+            print(f"Failed to load conversation {conversation_id}: {e}")
+            return None
+
+    async def update_conversation(self, conversation: Conversation) -> None:
+        """Update conversation with new messages."""
+        # Update the updated_at timestamp
+        conversation.updated_at = datetime.now()
+
+        # Save updated metadata
+        self._save_metadata(conversation)
+
+        # Get existing messages count to determine new message indices
+        existing_messages = self._load_messages(conversation.id)
+        existing_count = len(existing_messages)
+
+        # Only append new messages (ones not already saved)
+        for i, message in enumerate(
+            conversation.messages[existing_count:], start=existing_count
+        ):
+            self._append_message(conversation.id, message, i)
+
+    async def delete_conversation(self, conversation_id: str, user: User) -> bool:
+        """Delete conversation."""
+        conv_dir = self._get_conversation_dir(conversation_id)
+
+        if not conv_dir.exists():
+            return False
+
+        # Verify ownership before deleting
+        conversation = await self.get_conversation(conversation_id, user)
+        if not conversation:
+            return False
+
+        try:
+            # Delete all message files
+            messages_dir = self._get_messages_dir(conversation_id)
+            if messages_dir.exists():
+                for file_path in messages_dir.glob("*.json"):
+                    file_path.unlink()
+                messages_dir.rmdir()
+
+            # Delete metadata
+            metadata_path = self._get_metadata_path(conversation_id)
+            if metadata_path.exists():
+                metadata_path.unlink()
+
+            # Delete conversation directory
+            conv_dir.rmdir()
+
+            return True
+        except OSError as e:
+            print(f"Failed to delete conversation {conversation_id}: {e}")
+            return False
+
+    async def list_conversations(
+        self, user: User, limit: int = 50, offset: int = 0
+    ) -> List[Conversation]:
+        """List conversations for user."""
+        if not self.base_dir.exists():
+            return []
+
+        conversations = []
+
+        # Iterate through all conversation directories
+        for conv_dir in self.base_dir.iterdir():
+            if not conv_dir.is_dir():
+                continue
+
+            metadata_path = conv_dir / "metadata.json"
+            if not metadata_path.exists():
+                continue
+
+            try:
+                # Load metadata
+                with open(metadata_path, "r") as f:
+                    metadata = json.load(f)
+
+                # Skip conversations not owned by this user
+                if metadata["user"]["id"] != user.id:
+                    continue
+
+                # Load messages
+                messages = self._load_messages(conv_dir.name)
+
+                # Reconstruct conversation
+                conversation = Conversation(
+                    id=metadata["id"],
+                    user=User.model_validate(metadata["user"]),
+                    messages=messages,
+                    created_at=datetime.fromisoformat(metadata["created_at"]),
+                    updated_at=datetime.fromisoformat(metadata["updated_at"]),
+                )
+                conversations.append(conversation)
+            except (json.JSONDecodeError, ValueError, KeyError) as e:
+                print(f"Failed to load conversation from {conv_dir}: {e}")
+                continue
+
+        # Sort by updated_at desc
+        conversations.sort(key=lambda x: x.updated_at, reverse=True)
+
+        # Apply pagination
+        return conversations[offset : offset + limit]

vanna/integrations/local/storage.py

@@ -0,0 +1,62 @@
+"""
+In-memory conversation store implementation.
+
+This module provides a simple in-memory implementation of the ConversationStore
+interface, useful for testing and development.
+"""
+
+from typing import Dict, List, Optional
+
+from vanna.core.storage import ConversationStore, Conversation, Message
+from vanna.core.user import User
+
+
+class MemoryConversationStore(ConversationStore):
+    """In-memory conversation store."""
+
+    def __init__(self) -> None:
+        self._conversations: Dict[str, Conversation] = {}
+
+    async def create_conversation(
+        self, conversation_id: str, user: User, initial_message: str
+    ) -> Conversation:
+        """Create a new conversation with the specified ID."""
+        conversation = Conversation(
+            id=conversation_id,
+            user=user,
+            messages=[Message(role="user", content=initial_message)],
+        )
+        self._conversations[conversation_id] = conversation
+        return conversation
+
+    async def get_conversation(
+        self, conversation_id: str, user: User
+    ) -> Optional[Conversation]:
+        """Get conversation by ID, scoped to user."""
+        conversation = self._conversations.get(conversation_id)
+        if conversation and conversation.user.id == user.id:
+            return conversation
+        return None
+
+    async def update_conversation(self, conversation: Conversation) -> None:
+        """Update conversation with new messages."""
+        self._conversations[conversation.id] = conversation
+
+    async def delete_conversation(self, conversation_id: str, user: User) -> bool:
+        """Delete conversation."""
+        conversation = await self.get_conversation(conversation_id, user)
+        if conversation:
+            del self._conversations[conversation_id]
+            return True
+        return False
+
+    async def list_conversations(
+        self, user: User, limit: int = 50, offset: int = 0
+    ) -> List[Conversation]:
+        """List conversations for user."""
+        user_conversations = [
+            conv for conv in self._conversations.values() if conv.user.id == user.id
+        ]
+        # Sort by updated_at desc
+        user_conversations.sort(key=lambda x: x.updated_at, reverse=True)
+        return user_conversations[offset : offset + limit]

vanna/integrations/marqo/__init__.py

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

vanna/integrations/marqo/agent_memory.py

@@ -0,0 +1,354 @@
+"""
+Marqo vector database implementation of AgentMemory.
+
+This implementation uses Marqo 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:
+    import marqo
+
+    MARQO_AVAILABLE = True
+except ImportError:
+    MARQO_AVAILABLE = False
+
+from vanna.capabilities.agent_memory import (
+    AgentMemory,
+    TextMemory,
+    TextMemorySearchResult,
+    ToolMemory,
+    ToolMemorySearchResult,
+)
+from vanna.core.tool import ToolContext
+
+
+class MarqoAgentMemory(AgentMemory):
+    """Marqo-based implementation of AgentMemory."""
+
+    def __init__(
+        self,
+        url: str = "http://localhost:8882",
+        index_name: str = "tool-memories",
+        api_key: Optional[str] = None,
+    ):
+        if not MARQO_AVAILABLE:
+            raise ImportError(
+                "Marqo is required for MarqoAgentMemory. Install with: pip install marqo"
+            )
+
+        self.url = url
+        self.index_name = index_name
+        self.api_key = api_key
+        self._client = None
+        self._executor = ThreadPoolExecutor(max_workers=2)
+
+    def _get_client(self):
+        """Get or create Marqo client."""
+        if self._client is None:
+            self._client = marqo.Client(url=self.url, api_key=self.api_key)
+
+            # Create index if it doesn't exist
+            try:
+                self._client.get_index(self.index_name)
+            except Exception:
+                self._client.create_index(self.index_name)
+
+        return self._client
+
+    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()
+
+            document = {
+                "_id": memory_id,
+                "question": question,
+                "tool_name": tool_name,
+                "args": json.dumps(args),
+                "timestamp": timestamp,
+                "success": success,
+                "metadata": json.dumps(metadata or {}),
+            }
+
+            client.index(self.index_name).add_documents(
+                [document], tensor_fields=["question"]
+            )
+
+        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()
+
+            # Build filter
+            filter_string = "success:true"
+            if tool_name_filter:
+                filter_string += f" AND tool_name:{tool_name_filter}"
+
+            results = client.index(self.index_name).search(
+                q=question, limit=limit, filter_string=filter_string
+            )
+
+            search_results = []
+            for i, hit in enumerate(results["hits"]):
+                # Marqo returns score
+                similarity_score = hit.get("_score", 0)
+
+                if similarity_score >= similarity_threshold:
+                    args = json.loads(hit.get("args", "{}"))
+                    metadata_dict = json.loads(hit.get("metadata", "{}"))
+
+                    memory = ToolMemory(
+                        memory_id=hit["_id"],
+                        question=hit["question"],
+                        tool_name=hit["tool_name"],
+                        args=args,
+                        timestamp=hit.get("timestamp"),
+                        success=hit.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_client()
+
+            # Search with wildcard and sort by timestamp
+            results = client.index(self.index_name).search(
+                q="*", limit=limit, sort="timestamp:desc"
+            )
+
+            memories = []
+            for hit in results.get("hits", []):
+                args = json.loads(hit.get("args", "{}"))
+                metadata_dict = json.loads(hit.get("metadata", "{}"))
+
+                memory = ToolMemory(
+                    memory_id=hit["_id"],
+                    question=hit["question"],
+                    tool_name=hit["tool_name"],
+                    args=args,
+                    timestamp=hit.get("timestamp"),
+                    success=hit.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_client()
+
+            try:
+                client.index(self.index_name).delete_documents(ids=[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_client()
+
+            memory_id = str(uuid.uuid4())
+            timestamp = datetime.now().isoformat()
+
+            document = {
+                "_id": memory_id,
+                "content": content,
+                "timestamp": timestamp,
+                "is_text_memory": True,
+            }
+
+            client.index(self.index_name).add_documents(
+                [document], tensor_fields=["content"]
+            )
+
+            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()
+
+            filter_string = "is_text_memory:true"
+
+            results = client.index(self.index_name).search(
+                q=query, limit=limit, filter_string=filter_string
+            )
+
+            search_results = []
+            for i, hit in enumerate(results["hits"]):
+                similarity_score = hit.get("_score", 0)
+
+                if similarity_score >= similarity_threshold:
+                    memory = TextMemory(
+                        memory_id=hit["_id"],
+                        content=hit.get("content", ""),
+                        timestamp=hit.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_client()
+
+            results = client.index(self.index_name).search(
+                q="*",
+                limit=limit,
+                filter_string="is_text_memory:true",
+                sort="timestamp:desc",
+            )
+
+            memories = []
+            for hit in results.get("hits", []):
+                memory = TextMemory(
+                    memory_id=hit["_id"],
+                    content=hit.get("content", ""),
+                    timestamp=hit.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:
+                client.index(self.index_name).delete_documents(ids=[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_client()
+
+            # Build filter for search
+            filter_parts = []
+            if tool_name:
+                filter_parts.append(f"tool_name:{tool_name}")
+            if before_date:
+                filter_parts.append(f"timestamp:[* TO {before_date}]")
+
+            if filter_parts or (tool_name is None and before_date is None):
+                filter_string = " AND ".join(filter_parts) if filter_parts else None
+
+                if filter_string:
+                    # Search for documents to delete
+                    results = client.index(self.index_name).search(
+                        q="*",
+                        limit=1000,  # Max results
+                        filter_string=filter_string,
+                    )
+
+                    ids_to_delete = [hit["_id"] for hit in results.get("hits", [])]
+
+                    if ids_to_delete:
+                        client.index(self.index_name).delete_documents(
+                            ids=ids_to_delete
+                        )
+
+                    return len(ids_to_delete)
+                else:
+                    # Delete entire index and recreate
+                    try:
+                        client.delete_index(self.index_name)
+                        client.create_index(self.index_name)
+                    except Exception:
+                        pass
+                    return 0
+
+            return 0
+
+        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)

vanna/integrations/milvus/__init__.py

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