ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/stores/kb/__init__.py

@@ -0,0 +1,9 @@
+from .abstract import AbstractKnowledgeBase
+from .redis import RedisKnowledgeBase
+from .local import LocalKB
+
+__all__ = (
+    'AbstractKnowledgeBase',
+    'RedisKnowledgeBase',
+    'LocalKB',
+)

parrot/stores/kb/abstract.py

@@ -0,0 +1,68 @@
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any, Tuple
+from navconfig.logging import logging
+from ...utils.helpers import RequestContext
+
+
+class AbstractKnowledgeBase(ABC):
+    """Base class for all knowledge bases."""
+
+    def __init__(
+        self,
+        name: str,
+        category: str,
+        description: str = None,
+        activation_patterns: List[str] = None,
+        always_active: bool = False,
+        priority: int = 0
+    ):
+        self.name = name
+        self.category = category
+        self.activation_patterns = activation_patterns or []
+        self.description = description or f"{name} knowledge base"
+        self.always_active = always_active
+        self.priority = priority  # Higher = included first
+        self.logger = logging.getLogger(__name__)
+
+    @abstractmethod
+    async def should_activate(
+        self,
+        query: str,
+        context: Dict[str, Any]
+    ) -> Tuple[bool, float]:
+        """
+        Determine if this KB should be activated for the query.
+        Returns (should_activate, confidence_score)
+        """
+        pass
+
+    async def close(self):
+        """Cleanup resources if needed."""
+        pass
+
+    @abstractmethod
+    async def search(
+        self,
+        query: str,
+        k: int = 5,
+        score_threshold: float = 0.6,
+        user_id: str = None,
+        session_id: str = None,
+        ctx: RequestContext = None,
+        **kwargs
+    ) -> List[Dict[str, Any]]:
+        """
+        Search for relevant facts/information.
+        Different KBs implement this differently.
+        """
+        pass
+
+    def format_context(self, results: List[Dict]) -> str:
+        """Format results for prompt injection."""
+        if not results:
+            return ""
+
+        lines = [f"## {self.name}:"]
+        for result in results:
+            lines.append(f"* {result.get('content', result)}")
+        return "\n".join(lines)

parrot/stores/kb/cache.py

@@ -0,0 +1,165 @@
+import asyncio
+import hashlib
+import json
+from typing import Any, Dict, Optional, Callable
+from datetime import datetime, timedelta
+from dataclasses import dataclass, field
+
+@dataclass
+class CacheEntry:
+    """Cache entry with TTL support."""
+    value: Any
+    expires_at: datetime
+    hit_count: int = 0
+
+    @property
+    def is_expired(self) -> bool:
+        return datetime.now() > self.expires_at
+
+    def accessed(self):
+        """Track cache hits for analytics."""
+        self.hit_count += 1
+
+class TTLCache:
+    """Thread-safe TTL cache with memory management."""
+
+    def __init__(
+        self,
+        max_size: int = 1000,
+        default_ttl: int = 300,  # 5 minutes default
+        cleanup_interval: int = 60  # Run cleanup every minute
+    ):
+        self._cache: Dict[str, CacheEntry] = {}
+        self._max_size = max_size
+        self._default_ttl = default_ttl
+        self._lock = asyncio.Lock()
+        self._cleanup_task = None
+        self._stats = {
+            'hits': 0,
+            'misses': 0,
+            'evictions': 0
+        }
+        self.cleanup_interval = cleanup_interval
+
+    async def start(self):
+        """Start the cleanup background task."""
+        if not self._cleanup_task:
+            self._cleanup_task = asyncio.create_task(
+                self._periodic_cleanup(self.cleanup_interval)
+            )
+
+    async def stop(self):
+        """Stop the cleanup task."""
+        if self._cleanup_task:
+            self._cleanup_task.cancel()
+            try:
+                await self._cleanup_task
+            except asyncio.CancelledError:
+                pass
+
+    def _generate_key(self, *args, **kwargs) -> str:
+        """Generate cache key from arguments."""
+        key_data = {
+            'args': args,
+            'kwargs': sorted(kwargs.items())
+        }
+        key_str = json.dumps(key_data, sort_keys=True, default=str)
+        return hashlib.md5(key_str.encode()).hexdigest()
+
+    async def get(
+        self,
+        key: str,
+        default: Any = None
+    ) -> Optional[Any]:
+        """Get value from cache."""
+        async with self._lock:
+            entry = self._cache.get(key)
+
+            if entry is None:
+                self._stats['misses'] += 1
+                return default
+
+            if entry.is_expired:
+                del self._cache[key]
+                self._stats['misses'] += 1
+                return default
+
+            entry.accessed()
+            self._stats['hits'] += 1
+            return entry.value
+
+    async def set(
+        self,
+        key: str,
+        value: Any,
+        ttl: Optional[int] = None
+    ):
+        """Set value in cache with TTL."""
+        ttl = ttl or self._default_ttl
+        expires_at = datetime.now() + timedelta(seconds=ttl)
+
+        async with self._lock:
+            # Check size limit
+            if len(self._cache) >= self._max_size:
+                await self._evict_lru()
+
+            self._cache[key] = CacheEntry(
+                value=value,
+                expires_at=expires_at
+            )
+
+    async def invalidate(self, pattern: str = None):
+        """Invalidate cache entries matching pattern."""
+        async with self._lock:
+            if pattern:
+                keys_to_remove = [
+                    k for k in self._cache.keys()
+                    if pattern in k
+                ]
+                for key in keys_to_remove:
+                    del self._cache[key]
+            else:
+                self._cache.clear()
+
+    async def _evict_lru(self):
+        """Evict least recently used entries."""
+        if not self._cache:
+            return
+
+        # Find oldest entry
+        oldest_key = min(
+            self._cache.keys(),
+            key=lambda k: self._cache[k].expires_at
+        )
+        del self._cache[oldest_key]
+        self._stats['evictions'] += 1
+
+    async def _periodic_cleanup(self, interval: int):
+        """Periodically clean up expired entries."""
+        while True:
+            try:
+                await asyncio.sleep(interval)
+                await self._cleanup_expired()
+            except asyncio.CancelledError:
+                break
+
+    async def _cleanup_expired(self):
+        """Remove expired entries."""
+        async with self._lock:
+            expired_keys = [
+                k for k, v in self._cache.items()
+                if v.is_expired
+            ]
+            for key in expired_keys:
+                del self._cache[key]
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get cache statistics."""
+        total = self._stats['hits'] + self._stats['misses']
+        hit_rate = (self._stats['hits'] / total * 100) if total > 0 else 0
+
+        return {
+            **self._stats,
+            'size': len(self._cache),
+            'hit_rate': f"{hit_rate:.2f}%"
+        }

parrot/stores/kb/doc.py

@@ -0,0 +1,325 @@
+from typing import Any, Dict, List, Optional
+from .redis import RedisKnowledgeBase
+
+
+class UserContext(RedisKnowledgeBase):
+    """Knowledge Base for user context and session data."""
+
+    def __init__(self, **kwargs):
+        super().__init__(
+            name="User Context",
+            category="context",
+            namespace="user_context",
+            activation_patterns=[
+                "last conversation", "previously", "we discussed",
+                "earlier you mentioned", "remember when"
+            ],
+            use_hash_storage=True,
+            ttl=86400 * 30,  # 30 days expiration
+            **kwargs
+        )
+
+    async def search(
+        self,
+        query: str,
+        user_id: Optional[str] = None,
+        **kwargs
+    ) -> List[Dict[str, Any]]:
+        """Retrieve user context matching the query."""
+        if not user_id:
+            return []
+
+        context = await self.get(user_id)
+        if not context:
+            return []
+
+        facts = []
+        query_lower = query.lower()
+
+        # Search through context fields
+        for key, value in context.items():
+            if query_lower in key.lower() or query_lower in str(value).lower():
+                facts.append({
+                    'content': f"User context - {key}: {value}",
+                    'metadata': {
+                        'context_key': key,
+                        'user_id': user_id,
+                        'value': value
+                    },
+                    'source': 'user_context',
+                    'relevance': 1.0 if query_lower == key.lower() else 0.5
+                })
+
+        return facts
+
+    async def update_context(
+        self,
+        user_id: str,
+        context_data: Dict[str, Any]
+    ) -> bool:
+        """
+        Update user context with new data.
+
+        Args:
+            user_id: User identifier
+            context_data: Context data to update
+
+        Returns:
+            True if successful
+        """
+        return await self.update(user_id, context_data)
+
+    async def get_context(self, user_id: str) -> Dict[str, Any]:
+        """Get all context for a user."""
+        return await self.get(user_id, default={})
+
+    async def set_context_field(
+        self,
+        user_id: str,
+        field: str,
+        value: Any
+    ) -> bool:
+        """Set a specific context field."""
+        return await self.insert(user_id, value, field=field)
+
+    async def get_context_field(
+        self,
+        user_id: str,
+        field: str,
+        default: Any = None
+    ) -> Any:
+        """Get a specific context field."""
+        return await self.get(user_id, field=field, default=default)
+
+    async def clear_context(self, user_id: str) -> bool:
+        """Clear all context for a user."""
+        return await self.delete(user_id)
+
+
+class ChatbotSettings(RedisKnowledgeBase):
+    """Knowledge Base for chatbot-specific settings."""
+
+    def __init__(self, **kwargs):
+        super().__init__(
+            name="Chatbot Settings",
+            category="settings",
+            namespace="bot_settings",
+            activation_patterns=[],  # Not activated by patterns
+            use_hash_storage=True,
+            **kwargs
+        )
+
+    async def search(
+        self,
+        query: str,
+        chatbot_id: Optional[str] = None,
+        **kwargs
+    ) -> List[Dict[str, Any]]:
+        """Retrieve chatbot settings."""
+        if not chatbot_id:
+            return []
+
+        settings = await self.get(chatbot_id)
+        if not settings:
+            return []
+
+        return [{
+            'content': f"Chatbot settings: {settings}",
+            'metadata': {
+                'chatbot_id': chatbot_id,
+                'settings': settings
+            },
+            'source': 'chatbot_settings'
+        }]
+
+    async def get_setting(
+        self,
+        chatbot_id: str,
+        setting: str,
+        default: Any = None
+    ) -> Any:
+        """Get a specific chatbot setting."""
+        return await self.get(chatbot_id, field=setting, default=default)
+
+    async def set_setting(
+        self,
+        chatbot_id: str,
+        setting: str,
+        value: Any
+    ) -> bool:
+        """Set a specific chatbot setting."""
+        return await self.insert(chatbot_id, value, field=setting)
+
+class DocumentMetadata(RedisKnowledgeBase):
+    """Knowledge Base for document metadata and indexing."""
+
+    def __init__(self, **kwargs):
+        super().__init__(
+            name="Document Metadata",
+            category="documents",
+            namespace="doc_meta",
+            activation_patterns=[
+                "document", "file", "uploaded", "attachment"
+            ],
+            use_hash_storage=True,
+            **kwargs
+        )
+
+    async def search(
+        self,
+        query: str,
+        user_id: Optional[str] = None,
+        **kwargs
+    ) -> List[Dict[str, Any]]:
+        """Search document metadata.
+
+        FIXED: Properly handle user_id filtering since it's stored as a field,
+        not part of the key structure.
+        """
+        query_lower = query.lower()
+        results = []
+
+        # Search across ALL documents (don't pass user_id to base search)
+        # The base search would try to use it in the key pattern which is wrong
+        all_results = await super().search(
+            query,
+            identifier=None,  # Search all documents
+            field_filter=['title', 'description', 'tags', 'filename'],
+            limit=kwargs.get('limit', 100)
+        )
+
+        # Now filter by user_id if provided
+        for result in all_results:
+            doc_data = result.get('data', {})
+
+            # Filter by user_id if specified
+            if user_id and doc_data.get('user_id') != user_id:
+                continue
+
+            # Format results for document context
+            results.append({
+                'content': f"Document: {doc_data.get('title', 'Unknown')} - {doc_data.get('description', '')}",
+                'metadata': {
+                    'document_id': result.get('identifier'),
+                    'user_id': doc_data.get('user_id'),
+                    'title': doc_data.get('title'),
+                    'filename': doc_data.get('filename'),
+                    'tags': doc_data.get('tags', []),
+                    **{k: v for k, v in doc_data.items() if k not in ['title', 'filename', 'tags', 'user_id']}
+                },
+                'source': 'document_metadata',
+                'relevance': result.get('relevance', 0.0)
+            })
+
+        # Sort by relevance
+        results.sort(key=lambda x: x.get('relevance', 0.0), reverse=True)
+        return results
+
+    async def add_document(
+        self,
+        doc_id: str,
+        user_id: str,
+        title: str,
+        filename: str,
+        description: str = "",
+        tags: List[str] = None,
+        **metadata
+    ) -> bool:
+        """
+        Add document metadata.
+
+        Args:
+            doc_id: Document identifier
+            user_id: Owner user ID
+            title: Document title
+            filename: Original filename
+            description: Document description
+            tags: Document tags
+            **metadata: Additional metadata
+
+        Returns:
+            True if successful
+        """
+        doc_data = {
+            'user_id': user_id,
+            'title': title,
+            'filename': filename,
+            'description': description,
+            'tags': tags or [],
+            **metadata
+        }
+        return await self.insert(doc_id, doc_data)
+
+    async def get_document(self, doc_id: str) -> Optional[Dict[str, Any]]:
+        """Get a specific document's metadata."""
+        return await self.get(doc_id)
+
+    async def delete_document(self, doc_id: str) -> bool:
+        """ Delete a document's metadata."""
+        return await self.delete(doc_id)
+
+    async def list_user_documents(
+        self,
+        user_id: str,
+        limit: int = 100
+    ) -> List[Dict[str, Any]]:
+        """List all documents for a specific user."""
+        all_docs = await self.list_all(limit=limit)
+
+        user_docs = []
+        for doc in all_docs:
+            doc_data = doc.get('data', {})
+            if doc_data.get('user_id') == user_id:
+                user_docs.append({
+                    'document_id': doc['identifier'],
+                    **doc_data
+                })
+
+        return user_docs
+
+    async def search_by_tags(
+        self,
+        tags: List[str],
+        user_id: Optional[str] = None,
+        match_all: bool = False
+    ) -> List[Dict[str, Any]]:
+        """
+        Search documents by tags.
+
+        Args:
+            tags: List of tags to search for
+            user_id: Optional user filter
+            match_all: If True, document must have all tags; if False, any tag matches
+
+        Returns:
+            List of matching documents
+        """
+        all_docs = await self.list_all()
+        results = []
+
+        for doc in all_docs:
+            doc_data = doc.get('data', {})
+
+            # Filter by user if specified
+            if user_id and doc_data.get('user_id') != user_id:
+                continue
+
+            doc_tags = doc_data.get('tags', [])
+
+            # Check tag matching
+            if match_all:
+                # Document must have all specified tags
+                if all(tag in doc_tags for tag in tags):
+                    results.append({
+                        'document_id': doc['identifier'],
+                        **doc_data
+                    })
+            else:
+                # Document must have at least one tag
+                if any(tag in doc_tags for tag in tags):
+                    results.append({
+                        'document_id': doc['identifier'],
+                        **doc_data
+                    })
+
+        return results