remdb 0.3.242__py3-none-any.whl

rem/agentic/query.py

@@ -0,0 +1,117 @@
+"""
+Agent query model for structured agent input.
+
+Design pattern for standardized agent query structure with:
+- Primary query (user question/task)
+- Knowledge context (retrieved context, documentation)
+- Scratchpad (working memory, session state)
+
+Key Design Pattern 
+- Separates query from retrieval (query is what user asks, knowledge is what we retrieve)
+- Scratchpad enables multi-turn reasoning and state tracking
+- Supports markdown + fenced JSON for structured data
+- Converts to single prompt string for agent consumption
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class AgentQuery(BaseModel):
+    """
+    Standard query structure for agent execution.
+
+    Provides consistent structure for queries, knowledge context, and
+    working memory across all agent types.
+
+    Design Pattern 
+    - query: User's question/task (markdown + fenced JSON)
+    - knowledge: Retrieved context from REM queries (markdown + fenced JSON)
+    - scratchpad: Working memory for multi-turn reasoning (JSON or markdown)
+
+    Example:
+        query = AgentQuery(
+            query="Find all documents Sarah authored",
+            knowledge=\"\"\"
+            # Entity Information
+            Sarah Chen (person/employee)
+            - Role: Senior Engineer
+            - Projects: [Project Alpha, TiDB Migration]
+            \"\"\",
+            scratchpad={"current_case": "TAP-1234", "stage": "entity_lookup"}
+        )
+
+        prompt = query.to_prompt()
+        result = await agent.run(prompt)
+    """
+
+    query: str = Field(
+        ...,
+        description="Primary user query or task (markdown format, may include fenced JSON)",
+        examples=[
+            "Find all documents Sarah authored",
+            "What happened in Q4 retrospective?",
+            "TRAVERSE manages WITH LOOKUP sarah-chen DEPTH 2",
+        ],
+    )
+
+    knowledge: str = Field(
+        default="",
+        description="Background knowledge and context (markdown, may include fenced JSON/code)",
+        examples=[
+            "# Entity: sarah-chen\nType: person/employee\nRole: Senior Engineer",
+            "Retrieved resources:\n```json\n[{...}]\n```",
+        ],
+    )
+
+    scratchpad: str | dict[str, Any] = Field(
+        default="",
+        description="Working memory for session state (JSON object or markdown with fenced JSON)",
+        examples=[
+            {"stage": "lookup", "visited_entities": ["sarah-chen"]},
+            "# Session State\n\nStage: TRAVERSE depth 1\n\n```json\n{\"nodes\": [...]}\n```",
+        ],
+    )
+
+    def to_prompt(self) -> str:
+        """
+        Convert query components to single prompt string.
+
+        Combines query, knowledge, and scratchpad into formatted prompt
+        for agent consumption.
+
+        Returns:
+            Formatted prompt string with sections
+
+        Example:
+            # Query
+
+            Find all documents Sarah authored
+
+            # Knowledge
+
+            Entity: sarah-chen
+            Type: person/employee
+
+            # Scratchpad
+
+            ```json
+            {"stage": "lookup"}
+            ```
+        """
+        parts = [f"# Query\n\n{self.query}"]
+
+        if self.knowledge:
+            parts.append(f"\n# Knowledge\n\n{self.knowledge}")
+
+        if self.scratchpad:
+            if isinstance(self.scratchpad, dict):
+                import json
+
+                scratchpad_str = json.dumps(self.scratchpad, indent=2)
+                parts.append(f"\n# Scratchpad\n\n```json\n{scratchpad_str}\n```")
+            else:
+                parts.append(f"\n# Scratchpad\n\n{self.scratchpad}")
+
+        return "\n".join(parts)

rem/agentic/query_helper.py

@@ -0,0 +1,89 @@
+"""
+Helper functions for REM Query Agent.
+
+This module provides convenience functions that load the REM Query Agent
+from YAML schema and execute queries.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel
+
+from .context import AgentContext
+from .providers.pydantic_ai import create_agent_from_schema_file
+
+
+class REMQueryOutput(BaseModel):
+    """
+    REM Query Agent structured output.
+
+    Matches the schema defined in schemas/agents/core/rem-query-agent.yaml
+    """
+
+    query: str
+    confidence: float
+    reasoning: str = ""
+
+
+async def ask_rem(
+    natural_query: str,
+    user_id: str = "system",
+    llm_model: str | None = None,
+) -> REMQueryOutput:
+    """
+    Convert natural language query to structured REM query.
+
+    Loads the REM Query Agent from YAML schema and executes the query.
+
+    Args:
+        natural_query: User's question in natural language
+        user_id: User ID for context (defaults to "system")
+        llm_model: Optional LLM model override
+
+    Returns:
+        REMQueryOutput with query, confidence, and reasoning
+
+    Example:
+        result = await ask_rem("Show me Sarah Chen")
+        # REMQueryOutput(
+        #     query="LOOKUP sarah-chen",
+        #     confidence=1.0,
+        #     reasoning=""
+        # )
+    """
+    # Create context (only pass default_model if llm_model is provided)
+    context_kwargs = {"user_id": user_id}
+    if llm_model is not None:
+        context_kwargs["default_model"] = llm_model
+    context = AgentContext(**context_kwargs)
+
+    # Load agent from YAML schema
+    agent = await create_agent_from_schema_file(
+        schema_name_or_path="rem-query-agent",
+        context=context,
+        model_override=llm_model,  # type: ignore[arg-type]
+    )
+
+    # Run query
+    result = await agent.run(natural_query)
+
+    # Handle different Pydantic AI versions
+    if hasattr(result, "data"):
+        output = result.data
+    elif hasattr(result, "output"):
+        output = result.output
+    else:
+        output = result
+
+    # Convert to REMQueryOutput if not already
+    if isinstance(output, dict):
+        return REMQueryOutput(**output)
+    elif isinstance(output, REMQueryOutput):
+        return output
+    else:
+        # Fallback: try to extract fields
+        return REMQueryOutput(
+            query=getattr(output, "query", str(output)),
+            confidence=getattr(output, "confidence", 0.5),
+            reasoning=getattr(output, "reasoning", ""),
+        )