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

parrot/stores/kb/store.py

@@ -0,0 +1,115 @@
+from collections import defaultdict
+from typing import List, Dict, Any
+
+class KnowledgeBaseStore:
+    """Lightweight in-memory store for validated facts."""
+
+    def __init__(
+        self,
+        embedding_model: str = "all-MiniLM-L6-v2",  # 384D model
+        dimension: int = 384,
+        index_type: str = "Flat",  # or "HNSW" for larger KBs
+    ):
+        try:
+            from sentence_transformers import SentenceTransformer
+            import faiss
+        except ImportError as e:
+            raise ImportError(
+                "Please install 'sentence-transformers' and 'faiss-cpu' to use KnowledgeBaseStore."
+            ) from e
+        self.embeddings = SentenceTransformer(embedding_model)
+        self.dimension = dimension
+        self.score_threshold = 0.5
+        # FAISS index
+        if index_type == "FlatIP":
+            self.index = faiss.IndexFlatIP(dimension)
+        else:
+            # HNSW with IP metric
+            self.index = faiss.IndexHNSWFlat(dimension, 32, faiss.METRIC_INNER_PRODUCT)
+
+        # Store facts and metadata
+        self.facts: List[str] = []
+        self.fact_metadata: List[dict] = []
+        self.category_index = defaultdict(list)  # Fast category lookup
+        self.entity_index = defaultdict(list)    # Entity-based retrieval
+
+    async def add_fact(self, fact: Dict[str, Any]):
+        """Add a single validated fact to the KB."""
+        await self.add_facts([fact])
+
+    async def add_facts(self, facts: List[Dict[str, Any]]):
+        """Add validated facts to the KB."""
+        if not facts:
+            return
+
+        texts = []
+        for fact in facts:
+            fact_id = len(self.facts)
+            self.facts.append(fact)
+            texts.append(fact['content'])
+            if category := fact.get('metadata', {}).get('category'):
+                self.category_index[category].append(fact_id)
+            # Index by entities
+            for key in ['subject', 'object']:
+                if entity := fact.get('metadata', {}).get(key):
+                    self.entity_index[entity].append(fact_id)
+
+        embeddings = self.embeddings.encode(texts, normalize_embeddings=True)
+        self.index.add(embeddings)
+
+        self.fact_metadata.extend(
+            [f.get('metadata', {}) for f in facts]
+        )
+
+    def _tokenize(self, text: str) -> set:
+        return {t.lower() for t in text.split()}
+
+    async def search_facts(
+        self,
+        query: str,
+        k: int = 5,
+        score_threshold: float = 0.5
+    ) -> List[Dict[str, Any]]:
+        """Ultra-fast fact retrieval."""
+        query_embedding = self.embeddings.encode(
+            [query],
+            normalize_embeddings=True
+        )
+        # Important: k should not exceed number of facts
+        actual_k = min(k, len(self.facts))
+        scores, indices = self.index.search(query_embedding, actual_k)
+        threshold = score_threshold or self.score_threshold
+
+        results = []
+        for score, idx in zip(scores[0], indices[0]):
+            if idx == -1:
+                continue
+            if float(score) >= threshold:
+                results.append({
+                    'fact': self.facts[idx],
+                    'score': float(score),
+                    'metadata': self.fact_metadata[idx]
+                })
+        # doing a re-ranking based on token overlap
+        # after collecting FAISS candidates as `results` with "score" = cosine
+        q_tokens = self._tokenize(query)
+        for r in results:
+            tags = set((r["metadata"].get("tags") or []))
+            overlap = len(q_tokens & {t.lower() for t in tags})
+            r["score"] += 0.05 * overlap  # tiny boost per overlapping tag
+        results.sort(key=lambda x: x["score"], reverse=True)
+        return results
+
+    def get_facts_by_category(self, category: str) -> List[Dict]:
+        """Retrieve all facts in a category."""
+        indices = self.category_index.get(category, [])
+        return [self.facts[i] for i in indices]
+
+    def get_entity_facts(self, entity: str) -> List[Dict]:
+        """Get all facts related to an entity."""
+        indices = self.entity_index.get(entity, [])
+        return [self.facts[i] for i in indices]
+
+    async def close(self):
+        """Cleanup resources if needed."""
+        pass

parrot/stores/kb/user.py

@@ -0,0 +1,374 @@
+from typing import Tuple, List, Dict, Any, Optional
+from asyncdb import AsyncDB
+from querysource.conf import default_dsn
+from .abstract import AbstractKnowledgeBase
+from .redis import RedisKnowledgeBase
+from ...utils.helpers import RequestContext
+from .cache import TTLCache
+
+
+class UserInfo(AbstractKnowledgeBase):
+    """Class to manage user information."""
+
+    def __init__(self, **kwargs):
+        super().__init__(
+            name="User Info",
+            category="user_context",
+            always_active=True,
+            activation_patterns=[
+                "about me"
+            ],
+            priority=10,  # High priority
+            **kwargs
+        )
+        self.db = AsyncDB('pg', dsn=default_dsn)
+        self.cache =  TTLCache(
+            max_size=500,
+            default_ttl=600  # 10 minutes for user data
+        )
+        self._initialized = False
+
+    async def _lazy_init(self):
+        """Lazy initialization of resources."""
+        if not self._initialized:
+            await self.cache.start()
+            self._initialized = True
+
+    async def should_activate(self, query: str, context: dict) -> Tuple[bool, float]:
+        # Implement activation logic based on user info
+        return True, 1.0  # Always activate with high confidence
+
+    async def search(self, query: str, user_id: int, **kwargs) -> List[Dict]:
+        """Query Database for User Information."""
+        if not user_id:
+            return []
+
+        try:
+            async with await self.db.connection() as conn:  # pylint: disable=E1101 # noqa
+                result = await conn.fetch_one(
+                    """
+                    SELECT user_id, display_name, username, email, job_code, associate_id as employee_id,
+                    associate_id, associate_oid, title, worker_type, manager_id
+                    FROM auth.vw_users WHERE user_id = $1
+                    """,
+                    user_id
+                )
+                # Ok, Result is an asyncpg Record object, iterate over the columns:
+                # iterate over the result and convert in a list of facts:
+                facts = []
+                user = dict(result)
+                for key, value in user.items():
+                    if value:
+                        facts.append({
+                            "content": f"{key.replace('_', ' ').title()}: {value}",
+                            "metadata": {
+                                    "category": "user_info",
+                                    "entity_type": key,
+                                    "confidence": 1.0,
+                                    "tags": ["user", "profile", key]
+                                }
+                            })
+                return facts
+        except Exception as e:
+            print(f"Error fetching user info: {e}")
+        return []
+
+class UserProfileKB(AbstractKnowledgeBase):
+    """KB that queries database for user information."""
+
+    def __init__(self, **kwargs):
+        super().__init__(
+            name="User Profile",
+            category="user_context",
+            activation_patterns=[
+                "my preferences", "my profile", "my settings",
+                "about me", "my role", "my permissions"
+            ],
+            priority=10  # High priority
+        )
+        self.db = AsyncDB('pg', dsn=default_dsn)
+
+    async def should_activate(self, query: str, context: Dict) -> Tuple[bool, float]:
+        """Check if query references user-specific information."""
+        query_lower = query.lower()
+
+        # Direct patterns
+        for pattern in self.activation_patterns:
+            if pattern in query_lower:
+                return True, 1.0
+
+        # Heuristic checks
+        personal_indicators = ["my", "me", "i", "our"]
+        if any(word in query_lower.split() for word in personal_indicators):
+            return True, 0.7
+
+        return False, 0.0
+
+    async def search(self, query: str, user_id: str = None, **kwargs) -> List[Dict]:
+        """Query database for user information."""
+        if not user_id:
+            return []
+
+        async with await self.db.connection() as conn:  # pylint: disable=E1101 # noqa
+            user_data = await conn.fetch_one("""
+                SELECT
+                    first_name, last_name, email,
+                    job_code, title, department_code,
+                    groups,
+                    programs
+                FROM auth.vw_users u
+                WHERE u.user_id = $1
+            """, user_id)
+
+        if not user_data:
+            return []
+
+        # Convert to facts
+        facts = [
+            {
+                'content': f"User's name is {user_data['first_name']} {user_data['last_name']}",
+                'metadata': {'field': 'name'},
+            },
+            {
+                'content': f"User's job title is {user_data['job_code']} in {user_data['department_code']}",
+                'metadata': {'field': 'position'},
+            },
+        ]
+
+        if user_data['groups']:
+            facts.append({
+                'content': f"User has groups: {', '.join(user_data['groups'])}",
+                'metadata': {'field': 'groups'}
+            })
+
+        if user_data['programs']:
+            facts.append({
+                'content': f"User is enrolled in programs: {', '.join(user_data['programs'])}",
+                'metadata': {'field': 'programs'}
+            })
+
+        return facts
+
+class SessionStateKB(AbstractKnowledgeBase):
+    """KB that retrieves from session state."""
+    def __init__(self, **kwargs):
+        super().__init__(
+            name="User Session",
+            category="session",
+            activation_patterns=[
+                "preferences", "prefer", "like", "favorite",
+                "based on what I like"
+            ]
+        )
+
+    async def search(self, query: str, ctx: RequestContext = None, **kwargs) -> List[Dict]:
+        """Extract relevant session state."""
+        if not ctx or not ctx.request:
+            return []
+
+        session = ctx.request.session
+        facts = []
+
+        # Extract current workflow state
+        if 'current_workflow' in session:
+            facts.append({
+                'content': f"User is currently working on: {session['current_workflow']}",
+                'metadata': {'source': 'session'}
+            })
+
+        return facts
+
+class UserPreferences(RedisKnowledgeBase):
+    """KB for user preferences stored in Redis."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(
+            name="User Preferences",
+            category="preferences",
+            activation_patterns=[
+                "preferences", "prefer", "like", "favorite",
+                "based on what I like", "my preferences"
+            ],
+            namespace="user_prefs",
+            use_hash_storage=True,
+            ttl=86400 * 365,  # 1 year expiration
+            **kwargs,
+        )
+
+    async def search(
+        self,
+        query: str,
+        user_id: Optional[str] = None,
+        **kwargs
+    ) -> List[Dict[str, Any]]:
+        """
+        Retrieve user preferences matching the query.
+
+        Args:
+            query: Search query
+            user_id: User identifier (required)
+            **kwargs: Additional parameters
+
+        Returns:
+            List of preference facts
+        """
+        if not user_id:
+            return []
+
+        # Get all preferences for user
+        prefs = await self.get(user_id)
+
+        if not prefs:
+            return []
+
+        facts = []
+        query_lower = query.lower()
+
+        for key, value in prefs.items():
+            # Check if query matches key or value
+            if query_lower in key.lower() or query_lower in str(value).lower():
+                # Calculate relevance
+                relevance = 1.0 if query_lower == key.lower() else 0.5
+                if query_lower in str(value).lower():
+                    relevance += 0.3
+
+                facts.append({
+                    'content': f"User prefers {key}: {value}",
+                    'metadata': {
+                        'preference': key,
+                        'value': value,
+                        'user_id': user_id
+                    },
+                    'source': 'user_preferences',
+                    'relevance': min(relevance, 1.0)  # Cap at 1.0
+                })
+
+        # Sort by relevance
+        facts.sort(key=lambda x: x['relevance'], reverse=True)
+        return facts
+
+    async def set_preference(
+        self,
+        user_id: str,
+        preference: str,
+        value: Any
+    ) -> bool:
+        """
+        Set a single user preference.
+
+        Args:
+            user_id: User identifier
+            preference: Preference name
+            value: Preference value
+
+        Returns:
+            True if successful
+        """
+        return await self.insert(user_id, value, field=preference)
+
+    async def get_preference(
+        self,
+        user_id: str,
+        preference: str,
+        default: Any = None
+    ) -> Any:
+        """
+        Get a single user preference.
+
+        Args:
+            user_id: User identifier
+            preference: Preference name
+            default: Default value if not found
+
+        Returns:
+            Preference value or default
+        """
+        return await self.get(user_id, field=preference, default=default)
+
+    async def delete_preference(
+        self,
+        user_id: str,
+        preference: str
+    ) -> bool:
+        """
+        Delete a single user preference.
+
+        Args:
+            user_id: User identifier
+            preference: Preference name
+
+        Returns:
+            True if successful
+        """
+        return await self.delete(user_id, field=preference)
+
+    async def get_all_preferences(self, user_id: str) -> Dict[str, Any]:
+        """
+        Get all preferences for a user.
+
+        Args:
+            user_id: User identifier
+
+        Returns:
+            Dict of all preferences
+        """
+        return await self.get(user_id, default={})
+
+    async def clear_all_preferences(self, user_id: str) -> bool:
+        """
+        Clear all preferences for a user.
+
+        Args:
+            user_id: User identifier
+
+        Returns:
+            True if successful
+        """
+        return await self.delete(user_id)
+
+    async def update_preferences(
+        self,
+        user_id: str,
+        preferences: Dict[str, Any]
+    ) -> bool:
+        """
+        Update multiple preferences at once.
+
+        Args:
+            user_id: User identifier
+            preferences: Dictionary of preferences to update
+
+        Returns:
+            True if successful
+        """
+        return await self.update(user_id, preferences)
+
+    async def has_preference(
+        self,
+        user_id: str,
+        preference: str
+    ) -> bool:
+        """
+        Check if a user has a specific preference set.
+
+        Args:
+            user_id: User identifier
+            preference: Preference name
+
+        Returns:
+            True if preference exists
+        """
+        return await self.exists(user_id, field=preference)
+
+    async def list_user_preference_keys(self, user_id: str) -> List[str]:
+        """
+        Get list of all preference keys for a user.
+
+        Args:
+            user_id: User identifier
+
+        Returns:
+            List of preference keys
+        """
+        prefs = await self.get_all_preferences(user_id)
+        return list(prefs.keys()) if prefs else []

parrot/stores/models.py

@@ -0,0 +1,59 @@
+from typing import Dict, Any, Optional, Union
+from enum import Enum
+from dataclasses import dataclass, field
+from pydantic import BaseModel, Field
+
+
+class SearchResult(BaseModel):
+    """
+    Data model for a single document returned from a vector search.
+    """
+    id: str
+    content: str
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    score: float
+    ensemble_score: float = None
+    search_source: str = None
+    similarity_rank: Optional[int] = None
+    mmr_rank: Optional[int] = None
+
+
+class Document(BaseModel):
+    """
+    A simple document model for adding data to the vector store.
+    This replaces langchain.docstore.document.Document.
+    """
+    page_content: str
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class DistanceStrategy(str, Enum):
+    """Enumerator of the Distance strategies for calculating distances
+    between vectors."""
+
+    EUCLIDEAN_DISTANCE = "EUCLIDEAN_DISTANCE"
+    MAX_INNER_PRODUCT = "MAX_INNER_PRODUCT"
+    DOT_PRODUCT = "DOT_PRODUCT"
+    JACCARD = "JACCARD"
+    COSINE = "COSINE"
+
+
+@dataclass
+class StoreConfig:
+    """Vector Store configuration dataclass."""
+    vector_store: str = 'postgres'  # postgres, faiss, arango, etc.
+    table: Optional[str] = None
+    schema: str = 'public'
+    embedding_model: Union[str, dict] = field(
+        default_factory=lambda: {
+            "model": "sentence-transformers/all-mpnet-base-v2",
+            "model_type": "huggingface"
+        }
+    )
+    dimension: int = 768
+    dsn: Optional[str] = None
+    distance_strategy: str = 'COSINE'
+    metric_type: str = 'COSINE'
+    index_type: str = 'IVF_FLAT'
+    auto_create: bool = False  # Auto-create collection on configure
+    extra: Dict[str, Any] = field(default_factory=dict)

parrot/stores/pgvector.py

@@ -0,0 +1,3 @@
+from .postgres import PgVectorStore
+
+__all__ = ["PgVectorStore"]