vi/services/memory/operations/facts_ops.py

"""
Facts memory operations.

Provides query and CRUD operations for factual memory (ChromaDB).
"""
import json
from uuid import uuid4
from datetime import datetime
from typing import List, Dict, Any, Optional, Tuple
from core.logger import setup_logger

logger = setup_logger('facts_ops', service_name='memory_service')


class FactsOperations:
    """Handles facts queries and CRUD operations"""

    def __init__(self, chroma_store):
        """
        Initialize facts operations.

        Args:
            chroma_store: ChromaStore instance
        """
        self.chroma_store = chroma_store

    def query(
        self,
        query: str = '',
        limit: int = 5,
        category: Optional[str] = None,
        identity_id: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Query facts from ChromaDB with optional category filtering.

        Args:
            query: Semantic search query (empty = get all matching filters)
            limit: Maximum number of facts to return
            category: Filter by fact category
            identity_id: Filter by identity ID

        Returns:
            List of fact dictionaries with metadata
        """
        try:
            collection = self.chroma_store.get_facts_collection()

            # Build metadata filters
            where_filters = {}
            if category:
                where_filters["category"] = category
            if identity_id:
                where_filters["identity_id"] = identity_id

            # Query ChromaDB
            if query:
                # Semantic search with query text
                results = collection.query(
                    query_texts=[query],
                    n_results=limit,
                    where=where_filters if where_filters else None
                )
            else:
                # Get all facts matching filters
                results = collection.get(
                    limit=limit,
                    where=where_filters if where_filters else None
                )

            # Format results
            facts = []
            if query and results['ids']:
                # Query results
                for i, doc_id in enumerate(results['ids'][0]):
                    metadata = results['metadatas'][0][i] if results['metadatas'] else {}
                    # Parse identities from JSON string
                    identities_str = metadata.get('identities', '[]')
                    identities = json.loads(identities_str) if isinstance(identities_str, str) else identities_str
                    facts.append({
                        "id": doc_id,
                        "content": results['documents'][0][i],
                        "category": metadata.get('category', 'general'),
                        "mutable": metadata.get('mutable', True),
                        "identities": identities,
                        "metadata": metadata,
                        "distance": results['distances'][0][i] if results['distances'] else None,
                        "source": "facts"
                    })
            elif not query and results['ids']:
                # Get results (no query)
                for i, doc_id in enumerate(results['ids']):
                    metadata = results['metadatas'][i] if results['metadatas'] else {}
                    # Parse identities from JSON string
                    identities_str = metadata.get('identities', '[]')
                    identities = json.loads(identities_str) if isinstance(identities_str, str) else identities_str
                    facts.append({
                        "id": doc_id,
                        "content": results['documents'][i],
                        "category": metadata.get('category', 'general'),
                        "mutable": metadata.get('mutable', True),
                        "identities": identities,
                        "metadata": metadata,
                        "source": "facts"
                    })

            logger.debug(f"[μ] Retrieved {len(facts)} facts (query='{query}', category={category}, limit={limit})")
            return facts

        except Exception as e:
            logger.error(f"[μ] Failed to query facts: {e}")
            return []

    def create(
        self,
        content: str,
        category: str = 'general',
        identities: List[str] = None,
        mutable: bool = True,
        metadata: Dict[str, Any] = None
    ) -> str:
        """
        Create a new fact in ChromaDB facts collection.

        Args:
            content: Fact content
            category: Fact category
            identities: List of identity IDs associated with this fact
            mutable: Whether the fact can be updated
            metadata: Additional metadata

        Returns:
            Created fact ID
        """
        fact_id = str(uuid4())
        identities = identities or []
        metadata = metadata or {}

        # Prepare metadata (ChromaDB only accepts scalar types, not lists)
        fact_metadata = {
            "category": category,
            "mutable": mutable,
            "identities": json.dumps(identities),  # Convert list to JSON string
            "created_at": datetime.utcnow().isoformat(),
            **metadata
        }

        # Add to ChromaDB (will automatically generate embedding)
        collection = self.chroma_store.get_facts_collection()
        collection.add(
            ids=[fact_id],
            documents=[content],
            metadatas=[fact_metadata]
        )

        logger.info(f"[μ] Created fact {fact_id}: category={category}, content='{content[:60]}...'")
        return fact_id

    def update(
        self,
        fact_id: str,
        new_content: str,
        identity_id: str = None,
        metadata: Dict[str, Any] = None
    ) -> Tuple[bool, str]:
        """
        Update an existing fact in ChromaDB if it's mutable and owned by identity.

        Args:
            fact_id: ID of fact to update
            new_content: New content for the fact
            identity_id: Identity ID for ownership validation (optional)
            metadata: Additional metadata to update

        Returns:
            (success, error_message) - error_message is empty string if successful
        """
        try:
            collection = self.chroma_store.get_facts_collection()

            # Get the fact to check if it exists and is mutable
            result = collection.get(ids=[fact_id])

            if not result['ids']:
                logger.warning(f"[μ] Fact {fact_id} not found")
                return False, "Fact not found"

            fact_metadata = result['metadatas'][0]

            # Check if fact is mutable
            if not fact_metadata.get('mutable', True):
                logger.warning(f"[μ] Fact {fact_id} is not mutable")
                return False, "Fact is not mutable"

            # Validate identity ownership if identity_id provided
            if identity_id:
                fact_identities = fact_metadata.get('identities', [])
                if isinstance(fact_identities, str):
                    try:
                        fact_identities = json.loads(fact_identities)
                    except:
                        fact_identities = [fact_identities]

                if identity_id not in fact_identities:
                    logger.warning(f"[μ] Identity {identity_id} does not own fact {fact_id}")
                    return False, f"Fact does not belong to identity {identity_id}"

            # Update metadata
            updated_metadata = dict(fact_metadata)
            updated_metadata['updated_at'] = datetime.utcnow().isoformat()
            if metadata:
                updated_metadata.update(metadata)

            # Update in ChromaDB
            collection.update(
                ids=[fact_id],
                documents=[new_content],
                metadatas=[updated_metadata]
            )

            logger.info(f"[μ] Updated fact {fact_id}")
            return True, ""

        except Exception as e:
            logger.error(f"[μ] Failed to update fact {fact_id}: {e}")
            return False, str(e)