remdb 0.2.6__py3-none-any.whl

rem/agentic/serialization.py

@@ -0,0 +1,245 @@
+"""
+Pydantic Serialization Utilities for Agent Results.
+
+Critical Pattern:
+When returning Pydantic model instances from agent results (especially in MCP tools,
+API responses, or any serialization context), ALWAYS serialize them explicitly using
+.model_dump() or .model_dump_json() before returning.
+
+Why This Matters:
+- FastMCP, FastAPI, and other frameworks may use their own serialization logic
+- Pydantic models returned directly may not include all fields during serialization
+- Newly added fields might be silently dropped if not explicitly serialized
+- result.output or result.data may be a Pydantic model instance, not a dict
+
+Common Anti-Patterns to Avoid:
+```python
+# ❌ BAD: Returns Pydantic model directly
+return {
+    "status": "success",
+    "response": result.output,  # Pydantic model instance!
+}
+
+# ✅ GOOD: Explicitly serialize first
+return {
+    "status": "success",
+    "response": result.output.model_dump(),  # Serialized dict
+}
+```
+
+Design Rules:
+1. Always check if object has .model_dump() or .model_dump_json()
+2. Use serialize_agent_result() for consistent handling
+3. In streaming contexts, use .model_dump_json() for SSE
+4. Document when functions return Pydantic models vs dicts
+"""
+
+from typing import Any, cast
+
+from pydantic import BaseModel
+
+
+def serialize_agent_result(result: Any) -> dict[str, Any] | str:
+    """
+    Safely serialize an agent result, handling Pydantic models correctly.
+
+    This function ensures that Pydantic model instances are properly serialized
+    before being returned from API endpoints, MCP tools, or any other context
+    where serialization is critical.
+
+    Args:
+        result: Agent result which may be:
+            - Pydantic model instance (has .model_dump())
+            - Dict (already serialized)
+            - Primitive type (str, int, bool, None)
+            - List or other collection
+
+    Returns:
+        Serialized result as dict or primitive type
+
+    Examples:
+        >>> # With Pydantic model result
+        >>> agent_result = await agent.run(query)
+        >>> serialized = serialize_agent_result(agent_result.output)
+        >>> return {"response": serialized}  # Safe to serialize
+
+        >>> # With already-serialized result
+        >>> data = {"key": "value"}
+        >>> serialized = serialize_agent_result(data)
+        >>> assert serialized == data  # No-op for dicts
+
+        >>> # With primitive result
+        >>> result = "Hello world"
+        >>> serialized = serialize_agent_result(result)
+        >>> assert serialized == result  # No-op for primitives
+    """
+    # Check if this is a Pydantic model instance
+    if isinstance(result, BaseModel):
+        return result.model_dump()
+
+    # Check if this has a model_dump method (duck typing)
+    if hasattr(result, "model_dump") and callable(getattr(result, "model_dump")):
+        return cast(dict[str, Any] | str, result.model_dump())
+
+    # Already a dict or primitive - return as-is
+    return cast(dict[str, Any] | str, result)
+
+
+def serialize_agent_result_json(result: Any) -> str:
+    """
+    Safely serialize an agent result to JSON string, handling Pydantic models correctly.
+
+    Use this variant when you need a JSON string output (e.g., for SSE streaming,
+    JSON responses, or storage).
+
+    Args:
+        result: Agent result which may be:
+            - Pydantic model instance (has .model_dump_json())
+            - Dict or other JSON-serializable object
+            - Primitive type
+
+    Returns:
+        JSON string representation
+
+    Examples:
+        >>> # With Pydantic model result
+        >>> agent_result = await agent.run(query)
+        >>> json_str = serialize_agent_result_json(agent_result.output)
+        >>> return Response(content=json_str, media_type="application/json")
+
+        >>> # For SSE streaming
+        >>> chunk = serialize_agent_result_json(result.output)
+        >>> yield f"data: {chunk}\\n\\n"
+    """
+    import json
+
+    # Check if this is a Pydantic model instance with model_dump_json
+    if isinstance(result, BaseModel):
+        return result.model_dump_json()
+
+    # Check if this has a model_dump_json method (duck typing)
+    if hasattr(result, "model_dump_json") and callable(
+        getattr(result, "model_dump_json")
+    ):
+        return cast(str, result.model_dump_json())
+
+    # Fall back to standard json.dumps
+    return json.dumps(result)
+
+
+def is_pydantic_model(obj: Any) -> bool:
+    """
+    Check if an object is a Pydantic model instance.
+
+    Args:
+        obj: Object to check
+
+    Returns:
+        True if object is a Pydantic model instance
+
+    Examples:
+        >>> from pydantic import BaseModel
+        >>> class MyModel(BaseModel):
+        ...     value: str
+        >>> instance = MyModel(value="test")
+        >>> assert is_pydantic_model(instance) == True
+        >>> assert is_pydantic_model({"value": "test"}) == False
+    """
+    return isinstance(obj, BaseModel) or (
+        hasattr(obj, "model_dump") and hasattr(obj, "model_fields")
+    )
+
+
+def safe_serialize_dict(data: dict[str, Any]) -> dict[str, Any]:
+    """
+    Recursively serialize a dict that may contain Pydantic models.
+
+    Use this when you have a dict that may contain Pydantic model instances
+    nested within it (e.g., as values).
+
+    Args:
+        data: Dict that may contain Pydantic models
+
+    Returns:
+        Dict with all Pydantic models serialized to dicts
+
+    Examples:
+        >>> # Dict with nested Pydantic model
+        >>> data = {
+        ...     "status": "success",
+        ...     "result": some_pydantic_model,  # Will be serialized
+        ...     "metadata": {"count": 5}
+        ... }
+        >>> serialized = safe_serialize_dict(data)
+        >>> # All Pydantic models are now dicts
+    """
+    result = {}
+    for key, value in data.items():
+        if is_pydantic_model(value):
+            result[key] = serialize_agent_result(value)
+        elif isinstance(value, dict):
+            result[key] = safe_serialize_dict(value)
+        elif isinstance(value, list):
+            result[key] = [
+                serialize_agent_result(item) if is_pydantic_model(item) else item
+                for item in value
+            ]
+        else:
+            result[key] = value
+    return result
+
+
+# Example usage patterns for documentation
+USAGE_EXAMPLES = """
+# Example 1: MCP Tool returning agent result
+async def ask_rem_tool(query: str) -> dict[str, Any]:
+    from rem.agentic.serialization import serialize_agent_result
+
+    agent = await create_agent()
+    result = await agent.run(query)
+
+    # ✅ GOOD: Serialize before returning
+    return {
+        "status": "success",
+        "response": serialize_agent_result(result.output),
+        "model": result.model,
+    }
+
+# Example 2: API endpoint with Pydantic result
+@app.post("/query")
+async def query_endpoint(body: QueryRequest):
+    from rem.agentic.serialization import serialize_agent_result
+
+    agent = await create_agent()
+    result = await agent.run(body.query)
+
+    # ✅ GOOD: Serialize for FastAPI response
+    return {
+        "data": serialize_agent_result(result.output),
+        "usage": result.usage().model_dump() if result.usage() else None,
+    }
+
+# Example 3: Streaming with SSE
+async def stream_results(agent, query):
+    from rem.agentic.serialization import serialize_agent_result_json
+
+    async with agent.iter(query) as run:
+        async for event in run:
+            if isinstance(event, SomeEvent):
+                # ✅ GOOD: Serialize to JSON string for SSE
+                json_str = serialize_agent_result_json(event.data)
+                yield f"data: {json_str}\\n\\n"
+
+# Example 4: Service layer returning to MCP tool
+async def ask_rem(query: str, tenant_id: str) -> dict[str, Any]:
+    from rem.agentic.serialization import serialize_agent_result
+
+    agent = await create_agent()
+    result = await agent.run(query)
+
+    # ✅ GOOD: Serialize in service layer
+    return {
+        "query_output": serialize_agent_result(result.data),
+        "natural_query": query,
+    }
+"""

rem/agentic/tools/__init__.py

@@ -0,0 +1,5 @@
+"""Agent tools for REM operations."""
+
+from .rem_tools import search_rem_tool, ingest_file_tool
+
+__all__ = ["search_rem_tool", "ingest_file_tool"]

rem/agentic/tools/rem_tools.py

@@ -0,0 +1,231 @@
+"""
+REM tools for agent execution (CLI and API compatible).
+
+These tools work in both CLI and API contexts by initializing services on-demand.
+They wrap the service layer directly, not MCP tools.
+"""
+
+from typing import Any, Literal, cast
+
+from loguru import logger
+
+from ...models.core import (
+    FuzzyParameters,
+    LookupParameters,
+    QueryType,
+    RemQuery,
+    SearchParameters,
+    SQLParameters,
+    TraverseParameters,
+)
+from ...services.content import ContentService
+from ...services.postgres import get_postgres_service
+from ...services.rem import RemService
+
+
+# Service cache for reuse within agent execution
+_service_cache: dict[str, Any] = {}
+
+
+async def _get_rem_service() -> RemService:
+    """Get or create RemService instance."""
+    if "rem_service" not in _service_cache:
+        db = get_postgres_service()
+        if not db:
+            raise RuntimeError("PostgreSQL is disabled. Cannot use REM service.")
+
+        await db.connect()
+        _service_cache["postgres"] = db
+        _service_cache["rem_service"] = RemService(postgres_service=db)
+        logger.debug("Initialized RemService for agent tools")
+    return cast(RemService, _service_cache["rem_service"])
+
+
+async def search_rem_tool(
+    query_type: Literal["lookup", "fuzzy", "search", "sql", "traverse"],
+    user_id: str,
+    # LOOKUP parameters
+    entity_key: str | None = None,
+    # FUZZY parameters
+    query_text: str | None = None,
+    threshold: float = 0.7,
+    # SEARCH parameters
+    table: str | None = None,
+    limit: int = 20,
+    # SQL parameters
+    sql_query: str | None = None,
+    # TRAVERSE parameters
+    initial_query: str | None = None,
+    edge_types: list[str] | None = None,
+    depth: int = 1,
+) -> dict[str, Any]:
+    """
+    Execute REM queries for entity lookup, semantic search, and graph traversal.
+
+    This tool works in both CLI and API contexts by initializing services on-demand.
+
+    Args:
+        query_type: Type of query (lookup, fuzzy, search, sql, traverse)
+        user_id: User identifier for data scoping
+        entity_key: Entity key for LOOKUP (e.g., "Sarah Chen")
+        query_text: Search text for FUZZY or SEARCH
+        threshold: Similarity threshold for FUZZY (0.0-1.0)
+        table: Target table for SEARCH (resources, moments, users, etc.)
+        limit: Max results for SEARCH
+        sql_query: SQL query string for SQL type
+        initial_query: Starting entity for TRAVERSE
+        edge_types: Edge types to follow for TRAVERSE
+        depth: Traversal depth for TRAVERSE
+
+    Returns:
+        Dict with query results and metadata
+    """
+    try:
+        rem_service = await _get_rem_service()
+
+        # Build RemQuery based on query_type
+        if query_type == "lookup":
+            if not entity_key:
+                return {"status": "error", "error": "entity_key required for LOOKUP"}
+
+            query = RemQuery(
+                query_type=QueryType.LOOKUP,
+                parameters=LookupParameters(
+                    key=entity_key,
+                    user_id=user_id,
+                ),
+                user_id=user_id,
+            )
+
+        elif query_type == "fuzzy":
+            if not query_text:
+                return {"status": "error", "error": "query_text required for FUZZY"}
+
+            query = RemQuery(
+                query_type=QueryType.FUZZY,
+                parameters=FuzzyParameters(
+                    query_text=query_text,
+                    threshold=threshold,
+                    limit=limit, # Implied parameter
+                ),
+                user_id=user_id,
+            )
+
+        elif query_type == "search":
+            if not query_text:
+                return {"status": "error", "error": "query_text required for SEARCH"}
+            if not table:
+                return {"status": "error", "error": "table required for SEARCH"}
+
+            query = RemQuery(
+                query_type=QueryType.SEARCH,
+                parameters=SearchParameters(
+                    query_text=query_text,
+                    table_name=table,
+                    limit=limit,
+                ),
+                user_id=user_id,
+            )
+
+        elif query_type == "sql":
+            if not sql_query:
+                return {"status": "error", "error": "sql_query required for SQL"}
+            
+            if not table:
+                 return {"status": "error", "error": "table required for SQL queries"}
+
+            query = RemQuery(
+                query_type=QueryType.SQL,
+                parameters=SQLParameters(
+                    table_name=table,
+                    where_clause=sql_query,
+                    limit=limit, # SQLParams accepts limit
+                ),
+                user_id=user_id,
+            )
+
+        elif query_type == "traverse":
+            if not initial_query:
+                return {"status": "error", "error": "initial_query required for TRAVERSE"}
+
+            query = RemQuery(
+                query_type=QueryType.TRAVERSE,
+                parameters=TraverseParameters(
+                    initial_query=initial_query,
+                    edge_types=edge_types or [],
+                    max_depth=depth,
+                ),
+                user_id=user_id,
+            )
+
+        else:
+            return {"status": "error", "error": f"Unknown query_type: {query_type}"}
+
+        # Execute query
+        logger.info(f"Executing REM query: {query_type} for user {user_id}")
+        result = await rem_service.execute_query(query)
+
+        logger.info(f"Query completed: {query_type}")
+        return {
+            "status": "success",
+            "query_type": query_type,
+            "results": result,
+        }
+
+    except Exception as e:
+        logger.error(f"search_rem_tool failed: {e}", exc_info=True)
+        return {
+            "status": "error",
+            "error": str(e),
+        }
+
+
+async def ingest_file_tool(
+    file_uri: str,
+    user_id: str,
+    category: str | None = None,
+    tags: list[str] | None = None,
+    is_local_server: bool = True,  # CLI is always local
+) -> dict[str, Any]:
+    """
+    Ingest file into REM (read + store + parse + chunk + embed).
+
+    This tool works in both CLI and API contexts.
+
+    Args:
+        file_uri: File location (local path, s3:// URI, or http(s):// URL)
+        user_id: User identifier for data scoping
+        category: Optional category (document, code, audio, etc.)
+        tags: Optional tags for file
+        is_local_server: True if running as local CLI (default)
+
+    Returns:
+        Dict with file_id, processing_status, resources_created, etc.
+    """
+    try:
+        content_service = ContentService()
+        result = await content_service.ingest_file(
+            file_uri=file_uri,
+            user_id=user_id,
+            category=category,
+            tags=tags,
+            is_local_server=is_local_server,
+        )
+
+        logger.info(
+            f"File ingestion complete: {result['file_name']} "
+            f"(status: {result['processing_status']}, "
+            f"resources: {result['resources_created']})"
+        )
+
+        return {
+            "status": "success",
+            **result,
+        }
+
+    except Exception as e:
+        logger.error(f"ingest_file_tool failed: {e}", exc_info=True)
+        return {
+            "status": "error",
+            "error": str(e),
+        }