remdb 0.3.230__py3-none-any.whl → 0.3.258__py3-none-any.whl

rem/api/routers/chat/streaming.py

@@ -404,52 +404,47 @@ async def stream_openai_response(
                                     tool_calls_out.append(tool_data)
                                     del state.pending_tool_data[tool_id]
 
-                                if not is_metadata_event:
-                                    # NOTE: text_response fallback is DISABLED
-                                    # Child agents now stream content via child_content events (above)
-                                    # which provides real-time streaming. The text_response in tool
-                                    # result would duplicate that content, so we skip it entirely.
-
-                                    # Normal tool completion - emit ToolCallEvent
-                                    # For finalize_intake, send full result dict for frontend
-                                    if tool_name == "finalize_intake" and isinstance(result_content, dict):
-                                        result_for_sse = result_content
+                                # Always emit ToolCallEvent completion for frontend tracking
+                                # Send full result for dict/list types, stringify others
+                                if isinstance(result_content, (dict, list)):
+                                    result_for_sse = result_content
+                                else:
+                                    result_for_sse = str(result_content)
+
+                                # Log result count for search_rem
+                                if tool_name == "search_rem" and isinstance(result_content, dict):
+                                    results = result_content.get("results", {})
+                                    # Handle nested result structure: results may be a dict with 'results' list and 'count'
+                                    if isinstance(results, dict):
+                                        count = results.get("count", len(results.get("results", [])))
+                                        query_type = results.get("query_type", "?")
+                                        query_text = results.get("query_text", results.get("key", ""))
+                                        table = results.get("table_name", "")
+                                    elif isinstance(results, list):
+                                        count = len(results)
+                                        query_type = "?"
+                                        query_text = ""
+                                        table = ""
                                     else:
-                                        result_str = str(result_content)
-                                        result_for_sse = result_str[:200] + "..." if len(result_str) > 200 else result_str
-
-                                    # Log result count for search_rem
-                                    if tool_name == "search_rem" and isinstance(result_content, dict):
-                                        results = result_content.get("results", {})
-                                        # Handle nested result structure: results may be a dict with 'results' list and 'count'
-                                        if isinstance(results, dict):
-                                            count = results.get("count", len(results.get("results", [])))
-                                            query_type = results.get("query_type", "?")
-                                            query_text = results.get("query_text", results.get("key", ""))
-                                            table = results.get("table_name", "")
-                                        elif isinstance(results, list):
-                                            count = len(results)
-                                            query_type = "?"
-                                            query_text = ""
-                                            table = ""
-                                        else:
-                                            count = "?"
-                                            query_type = "?"
-                                            query_text = ""
-                                            table = ""
-                                        status = result_content.get("status", "unknown")
-                                        # Truncate query text for logging
-                                        if query_text and len(str(query_text)) > 40:
-                                            query_text = str(query_text)[:40] + "..."
-                                        logger.info(f"  ↳ {tool_name} {query_type} '{query_text}' table={table} → {count} results")
-
-                                    yield format_sse_event(ToolCallEvent(
-                                        tool_name=tool_name,
-                                        tool_id=tool_id,
-                                        status="completed",
-                                        arguments=completed_args,
-                                        result=result_for_sse
-                                    ))
+                                        count = "?"
+                                        query_type = "?"
+                                        query_text = ""
+                                        table = ""
+                                    status = result_content.get("status", "unknown")
+                                    # Truncate query text for logging
+                                    if query_text and len(str(query_text)) > 40:
+                                        query_text = str(query_text)[:40] + "..."
+                                    logger.info(f"  ↳ {tool_name} {query_type} '{query_text}' table={table} → {count} results")
+
+                                # Always emit ToolCallEvent completion for frontend tracking
+                                # This includes register_metadata calls so they turn green in the UI
+                                yield format_sse_event(ToolCallEvent(
+                                    tool_name=tool_name,
+                                    tool_id=tool_id,
+                                    status="completed",
+                                    arguments=completed_args,
+                                    result=result_for_sse
+                                ))
 
                                 # Update progress after tool completion
                                 state.current_step = 3

rem/api/routers/chat/streaming_utils.py

@@ -177,19 +177,17 @@ def build_tool_complete_event(
     arguments: dict | None = None,
     result: Any = None,
 ) -> str:
-    """Build a tool call completed SSE event."""
-    result_str = None
-    if result is not None:
-        result_str = str(result)
-        if len(result_str) > 200:
-            result_str = result_str[:200] + "..."
+    """Build a tool call completed SSE event.
 
+    Note: Full result is sent in SSE events for UI display.
+    Truncation only happens in log_tool_result() for log readability.
+    """
     return format_sse_event(ToolCallEvent(
         tool_name=tool_name,
         tool_id=tool_id,
         status="completed",
         arguments=arguments,
-        result=result_str,
+        result=result,
     ))
 
 

rem/api/routers/feedback.py

@@ -21,11 +21,11 @@ IMPORTANT - Testing Requirements:
     ║  2. Session IDs MUST be UUIDs - use python3 -c "import uuid; print(uuid.uuid4())"                  ║
     ║  3. Port-forward OTEL collector: kubectl port-forward -n observability                             ║
     ║       svc/otel-collector-collector 4318:4318                                                       ║
-    ║  4. Port-forward Phoenix: kubectl port-forward -n siggy svc/phoenix 6006:6006                      ║
+    ║  4. Port-forward Phoenix: kubectl port-forward -n rem svc/phoenix 6006:6006                        ║
     ║  5. Set environment variables when starting the API:                                               ║
     ║       OTEL__ENABLED=true PHOENIX__ENABLED=true PHOENIX_API_KEY=<jwt> uvicorn ...                   ║
     ║  6. Get PHOENIX_API_KEY:                                                                           ║
-    ║       kubectl get secret -n siggy rem-phoenix-api-key -o jsonpath='{.data.PHOENIX_API_KEY}'        ║
+    ║       kubectl get secret -n rem rem-phoenix-api-key -o jsonpath='{.data.PHOENIX_API_KEY}'          ║
     ║         | base64 -d                                                                                ║
     ╚════════════════════════════════════════════════════════════════════════════════════════════════════╝
 

rem/api/routers/query.py

@@ -90,8 +90,6 @@ from .common import ErrorResponse
 
 from ...services.postgres import get_postgres_service
 from ...services.rem.service import RemService
-from ...services.rem.parser import RemQueryParser
-from ...models.core import RemQuery
 from ...settings import settings
 
 router = APIRouter(prefix="/api/v1", tags=["query"])
@@ -331,7 +329,7 @@ async def execute_query(
             return response
 
         else:
-            # REM dialect mode - parse and execute directly
+            # REM dialect mode - use unified execute_query_string
             if not request.query:
                 raise HTTPException(
                     status_code=400,
@@ -340,17 +338,10 @@ async def execute_query(
 
             logger.info(f"REM dialect query: {request.query[:100]}...")
 
-            parser = RemQueryParser()
-            query_type, parameters = parser.parse(request.query)
-
-            # Create and execute RemQuery
-            rem_query = RemQuery.model_validate({
-                "query_type": query_type,
-                "parameters": parameters,
-                "user_id": effective_user_id,
-            })
-
-            result = await rem_service.execute_query(rem_query)
+            # Use the unified execute_query_string method
+            result = await rem_service.execute_query_string(
+                request.query, user_id=effective_user_id
+            )
 
             return QueryResponse(
                 query_type=result["query_type"],

rem/cli/commands/ask.py

@@ -164,9 +164,13 @@ async def run_agent_non_streaming(
     context: AgentContext | None = None,
     plan: bool = False,
     max_iterations: int | None = None,
+    user_message: str | None = None,
 ) -> dict[str, Any] | None:
     """
-    Run agent in non-streaming mode using agent.run() with usage limits.
+    Run agent in non-streaming mode using agent.iter() to capture tool calls.
+
+    This mirrors the streaming code path to ensure tool messages are properly
+    persisted to the database for state tracking across turns.
 
     Args:
         agent: Pydantic AI agent
@@ -176,77 +180,183 @@ async def run_agent_non_streaming(
         context: Optional AgentContext for session persistence
         plan: If True, output only the generated query (for query-agent)
         max_iterations: Maximum iterations/requests (from agent schema or settings)
+        user_message: The user's original message (for database storage)
 
     Returns:
         Output data if successful, None otherwise
     """
     from pydantic_ai import UsageLimits
+    from pydantic_ai.agent import Agent
+    from pydantic_ai.messages import (
+        FunctionToolResultEvent,
+        PartStartEvent,
+        PartEndEvent,
+        TextPart,
+        ToolCallPart,
+    )
     from rem.utils.date_utils import to_iso_with_z, utc_now
 
     logger.info("Running agent in non-streaming mode...")
 
     try:
-        # Run agent and get complete result with usage limits
-        usage_limits = UsageLimits(request_limit=max_iterations) if max_iterations else None
-        result = await agent.run(prompt, usage_limits=usage_limits)
+        # Track tool calls for persistence (same as streaming code path)
+        tool_calls: list = []
+        pending_tool_data: dict = {}
+        pending_tool_completions: list = []
+        accumulated_content: list = []
+
+        # Get the underlying pydantic-ai agent
+        pydantic_agent = agent.agent if hasattr(agent, 'agent') else agent
+
+        # Use agent.iter() to capture tool calls (same as streaming)
+        async with pydantic_agent.iter(prompt) as agent_run:
+            async for node in agent_run:
+                # Handle model request nodes (text + tool call starts)
+                if Agent.is_model_request_node(node):
+                    async with node.stream(agent_run.ctx) as request_stream:
+                        async for event in request_stream:
+                            # Capture text content
+                            if isinstance(event, PartStartEvent) and isinstance(event.part, TextPart):
+                                if event.part.content:
+                                    accumulated_content.append(event.part.content)
+
+                            # Capture tool call starts
+                            elif isinstance(event, PartStartEvent) and isinstance(event.part, ToolCallPart):
+                                tool_name = event.part.tool_name
+                                if tool_name == "final_result":
+                                    continue
+
+                                import uuid
+                                tool_id = f"call_{uuid.uuid4().hex[:8]}"
+                                pending_tool_completions.append((tool_name, tool_id))
+
+                                # Extract arguments
+                                args_dict = {}
+                                if hasattr(event.part, 'args'):
+                                    args = event.part.args
+                                    if isinstance(args, str):
+                                        try:
+                                            args_dict = json.loads(args)
+                                        except json.JSONDecodeError:
+                                            args_dict = {"raw": args}
+                                    elif isinstance(args, dict):
+                                        args_dict = args
+
+                                pending_tool_data[tool_id] = {
+                                    "tool_name": tool_name,
+                                    "tool_id": tool_id,
+                                    "arguments": args_dict,
+                                }
+
+                                # Print tool call for CLI visibility
+                                print(f"\n[Calling: {tool_name}]", flush=True)
+
+                            # Capture tool call end (update arguments if changed)
+                            elif isinstance(event, PartEndEvent) and isinstance(event.part, ToolCallPart):
+                                pass  # Arguments already captured at start
+
+                # Handle tool execution nodes (results)
+                elif Agent.is_call_tools_node(node):
+                    async with node.stream(agent_run.ctx) as tools_stream:
+                        async for event in tools_stream:
+                            if isinstance(event, FunctionToolResultEvent):
+                                # Get tool info from pending queue
+                                if pending_tool_completions:
+                                    tool_name, tool_id = pending_tool_completions.pop(0)
+                                else:
+                                    import uuid
+                                    tool_name = "tool"
+                                    tool_id = f"call_{uuid.uuid4().hex[:8]}"
+
+                                result_content = event.result.content if hasattr(event.result, 'content') else event.result
+
+                                # Capture tool call for persistence
+                                if tool_id in pending_tool_data:
+                                    tool_data = pending_tool_data[tool_id]
+                                    tool_data["result"] = result_content
+                                    tool_calls.append(tool_data)
+                                    del pending_tool_data[tool_id]
+
+            # Get final result
+            result = agent_run.result
 
         # Extract output data
         output_data = None
         assistant_content = None
-        if hasattr(result, "output"):
+        if result is not None and hasattr(result, "output"):
             output = result.output
             from rem.agentic.serialization import serialize_agent_result
             output_data = serialize_agent_result(output)
 
             if plan and isinstance(output_data, dict) and "query" in output_data:
-                # Plan mode: Output only the query
-                # Use sql formatting if possible or just raw string
                 assistant_content = output_data["query"]
                 print(assistant_content)
             else:
-                # Normal mode
-                assistant_content = json.dumps(output_data, indent=2)
+                # For string output, use it directly
+                if isinstance(output_data, str):
+                    assistant_content = output_data
+                else:
+                    assistant_content = json.dumps(output_data, indent=2)
                 print(assistant_content)
         else:
-            # Fallback for text-only results
-            assistant_content = str(result)
-            print(assistant_content)
+            assistant_content = str(result) if result else ""
+            if assistant_content:
+                print(assistant_content)
 
         # Save to file if requested
         if output_file and output_data:
             await _save_output_file(output_file, output_data)
 
-        # Save session messages (if session_id provided and postgres enabled)
+        # Save session messages including tool calls (same as streaming code path)
         if context and context.session_id and settings.postgres.enabled:
             from ...services.session.compression import SessionMessageStore
 
-            # Extract just the user query from prompt
-            # Prompt format from ContextBuilder: system + history + user message
-            # We need to extract the last user message
-            user_message_content = prompt.split("\n\n")[-1] if "\n\n" in prompt else prompt
+            timestamp = to_iso_with_z(utc_now())
+            messages_to_store = []
 
-            user_message = {
+            # Save user message first
+            user_message_content = user_message or (prompt.split("\n\n")[-1] if "\n\n" in prompt else prompt)
+            messages_to_store.append({
                 "role": "user",
                 "content": user_message_content,
-                "timestamp": to_iso_with_z(utc_now()),
-            }
-
-            assistant_message = {
-                "role": "assistant",
-                "content": assistant_content,
-                "timestamp": to_iso_with_z(utc_now()),
-            }
-
-            # Store messages with compression
+                "timestamp": timestamp,
+            })
+
+            # Save tool call messages (message_type: "tool") - CRITICAL for state tracking
+            for tool_call in tool_calls:
+                if not tool_call:
+                    continue
+                tool_message = {
+                    "role": "tool",
+                    "content": json.dumps(tool_call.get("result", {}), default=str),
+                    "timestamp": timestamp,
+                    "tool_call_id": tool_call.get("tool_id"),
+                    "tool_name": tool_call.get("tool_name"),
+                    "tool_arguments": tool_call.get("arguments"),
+                }
+                messages_to_store.append(tool_message)
+
+            # Save assistant message
+            if assistant_content:
+                messages_to_store.append({
+                    "role": "assistant",
+                    "content": assistant_content,
+                    "timestamp": timestamp,
+                })
+
+            # Store all messages
             store = SessionMessageStore(user_id=context.user_id or settings.test.effective_user_id)
             await store.store_session_messages(
                 session_id=context.session_id,
-                messages=[user_message, assistant_message],
+                messages=messages_to_store,
                 user_id=context.user_id,
-                compress=True,
+                compress=False,  # Store uncompressed; compression happens on reload
             )
 
-            logger.debug(f"Saved conversation to session {context.session_id}")
+            logger.debug(
+                f"Saved {len(tool_calls)} tool calls + user/assistant messages "
+                f"to session {context.session_id}"
+            )
 
         return output_data
 
@@ -332,8 +442,8 @@ async def _save_output_file(file_path: Path, data: dict[str, Any]) -> None:
 )
 @click.option(
     "--stream/--no-stream",
-    default=False,
-    help="Enable streaming mode (default: disabled)",
+    default=True,
+    help="Enable streaming mode (default: enabled)",
 )
 @click.option(
     "--user-id",
@@ -538,6 +648,7 @@ async def _ask_async(
             output_file=output_file,
             context=context,
             plan=plan,
+            user_message=query,
         )
 
     # Log session ID for reuse

rem/cli/commands/experiments.py

@@ -1568,7 +1568,7 @@ def export(
         rem experiments export my-experiment
 
         # Export to specific bucket
-        rem experiments export my-experiment --bucket siggy-data
+        rem experiments export my-experiment --bucket my-data-lake
 
         # Include results in export
         rem experiments export my-experiment --include-results

rem/cli/commands/process.py

@@ -193,7 +193,15 @@ def process_ingest(
             try:
                 # Read file content
                 content = file_path.read_text(encoding="utf-8")
-                entity_key = file_path.stem  # filename without extension
+
+                # Generate entity key from filename
+                # Special case: README files use parent directory as section name
+                if file_path.stem.lower() == "readme":
+                    # Use parent directory name, e.g., "drugs" for drugs/README.md
+                    # For nested paths like disorders/anxiety/README.md -> "anxiety"
+                    entity_key = file_path.parent.name
+                else:
+                    entity_key = file_path.stem  # filename without extension
 
                 # Build entity based on table
                 entity_data = {

rem/cli/commands/query.py

@@ -0,0 +1,109 @@
+"""
+REM query command.
+
+Usage:
+    rem query --sql 'LOOKUP "Sarah Chen"'
+    rem query --sql 'SEARCH resources "API design" LIMIT 10'
+    rem query --sql "SELECT * FROM resources LIMIT 5"
+    rem query --file queries/my_query.sql
+
+This tool connects to the configured PostgreSQL instance and executes the
+provided REM dialect query, printing results as JSON (default) or plain dicts.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from pathlib import Path
+from typing import List
+
+import click
+from loguru import logger
+
+from ...services.rem import QueryExecutionError
+from ...services.rem.service import RemService
+
+
+@click.command("query")
+@click.option("--sql", "-s", default=None, help="REM query string (LOOKUP, SEARCH, FUZZY, TRAVERSE, or SQL)")
+@click.option(
+    "--file",
+    "-f",
+    "sql_file",
+    type=click.Path(exists=True, path_type=Path),
+    default=None,
+    help="Path to file containing REM query",
+)
+@click.option("--no-json", is_flag=True, default=False, help="Print rows as Python dicts instead of JSON")
+@click.option("--user-id", "-u", default=None, help="Scope query to a specific user")
+def query_command(sql: str | None, sql_file: Path | None, no_json: bool, user_id: str | None):
+    """
+    Execute a REM query against the database.
+
+    Supports REM dialect queries (LOOKUP, SEARCH, FUZZY, TRAVERSE) and raw SQL.
+    Either --sql or --file must be provided.
+    """
+    if not sql and not sql_file:
+        click.secho("Error: either --sql or --file is required", fg="red")
+        raise click.Abort()
+
+    # Read query from file if provided
+    if sql_file:
+        query_text = sql_file.read_text(encoding="utf-8")
+    else:
+        query_text = sql  # type: ignore[assignment]
+
+    try:
+        asyncio.run(_run_query_async(query_text, not no_json, user_id))
+    except Exception as exc:  # pragma: no cover - CLI error path
+        logger.exception("Query failed")
+        click.secho(f"✗ Query failed: {exc}", fg="red")
+        raise click.Abort()
+
+
+async def _run_query_async(query_text: str, as_json: bool, user_id: str | None) -> None:
+    """
+    Execute the query using RemService.execute_query_string().
+    """
+    from ...services.postgres import get_postgres_service
+
+    db = get_postgres_service()
+    if not db:
+        click.secho("✗ PostgreSQL is disabled in settings. Enable with POSTGRES__ENABLED=true", fg="red")
+        raise click.Abort()
+
+    if db.pool is None:
+        await db.connect()
+
+    rem_service = RemService(db)
+
+    try:
+        # Use the unified execute_query_string method
+        result = await rem_service.execute_query_string(query_text, user_id=user_id)
+        output_rows = result.get("results", [])
+    except QueryExecutionError as qe:
+        logger.exception("Query execution failed")
+        click.secho(f"✗ Query execution failed: {qe}. Please check the query you provided and try again.", fg="red")
+        raise click.Abort()
+    except ValueError as ve:
+        # Parse errors from the query parser
+        click.secho(f"✗ Invalid query: {ve}", fg="red")
+        raise click.Abort()
+    except Exception as exc:  # pragma: no cover - CLI error path
+        logger.exception("Unexpected error during query execution")
+        click.secho("✗ An unexpected error occurred while executing the query. Please check the query you provided and try again.", fg="red")
+        raise click.Abort()
+
+    if as_json:
+        click.echo(json.dumps(output_rows, default=str, indent=2))
+    else:
+        for r in output_rows:
+            click.echo(str(r))
+
+
+def register_command(cli_group):
+    """Register the query command on the given CLI group (top-level)."""
+    cli_group.add_command(query_command)
+
+