remdb 0.3.242__py3-none-any.whl

rem/services/rem/executor.py

@@ -0,0 +1,293 @@
+"""
+REM Query Executor - Shared PostgreSQL function calling layer.
+
+This module provides the single source of truth for executing REM queries
+against PostgreSQL functions (rem_lookup, rem_search, rem_fuzzy, rem_traverse).
+
+Both REMQueryService (string-based) and RemService (Pydantic-based) delegate
+to these functions to avoid code duplication.
+
+Design:
+- One function per query type
+- All embedding generation happens here
+- Direct PostgreSQL function calls
+- Type-safe parameters via Pydantic models or dicts
+"""
+
+import asyncio
+from collections import defaultdict
+from typing import Any, Optional, cast
+from loguru import logger
+
+
+class REMQueryExecutor:
+    """
+    Executor for REM PostgreSQL functions.
+
+    Provides unified backend for both string-based and Pydantic-based query services.
+    """
+
+    def __init__(self, postgres_service: Any):
+        """
+        Initialize query executor.
+
+        Args:
+            postgres_service: PostgresService instance
+        """
+        self.db = postgres_service
+        logger.debug("Initialized REMQueryExecutor")
+
+    async def execute_lookup(
+        self,
+        entity_key: str,
+        user_id: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """
+        Execute rem_lookup() PostgreSQL function.
+
+        Args:
+            entity_key: Entity key to lookup
+            user_id: Optional user filter
+
+        Returns:
+            List of entity dicts from KV_STORE
+        """
+        sql = """
+            SELECT entity_type, data
+            FROM rem_lookup($1, $2, $3)
+        """
+
+        results = await self.db.execute(sql, (entity_key, user_id, user_id))
+        # Extract JSONB records from the data column and add aliases
+        entities = []
+        for row in results:
+            entity = dict(row["data"])
+            # Add entity_key as alias for name (for backward compat with tests)
+            if "name" in entity:
+                entity["entity_key"] = entity["name"]
+            # Add entity_id as alias for id (for backward compat with tests)
+            if "id" in entity:
+                entity["entity_id"] = entity["id"]
+            entities.append(entity)
+        logger.debug(f"LOOKUP '{entity_key}': {len(entities)} results")
+        return entities
+
+    async def execute_fetch(
+        self,
+        entity_keys: list[str],
+        entity_types: list[str],
+        user_id: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """
+        Execute rem_fetch() PostgreSQL function.
+
+        Fetches full entity records (all columns) from multiple tables by entity keys.
+        Groups by table internally, fetches all records, returns unified JSONB result set.
+        Returns complete entities, not just KV store metadata.
+
+        Args:
+            entity_keys: List of entity keys to fetch
+            entity_types: Parallel list of entity types (table names)
+            user_id: Optional user filter
+
+        Returns:
+            List of full entity records as dicts with entity_key, entity_type, and entity_record
+        """
+        if not entity_keys:
+            return []
+
+        # Build JSONB structure: {"resources": ["key1", "key2"], "moments": ["key3"]}
+        import json
+        entities_by_table: dict[str, list[str]] = {}
+        for key, table in zip(entity_keys, entity_types):
+            if table not in entities_by_table:
+                entities_by_table[table] = []
+            entities_by_table[table].append(key)
+
+        entities_json = json.dumps(entities_by_table)
+
+        sql = """
+            SELECT entity_key, entity_type, entity_record
+            FROM rem_fetch($1::jsonb, $2)
+        """
+
+        results = await self.db.execute(sql, (entities_json, user_id))
+
+        logger.debug(
+            f"FETCH: {len(results)}/{len(entity_keys)} records fetched from {len(set(entity_types))} tables"
+        )
+        return cast(list[dict[str, Any]], results)
+
+    async def execute_fuzzy(
+        self,
+        query_text: str,
+        user_id: str | None = None,
+        threshold: float = 0.3,
+        limit: int = 10,
+    ) -> list[dict[str, Any]]:
+        """
+        Execute rem_fuzzy() PostgreSQL function.
+
+        Args:
+            query_text: Text to fuzzy match
+            user_id: Optional user filter
+            threshold: Similarity threshold (0.0-1.0)
+            limit: Max results
+
+        Returns:
+            List of fuzzy-matched entities with similarity_score
+        """
+        sql = """
+            SELECT entity_type, data, similarity_score
+            FROM rem_fuzzy($1, $2, $3, $4, $5)
+        """
+
+        results = await self.db.execute(
+            sql, (query_text, user_id, threshold, limit, user_id)
+        )
+        # Extract JSONB records and add similarity_score + entity_key alias
+        entities = []
+        for row in results:
+            entity = dict(row["data"])
+            entity["similarity_score"] = row["similarity_score"]
+            # Add entity_key as alias for name (for backward compat)
+            if "name" in entity:
+                entity["entity_key"] = entity["name"]
+            entities.append(entity)
+        logger.debug(f"FUZZY '{query_text}': {len(entities)} results (threshold={threshold})")
+        return entities
+
+    async def execute_search(
+        self,
+        query_embedding: list[float],
+        table_name: str,
+        field_name: str,
+        provider: str,
+        min_similarity: float = 0.7,
+        limit: int = 10,
+        user_id: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """
+        Execute rem_search() PostgreSQL function.
+
+        Args:
+            query_embedding: Embedding vector for query
+            table_name: Table to search (resources, moments, users)
+            field_name: Field name to search
+            provider: Embedding provider (openai, anthropic)
+            min_similarity: Minimum cosine similarity
+            limit: Max results
+            user_id: Optional user filter
+
+        Returns:
+            List of similar entities with distance scores
+        """
+        # Convert embedding to PostgreSQL vector format
+        embedding_str = "[" + ",".join(str(x) for x in query_embedding) + "]"
+
+        sql = """
+            SELECT entity_type, similarity_score, data
+            FROM rem_search($1::vector(1536), $2, $3, $4, $5, $6, $7, $8)
+        """
+
+        results = await self.db.execute(
+            sql,
+            (
+                embedding_str,
+                table_name,
+                field_name,
+                user_id,  # tenant_id (backward compat)
+                provider,
+                min_similarity,
+                limit,
+                user_id,  # user_id
+            ),
+        )
+        # Extract JSONB records and add similarity_score + entity_key alias
+        entities = []
+        for row in results:
+            entity = dict(row["data"])
+            entity["similarity_score"] = row["similarity_score"]
+            entity["entity_type"] = row["entity_type"]
+            # Add entity_key as alias for name (for backward compat)
+            if "name" in entity:
+                entity["entity_key"] = entity["name"]
+            # Add distance as alias for similarity_score (for backward compat)
+            # Note: similarity_score is cosine similarity (higher = more similar)
+            # distance is inverse (lower = more similar), so: distance = 1 - similarity_score
+            entity["distance"] = 1.0 - row["similarity_score"]
+            entities.append(entity)
+        logger.debug(
+            f"SEARCH in {table_name}.{field_name}: {len(entities)} results (similarity≥{min_similarity})"
+        )
+        return entities
+
+    async def execute_traverse(
+        self,
+        start_key: str,
+        direction: str,
+        max_depth: int,
+        edge_types: list[str] | None,
+        user_id: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """
+        Execute rem_traverse() PostgreSQL function.
+
+        Args:
+            start_key: Starting entity key
+            direction: OUTBOUND, INBOUND, or BOTH (not used in current function)
+            max_depth: Maximum traversal depth
+            edge_types: Optional list of edge types to filter
+            user_id: Optional user filter
+
+        Returns:
+            List of traversed entities with path information
+        """
+        # Convert edge_types to PostgreSQL array or NULL
+        edge_types_sql = None
+        if edge_types:
+            edge_types_sql = "{" + ",".join(edge_types) + "}"
+
+        # Note: rem_traverse signature is (entity_key, tenant_id, user_id, max_depth, rel_type, keys_only)
+        # tenant_id is for backward compat, set to user_id
+        # direction parameter is not used by the current PostgreSQL function
+        # edge_types is single value, not array
+        edge_type_filter = edge_types[0] if edge_types else None
+
+        sql = """
+            SELECT depth, entity_key, entity_type, entity_id, rel_type, rel_weight, path, entity_record
+            FROM rem_traverse($1, $2, $3, $4, $5, $6)
+        """
+
+        results = await self.db.execute(
+            sql, (start_key, user_id, user_id, max_depth, edge_type_filter, False)
+        )
+        # Add edge_type alias for rel_type (backward compat)
+        processed_results = []
+        for row in results:
+            result = dict(row)
+            if "rel_type" in result:
+                result["edge_type"] = result["rel_type"]
+            processed_results.append(result)
+
+        logger.debug(
+            f"TRAVERSE from '{start_key}' (depth={max_depth}): {len(processed_results)} results"
+        )
+        return processed_results
+
+    async def execute_sql(
+        self,
+        query: str,
+    ) -> list[dict[str, Any]]:
+        """
+        Execute raw SQL query.
+
+        Args:
+            query: SQL query string
+
+        Returns:
+            Query results as list of dicts
+        """
+        results = await self.db.execute(query)
+        logger.debug(f"SQL query: {len(results)} results")
+        return cast(list[dict[str, Any]], results)

rem/services/rem/parser.py

@@ -0,0 +1,180 @@
+import shlex
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from ...models.core import QueryType
+
+
+class RemQueryParser:
+    """
+    Robust parser for REM query language using shlex for proper quoting support.
+    """
+
+    def parse(self, query_string: str) -> Tuple[QueryType, Dict[str, Any]]:
+        """
+        Parse a REM query string into a QueryType and a dictionary of parameters.
+
+        Args:
+            query_string: The raw query string (e.g., 'LOOKUP "Sarah Chen"').
+
+        Returns:
+            Tuple of (QueryType, parameters_dict).
+
+        Raises:
+            ValueError: If the query string is empty or has an invalid query type.
+        """
+        if not query_string or not query_string.strip():
+            raise ValueError("Empty query string")
+
+        try:
+            # Use shlex to handle quoted strings correctly
+            tokens = shlex.split(query_string)
+        except ValueError as e:
+            raise ValueError(f"Failed to parse query string: {e}")
+
+        if not tokens:
+            raise ValueError("Empty query string")
+
+        query_type_str = tokens[0].upper()
+
+        # Try to match REM query types first
+        try:
+            query_type = QueryType(query_type_str)
+        except ValueError:
+            # If not a known REM query type, treat as raw SQL
+            # This supports SELECT, INSERT, UPDATE, DELETE, WITH, DROP, CREATE, ALTER, etc.
+            query_type = QueryType.SQL
+            # Return raw SQL query directly in params
+            params = {"raw_query": query_string.strip()}
+            return query_type, params
+
+        params: Dict[str, Any] = {}
+        positional_args: List[str] = []
+
+        # For SQL queries, preserve the raw query (keywords like LIMIT are SQL keywords)
+        if query_type == QueryType.SQL:
+            # Everything after "SQL" is the raw SQL query
+            raw_sql = query_string[3:].strip()  # Skip "SQL" prefix
+            params["raw_query"] = raw_sql
+            return query_type, params
+
+        # Process remaining tokens, handling REM keywords
+        i = 1
+        while i < len(tokens):
+            token = tokens[i]
+            token_upper = token.upper()
+
+            # Handle REM keywords that take a value
+            if token_upper in ("LIMIT", "DEPTH", "THRESHOLD", "TYPE", "FROM", "WITH"):
+                if i + 1 < len(tokens):
+                    keyword_map = {
+                        "LIMIT": "limit",
+                        "DEPTH": "max_depth",
+                        "THRESHOLD": "threshold",
+                        "TYPE": "edge_types",
+                        "FROM": "initial_query",
+                        "WITH": "initial_query",
+                    }
+                    key = keyword_map[token_upper]
+                    value = tokens[i + 1]
+                    params[key] = self._convert_value(key, value)
+                    i += 2
+                    continue
+            elif "=" in token:
+                # It's a keyword argument
+                key, value = token.split("=", 1)
+                # Handle parameter aliases
+                mapped_key = self._map_parameter_alias(key)
+                params[mapped_key] = self._convert_value(mapped_key, value)
+            else:
+                # It's a positional argument part
+                positional_args.append(token)
+            i += 1
+
+        # Map positional arguments to specific fields based on QueryType
+        self._map_positional_args(query_type, positional_args, params)
+
+        return query_type, params
+
+    def _map_parameter_alias(self, key: str) -> str:
+        """
+        Map common aliases to internal model field names.
+        """
+        aliases = {
+            "table": "table_name",
+            "field": "field_name",
+            "where": "where_clause",
+            "depth": "max_depth",
+            "rel_type": "edge_types",
+            "rel_types": "edge_types",
+        }
+        return aliases.get(key, key)
+
+    def _convert_value(self, key: str, value: str) -> Union[str, int, float, List[str]]:
+        """
+        Convert string values to appropriate types based on the key name.
+        """
+        # Integer fields
+        if key in ("limit", "max_depth", "depth", "limit"):
+            try:
+                return int(value)
+            except ValueError:
+                return value  # Return as string if conversion fails (validation will catch it)
+
+        # Float fields
+        if key in ("threshold", "min_similarity"):
+            try:
+                return float(value)
+            except ValueError:
+                return value
+
+        # List fields (comma-separated)
+        if key in ("edge_types", "tags"):
+            return [v.strip() for v in value.split(",")]
+
+        # Default to string
+        return value
+
+    def _map_positional_args(
+        self, query_type: QueryType, positional_args: List[str], params: Dict[str, Any]
+    ) -> None:
+        """
+        Map accumulated positional arguments to the primary field for the query type.
+        """
+        if not positional_args:
+            return
+
+        # Join positional args with space to reconstruct the text
+        # This handles cases where the user didn't quote a multi-word string
+        # e.g. FUZZY Sarah Chen -> "Sarah Chen"
+        combined_value = " ".join(positional_args)
+
+        if query_type == QueryType.LOOKUP:
+            # LOOKUP supports list of keys, but as positional arg we treat as single key or comma-separated
+            # If the user provided "key1 key2", it might be interpreted as one key "key1 key2"
+            # or multiple keys. For now, let's assume it's a single key entity name unless it has commas.
+            if "," in combined_value:
+                 params["key"] = [k.strip() for k in combined_value.split(",")]
+            else:
+                 params["key"] = combined_value
+
+        elif query_type == QueryType.FUZZY:
+            params["query_text"] = combined_value
+
+        elif query_type == QueryType.SEARCH:
+            # SEARCH expects: SEARCH <table> <query_text> [LIMIT n]
+            # First positional arg is table name, rest is query text
+            if len(positional_args) >= 2:
+                params["table_name"] = positional_args[0]
+                params["query_text"] = " ".join(positional_args[1:])
+            elif len(positional_args) == 1:
+                # Could be table name or query text - assume query text if no table
+                params["query_text"] = positional_args[0]
+            # If no positional args, params stays empty
+
+        elif query_type == QueryType.TRAVERSE:
+            params["initial_query"] = combined_value
+        
+        elif query_type == QueryType.SQL:
+            # SQL with positional args means "SQL SELECT * FROM ..." form
+            # Treat the combined positional args as the raw SQL query
+            params["raw_query"] = combined_value

rem/services/rem/queries.py

@@ -0,0 +1,196 @@
+"""
+REM Query SQL Templates.
+
+All SQL queries for REM operations are defined here with proper parameterization.
+This separates query logic from business logic and makes queries easier to maintain.
+
+Design Pattern:
+- Each query is a named constant with $1, $2, etc. placeholders
+- Query parameters are documented in docstrings
+- Queries delegate to PostgreSQL functions for performance
+- All queries include tenant isolation
+"""
+
+# LOOKUP Query
+# Delegates to rem_lookup() PostgreSQL function
+# Returns raw JSONB data for LLM consumption
+LOOKUP_QUERY = """
+SELECT
+    entity_type,
+    data
+FROM rem_lookup($1, $2, $3)
+"""
+# Parameters:
+# $1: entity_key (str)
+# $2: tenant_id (str)
+# $3: user_id (str | None)
+# Returns:
+# - entity_type: Table name (e.g., "resources", "users")
+# - data: Complete entity record as JSONB
+
+
+# FUZZY Query
+# Delegates to rem_fuzzy() PostgreSQL function
+# Returns raw JSONB data with similarity scores
+FUZZY_QUERY = """
+SELECT
+    entity_type,
+    similarity_score,
+    data
+FROM rem_fuzzy($1, $2, $3, $4, $5)
+"""
+# Parameters:
+# $1: query_text (str)
+# $2: tenant_id (str)
+# $3: threshold (float)
+# $4: limit (int)
+# $5: user_id (str | None)
+# Returns:
+# - entity_type: Table name (e.g., "resources", "files")
+# - similarity_score: Fuzzy match score (0.0-1.0)
+# - data: Complete entity record as JSONB
+
+
+# SEARCH Query
+# Delegates to rem_search() PostgreSQL function
+# Returns raw JSONB data with similarity scores
+SEARCH_QUERY = """
+SELECT
+    entity_type,
+    similarity_score,
+    data
+FROM rem_search($1, $2, $3, $4, $5, $6, $7, $8)
+"""
+# Parameters:
+# $1: query_embedding (list[float])
+# $2: table_name (str)
+# $3: field_name (str)
+# $4: tenant_id (str)
+# $5: provider (str)
+# $6: min_similarity (float)
+# $7: limit (int)
+# $8: user_id (str | None)
+# Returns:
+# - entity_type: Table name (e.g., "resources", "moments")
+# - similarity_score: Vector similarity (0.0-1.0)
+# - data: Complete entity record as JSONB
+
+
+# TRAVERSE Query
+# Delegates to rem_traverse() PostgreSQL function
+TRAVERSE_QUERY = """
+SELECT
+    depth,
+    entity_key,
+    entity_type,
+    entity_id,
+    rel_type,
+    rel_weight,
+    path
+FROM rem_traverse($1, $2, $3, $4, $5, $6)
+"""
+# Parameters:
+# $1: start_key (str)
+# $2: tenant_id (str)
+# $3: user_id (str | None)
+# $4: max_depth (int)
+# $5: rel_type (str | None) - single type, not array
+# $6: keys_only (bool)
+
+
+# SQL Query Builder
+# Direct SQL queries with tenant isolation
+def build_sql_query(table_name: str, where_clause: str, tenant_id: str, limit: int | None = None) -> str:
+    """
+    Build SQL query with tenant isolation.
+
+    Args:
+        table_name: Table name (e.g., "resources", "moments")
+        where_clause: WHERE clause (e.g., "moment_type='meeting'")
+        tenant_id: Tenant identifier for isolation
+        limit: Optional result limit
+
+    Returns:
+        Parameterized SQL query string
+
+    Note:
+        This builds a dynamic query. Consider using prepared statements
+        or query builders like SQLAlchemy for production.
+    """
+    # Sanitize table name (basic validation)
+    allowed_tables = ["resources", "moments", "messages", "users", "files"]
+    if table_name not in allowed_tables:
+        raise ValueError(f"Invalid table name: {table_name}")
+
+    # Build query with tenant isolation
+    where_clause = where_clause or "1=1"
+    query = f"SELECT * FROM {table_name} WHERE tenant_id = $1 AND ({where_clause})"
+
+    if limit:
+        query += f" LIMIT {int(limit)}"
+
+    return query
+
+
+# Helper: Get query parameters for LOOKUP
+def get_lookup_params(entity_key: str, tenant_id: str, user_id: str | None = None) -> tuple:
+    """Get parameters for LOOKUP query."""
+    return (entity_key, tenant_id, user_id)
+
+
+# Helper: Get query parameters for FUZZY
+def get_fuzzy_params(
+    query_text: str,
+    tenant_id: str,
+    threshold: float = 0.7,
+    limit: int = 10,
+    user_id: str | None = None,
+) -> tuple:
+    """Get parameters for FUZZY query."""
+    return (query_text, tenant_id, threshold, limit, user_id)
+
+
+# Helper: Get query parameters for SEARCH
+def get_search_params(
+    query_embedding: list[float],
+    table_name: str,
+    field_name: str,
+    tenant_id: str,
+    provider: str,
+    min_similarity: float = 0.7,
+    limit: int = 10,
+    user_id: str | None = None,
+) -> tuple:
+    """
+    Get parameters for SEARCH query.
+
+    Note: provider parameter is required (no default) - should come from settings.
+    """
+    return (
+        str(query_embedding),
+        table_name,
+        field_name,
+        tenant_id,
+        provider,
+        min_similarity,
+        limit,
+        user_id,
+    )
+
+
+# Helper: Get query parameters for TRAVERSE
+def get_traverse_params(
+    start_key: str,
+    tenant_id: str,
+    user_id: str | None,
+    max_depth: int = 1,
+    rel_type: str | None = None,
+    keys_only: bool = False,
+) -> tuple:
+    """
+    Get parameters for TRAVERSE query.
+
+    Note: rel_type is singular (not array) - PostgreSQL function filters by single type.
+    If you need multiple types, call traverse multiple times or update the function.
+    """
+    return (start_key, tenant_id, user_id, max_depth, rel_type, keys_only)