beadboard/.agents/skills/rlm-mem/brain/scripts/repl_functions.py

"""
RLM-MEM - REPL Functions
Memory access functions available within the REPL sandbox.
"""

from typing import Dict, Any, List, Optional


import re


def read_chunk(chunk_id: str, chunk_store) -> Optional[Dict[str, Any]]:
    """
    Read a chunk by ID.

    Args:
        chunk_id: The chunk ID to read
        chunk_store: ChunkStore instance

    Returns:
        Chunk data dict or None if not found
    """
    # Validate chunk_id format - reject path traversal attempts
    if chunk_id is None:
        return None

    # Check for path traversal patterns
    if ".." in chunk_id or "/" in chunk_id or "\\" in chunk_id:
        return None

    # Only allow alphanumeric, hyphens, and underscores
    if not re.match(r"^[a-zA-Z0-9_-]+$", chunk_id):
        return None

    try:
        chunk = chunk_store.get_chunk(chunk_id)
        if chunk is None:
            return None

        # Convert Chunk dataclass to dict
        return {
            "id": chunk.id,
            "content": chunk.content,
            "tokens": chunk.tokens,
            "type": chunk.type,
            "metadata": chunk.metadata,
            "links": chunk.links,
            "tags": chunk.tags,
        }
    except (AttributeError, TypeError, KeyError, ValueError):
        return None


def search_chunks(query: str, chunk_store, limit: int = 10) -> List[str]:
    """
    Search for chunks matching query.

    Args:
        query: Search query string
        chunk_store: ChunkStore instance
        limit: Maximum results to return

    Returns:
        List of matching chunk IDs
    """
    try:
        # Simple keyword search for now
        # In production, this could use embeddings or more sophisticated search
        query_lower = query.lower()
        words = set(query_lower.split())

        all_chunks = chunk_store.list_chunks()
        results = []

        for chunk_id in all_chunks:
            chunk = chunk_store.get_chunk(chunk_id)
            if chunk is None:
                continue

            content_lower = chunk.content.lower()

            # Check if any query word appears in content
            if any(word in content_lower for word in words):
                results.append(chunk_id)

            if len(results) >= limit:
                break

        return results
    except (AttributeError, TypeError, KeyError, ValueError):
        return []


def list_chunks_by_tag(tags, chunk_store) -> List[str]:
    """
    List all chunks with given tag(s).

    Args:
        tags: Single tag string or list of tags to search for
        chunk_store: ChunkStore instance

    Returns:
        List of chunk IDs with the tag(s)
    """
    try:
        # Handle single tag or list of tags
        if isinstance(tags, str):
            return chunk_store.list_chunks(tags=[tags])
        elif isinstance(tags, list):
            return chunk_store.list_chunks(tags=tags)
        return []
    except (AttributeError, TypeError, KeyError, ValueError):
        return []


def get_linked_chunks(
    chunk_id: str, chunk_store, link_type: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    Get chunks linked to the given chunk.

    Args:
        chunk_id: Source chunk ID
        chunk_store: ChunkStore instance
        link_type: Optional link type filter (e.g., 'context_of', 'follows', 'related_to')

    Returns:
        List of linked chunk data dicts
    """
    try:
        chunk = chunk_store.get_chunk(chunk_id)
        if chunk is None:
            return []

        linked = []
        for link in chunk.links:
            # Filter by link type if specified
            if link_type and link.get("type") != link_type:
                continue

            target_id = link.get("target_id")
            if target_id:
                target_chunk = read_chunk(target_id, chunk_store)
                if target_chunk:
                    # Include link metadata
                    target_chunk["_link_type"] = link.get("type", "unknown")
                    target_chunk["_link_strength"] = link.get("strength", 0.5)
                    linked.append(target_chunk)

        return linked
    except (AttributeError, TypeError, KeyError, ValueError):
        return []