remdb 0.3.202__py3-none-any.whl → 0.3.245__py3-none-any.whl

rem/api/routers/shared_sessions.py

@@ -18,6 +18,8 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request
 from loguru import logger
 from pydantic import BaseModel, Field
 
+from .common import ErrorResponse
+
 from ..deps import get_current_user, require_auth
 from ...models.entities import (
     Message,
@@ -83,6 +85,10 @@ class ShareSessionResponse(BaseModel):
     response_model=ShareSessionResponse,
     status_code=201,
     tags=["sessions"],
+    responses={
+        400: {"model": ErrorResponse, "description": "Session already shared with this user"},
+        503: {"model": ErrorResponse, "description": "Database not enabled"},
+    },
 )
 async def share_session(
     request: Request,
@@ -175,6 +181,10 @@ async def share_session(
     "/sessions/{session_id}/share/{shared_with_user_id}",
     status_code=200,
     tags=["sessions"],
+    responses={
+        404: {"model": ErrorResponse, "description": "Share not found"},
+        503: {"model": ErrorResponse, "description": "Database not enabled"},
+    },
 )
 async def remove_session_share(
     request: Request,
@@ -250,6 +260,9 @@ async def remove_session_share(
     "/sessions/shared-with-me",
     response_model=SharedWithMeResponse,
     tags=["sessions"],
+    responses={
+        503: {"model": ErrorResponse, "description": "Database not enabled"},
+    },
 )
 async def get_shared_with_me(
     request: Request,
@@ -328,6 +341,9 @@ async def get_shared_with_me(
     "/sessions/shared-with-me/{owner_user_id}/messages",
     response_model=SharedMessagesResponse,
     tags=["sessions"],
+    responses={
+        503: {"model": ErrorResponse, "description": "Database not enabled"},
+    },
 )
 async def get_shared_messages(
     request: Request,

rem/cli/commands/ask.py

@@ -71,16 +71,18 @@ async def run_agent_streaming(
     max_turns: int = 10,
     context: AgentContext | None = None,
     max_iterations: int | None = None,
+    user_message: str | None = None,
 ) -> None:
     """
-    Run agent in streaming mode using agent.iter() with usage limits.
+    Run agent in streaming mode using the SAME code path as the API.
 
-    Design Pattern:
-    - Use agent.iter() for complete execution with tool call visibility
-    - run_stream() stops after first output, missing tool calls
-    - Stream tool call markers: [Calling: tool_name]
-    - Stream text content deltas as they arrive
-    - Show final structured result
+    This uses stream_openai_response_with_save from the API to ensure:
+    1. Tool calls are saved as separate "tool" messages (not embedded in content)
+    2. Assistant response is clean text only (no [Calling: ...] markers)
+    3. CLI testing is equivalent to API testing
+
+    The CLI displays tool calls as [Calling: tool_name] for visibility,
+    but these are NOT saved to the database.
 
     Args:
         agent: Pydantic AI agent
@@ -88,88 +90,66 @@ async def run_agent_streaming(
         max_turns: Maximum turns for agent execution (not used in current API)
         context: Optional AgentContext for session persistence
         max_iterations: Maximum iterations/requests (from agent schema or settings)
+        user_message: The user's original message (for database storage)
     """
-    from pydantic_ai import UsageLimits
-    from rem.utils.date_utils import to_iso_with_z, utc_now
+    import json
+    from rem.api.routers.chat.streaming import stream_openai_response_with_save, save_user_message
 
     logger.info("Running agent in streaming mode...")
 
     try:
-        # Import event types for streaming
-        from pydantic_ai import Agent as PydanticAgent
-        from pydantic_ai.messages import PartStartEvent, PartDeltaEvent, TextPartDelta, ToolCallPart
-
-        # Accumulate assistant response for session persistence
-        assistant_response_parts = []
-
-        # Use agent.iter() to get complete execution with tool calls
-        usage_limits = UsageLimits(request_limit=max_iterations) if max_iterations else None
-        async with agent.iter(prompt, usage_limits=usage_limits) as agent_run:
-            async for node in agent_run:
-                # Check if this is a model request node (includes tool calls and text)
-                if PydanticAgent.is_model_request_node(node):
-                    # Stream events from model request
-                    request_stream: Any
-                    async with node.stream(agent_run.ctx) as request_stream:
-                        async for event in request_stream:
-                            # Tool call start event
-                            if isinstance(event, PartStartEvent) and isinstance(
-                                event.part, ToolCallPart
-                            ):
-                                tool_marker = f"\n[Calling: {event.part.tool_name}]"
-                                print(tool_marker, flush=True)
-                                assistant_response_parts.append(tool_marker)
-
-                            # Text content delta
-                            elif isinstance(event, PartDeltaEvent) and isinstance(
-                                event.delta, TextPartDelta
-                            ):
-                                print(event.delta.content_delta, end="", flush=True)
-                                assistant_response_parts.append(event.delta.content_delta)
-
-        print("\n")  # Final newline after streaming
-
-        # Get final result from agent_run
-        result = agent_run.result
-        if hasattr(result, "output"):
-            logger.info("Final structured result:")
-            output = result.output
-            from rem.agentic.serialization import serialize_agent_result
-            output_json = json.dumps(serialize_agent_result(output), indent=2)
-            print(output_json)
-            assistant_response_parts.append(f"\n{output_json}")
-
-        # Save session messages (if session_id provided and postgres enabled)
-        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
-
-            user_message = {
-                "role": "user",
-                "content": user_message_content,
-                "timestamp": to_iso_with_z(utc_now()),
-            }
-
-            assistant_message = {
-                "role": "assistant",
-                "content": "".join(assistant_response_parts),
-                "timestamp": to_iso_with_z(utc_now()),
-            }
-
-            # Store messages with compression
-            store = SessionMessageStore(user_id=context.user_id or settings.test.effective_user_id)
-            await store.store_session_messages(
+        # Save user message BEFORE streaming (same as API, using shared utility)
+        if context and context.session_id and user_message:
+            await save_user_message(
                 session_id=context.session_id,
-                messages=[user_message, assistant_message],
                 user_id=context.user_id,
-                compress=True,
+                content=user_message,
             )
 
-            logger.debug(f"Saved conversation to session {context.session_id}")
+        # Use the API streaming code path for consistency
+        # This properly handles tool calls and message persistence
+        model_name = getattr(agent, 'model', 'unknown')
+        if hasattr(model_name, 'model_name'):
+            model_name = model_name.model_name
+        elif hasattr(model_name, 'name'):
+            model_name = model_name.name
+        else:
+            model_name = str(model_name)
+
+        async for chunk in stream_openai_response_with_save(
+            agent=agent.agent if hasattr(agent, 'agent') else agent,
+            prompt=prompt,
+            model=model_name,
+            session_id=context.session_id if context else None,
+            user_id=context.user_id if context else None,
+            agent_context=context,
+        ):
+            # Parse SSE chunks for CLI display
+            if chunk.startswith("event: tool_call"):
+                # Extract tool call info from next data line
+                continue
+            elif chunk.startswith("data: ") and not chunk.startswith("data: [DONE]"):
+                try:
+                    data_str = chunk[6:].strip()
+                    if data_str:
+                        data = json.loads(data_str)
+                        # Check for tool_call event
+                        if data.get("type") == "tool_call":
+                            tool_name = data.get("tool_name", "tool")
+                            status = data.get("status", "")
+                            if status == "started":
+                                print(f"\n[Calling: {tool_name}]", flush=True)
+                        # Check for text content (OpenAI format)
+                        elif "choices" in data and data["choices"]:
+                            delta = data["choices"][0].get("delta", {})
+                            content = delta.get("content")
+                            if content:
+                                print(content, end="", flush=True)
+                except (json.JSONDecodeError, KeyError, IndexError):
+                    pass
+
+        print("\n")  # Final newline after streaming
+        logger.info("Final structured result:")
 
     except Exception as e:
         logger.error(f"Agent execution failed: {e}")
@@ -184,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
@@ -196,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
 
@@ -352,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",
@@ -549,7 +639,7 @@ async def _ask_async(
 
     # Run agent with session persistence
     if stream:
-        await run_agent_streaming(agent, prompt, max_turns=max_turns, context=context)
+        await run_agent_streaming(agent, prompt, max_turns=max_turns, context=context, user_message=query)
     else:
         await run_agent_non_streaming(
             agent,
@@ -558,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/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 = {
@@ -206,9 +214,9 @@ def process_ingest(
                 if category:
                     entity_data["category"] = category
 
-                # Scoping: user_id for private data, None for public/shared
-                # tenant_id=None and user_id=None means PUBLIC data (visible to all)
-                entity_data["tenant_id"] = user_id  # None = public/shared
+                # Scoping: user_id for private data, "public" for shared
+                # tenant_id="public" is the default for shared knowledge bases
+                entity_data["tenant_id"] = user_id or "public"
                 entity_data["user_id"] = user_id  # None = public/shared
 
                 # For ontologies, add URI

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)
+
+