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

rem/api/routers/chat/streaming_utils.py

@@ -0,0 +1,327 @@
+"""
+Streaming Utilities.
+
+Pure functions and data structures for SSE streaming.
+No I/O, no database calls - just data transformation.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any
+
+from loguru import logger
+
+from .models import (
+    ChatCompletionMessageDelta,
+    ChatCompletionStreamChoice,
+    ChatCompletionStreamResponse,
+)
+from .sse_events import (
+    MetadataEvent,
+    ProgressEvent,
+    ReasoningEvent,
+    ToolCallEvent,
+    format_sse_event,
+)
+
+
+# =============================================================================
+# STREAMING STATE
+# =============================================================================
+
+@dataclass
+class StreamingState:
+    """
+    Tracks state during SSE streaming.
+
+    This is a pure data container - no methods that do I/O.
+    """
+    request_id: str
+    created_at: int
+    model: str
+    start_time: float = field(default_factory=time.time)
+
+    # Content tracking
+    is_first_chunk: bool = True
+    token_count: int = 0
+
+    # Child agent tracking - KEY FOR DUPLICATION FIX
+    child_content_streamed: bool = False
+    responding_agent: str | None = None
+
+    # Tool tracking
+    active_tool_calls: dict = field(default_factory=dict)  # index -> (name, id)
+    pending_tool_completions: list = field(default_factory=list)  # FIFO queue
+    pending_tool_data: dict = field(default_factory=dict)  # tool_id -> data
+
+    # Reasoning tracking
+    reasoning_step: int = 0
+
+    # Progress tracking
+    current_step: int = 0
+    total_steps: int = 3
+
+    # Metadata tracking
+    metadata_registered: bool = False
+
+    # Trace context (captured from OTEL)
+    trace_id: str | None = None
+    span_id: str | None = None
+
+    @classmethod
+    def create(cls, model: str, request_id: str | None = None) -> "StreamingState":
+        """Create a new streaming state."""
+        return cls(
+            request_id=request_id or f"chatcmpl-{uuid.uuid4().hex[:24]}",
+            created_at=int(time.time()),
+            model=model,
+        )
+
+    def latency_ms(self) -> int:
+        """Calculate latency since start."""
+        return int((time.time() - self.start_time) * 1000)
+
+
+# =============================================================================
+# SSE CHUNK BUILDERS
+# =============================================================================
+
+def build_content_chunk(state: StreamingState, content: str) -> str:
+    """
+    Build an SSE content chunk in OpenAI format.
+
+    Updates state.is_first_chunk and state.token_count.
+    """
+    state.token_count += len(content.split())
+
+    chunk = ChatCompletionStreamResponse(
+        id=state.request_id,
+        created=state.created_at,
+        model=state.model,
+        choices=[
+            ChatCompletionStreamChoice(
+                index=0,
+                delta=ChatCompletionMessageDelta(
+                    role="assistant" if state.is_first_chunk else None,
+                    content=content,
+                ),
+                finish_reason=None,
+            )
+        ],
+    )
+    state.is_first_chunk = False
+    return f"data: {chunk.model_dump_json()}\n\n"
+
+
+def build_final_chunk(state: StreamingState) -> str:
+    """Build the final SSE chunk with finish_reason=stop."""
+    chunk = ChatCompletionStreamResponse(
+        id=state.request_id,
+        created=state.created_at,
+        model=state.model,
+        choices=[
+            ChatCompletionStreamChoice(
+                index=0,
+                delta=ChatCompletionMessageDelta(),
+                finish_reason="stop",
+            )
+        ],
+    )
+    return f"data: {chunk.model_dump_json()}\n\n"
+
+
+def build_reasoning_event(state: StreamingState, content: str) -> str:
+    """Build a reasoning SSE event."""
+    return format_sse_event(ReasoningEvent(
+        content=content,
+        step=state.reasoning_step,
+    ))
+
+
+def build_progress_event(
+    step: int,
+    total_steps: int,
+    label: str,
+    status: str = "in_progress",
+) -> str:
+    """Build a progress SSE event."""
+    return format_sse_event(ProgressEvent(
+        step=step,
+        total_steps=total_steps,
+        label=label,
+        status=status,
+    ))
+
+
+def build_tool_start_event(
+    tool_name: str,
+    tool_id: str,
+    arguments: dict | None = None,
+) -> str:
+    """Build a tool call started SSE event."""
+    return format_sse_event(ToolCallEvent(
+        tool_name=tool_name,
+        tool_id=tool_id,
+        status="started",
+        arguments=arguments,
+    ))
+
+
+def build_tool_complete_event(
+    tool_name: str,
+    tool_id: str,
+    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] + "..."
+
+    return format_sse_event(ToolCallEvent(
+        tool_name=tool_name,
+        tool_id=tool_id,
+        status="completed",
+        arguments=arguments,
+        result=result_str,
+    ))
+
+
+def build_metadata_event(
+    message_id: str | None = None,
+    in_reply_to: str | None = None,
+    session_id: str | None = None,
+    agent_schema: str | None = None,
+    responding_agent: str | None = None,
+    confidence: float | None = None,
+    sources: list | None = None,
+    model_version: str | None = None,
+    latency_ms: int | None = None,
+    token_count: int | None = None,
+    trace_id: str | None = None,
+    span_id: str | None = None,
+    extra: dict | None = None,
+) -> str:
+    """Build a metadata SSE event."""
+    return format_sse_event(MetadataEvent(
+        message_id=message_id,
+        in_reply_to=in_reply_to,
+        session_id=session_id,
+        agent_schema=agent_schema,
+        responding_agent=responding_agent,
+        confidence=confidence,
+        sources=sources,
+        model_version=model_version,
+        latency_ms=latency_ms,
+        token_count=token_count,
+        trace_id=trace_id,
+        span_id=span_id,
+        extra=extra,
+    ))
+
+
+# =============================================================================
+# TOOL ARGUMENT EXTRACTION
+# =============================================================================
+
+def extract_tool_args(part) -> dict | None:
+    """
+    Extract arguments from a ToolCallPart.
+
+    Handles various formats:
+    - ArgsDict object with args_dict attribute
+    - Plain dict
+    - JSON string
+    """
+    if part.args is None:
+        return None
+
+    if hasattr(part.args, 'args_dict'):
+        return part.args.args_dict
+
+    if isinstance(part.args, dict):
+        return part.args
+
+    if isinstance(part.args, str) and part.args:
+        try:
+            return json.loads(part.args)
+        except json.JSONDecodeError:
+            logger.warning(f"Failed to parse tool args: {part.args[:100]}")
+
+    return None
+
+
+def log_tool_call(tool_name: str, args_dict: dict | None) -> None:
+    """Log a tool call with key parameters."""
+    if args_dict and tool_name == "search_rem":
+        query_type = args_dict.get("query_type", "?")
+        limit = args_dict.get("limit", 20)
+        table = args_dict.get("table", "")
+        query_text = args_dict.get("query_text", args_dict.get("entity_key", ""))
+        if query_text and len(str(query_text)) > 50:
+            query_text = str(query_text)[:50] + "..."
+        logger.info(f"🔧 {tool_name} {query_type.upper()} '{query_text}' table={table} limit={limit}")
+    else:
+        logger.info(f"🔧 {tool_name}")
+
+
+def log_tool_result(tool_name: str, result_content: Any) -> None:
+    """Log a tool result with key metrics."""
+    if tool_name == "search_rem" and isinstance(result_content, dict):
+        results = result_content.get("results", {})
+        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 = ""
+
+        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")
+
+
+# =============================================================================
+# METADATA EXTRACTION
+# =============================================================================
+
+def extract_metadata_from_result(result_content: Any) -> dict | None:
+    """
+    Extract metadata from a register_metadata tool result.
+
+    Returns dict with extracted fields or None if not a metadata event.
+    """
+    if not isinstance(result_content, dict):
+        return None
+
+    if not result_content.get("_metadata_event"):
+        return None
+
+    return {
+        "confidence": result_content.get("confidence"),
+        "sources": result_content.get("sources"),
+        "references": result_content.get("references"),
+        "flags": result_content.get("flags"),
+        "session_name": result_content.get("session_name"),
+        "risk_level": result_content.get("risk_level"),
+        "risk_score": result_content.get("risk_score"),
+        "risk_reasoning": result_content.get("risk_reasoning"),
+        "recommended_action": result_content.get("recommended_action"),
+        "agent_schema": result_content.get("agent_schema"),
+        "extra": result_content.get("extra"),
+    }

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/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 = {