pydantic-ai-rlm 0.1.1__tar.gz → 0.1.2__tar.gz

{pydantic_ai_rlm-0.1.1 → pydantic_ai_rlm-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-rlm
-Version: 0.1.1
+Version: 0.1.2
 Summary: Recursive Language Model (RLM) toolset for Pydantic AI - handle extremely large contexts
 Author: Pydantic AI RLM Contributors
 License-Expression: MIT
@@ -52,6 +52,8 @@ Description-Content-Type: text/markdown
    • 
   <b>Sub-Model Delegation</b>
   &nbsp;•&nbsp;
+  <b>Grounded Citations</b>
+  &nbsp;•&nbsp;
   <b>Fully Type-Safe</b>
 </p>
 
@@ -203,6 +205,30 @@ result = await agent.run(
 )
 ```
 
+### Grounded Responses with Citations
+
+Get answers with traceable citations back to the source:
+
+```python
+from pydantic_ai_rlm import run_rlm_analysis
+
+# Enable grounding for citation tracking
+result = await run_rlm_analysis(
+    context=financial_report,
+    query="What were the key revenue changes?",
+    model="openai:gpt-5",
+    grounded=True,  # Returns GroundedResponse instead of str
+)
+
+# Response contains citation markers
+print(result.info)
+# "Revenue increased [1] primarily due to [2]"
+
+# Grounding maps markers to exact quotes from the source
+print(result.grounding)
+# {"1": "by 45% year-over-year", "2": "expansion into Asian markets"}
+```
+
 ---
 
 ## API Reference
@@ -217,6 +243,7 @@ agent = create_rlm_agent(
     sub_model="openai:gpt-5-mini",  # Model for llm_query() (optional)
     code_timeout=60.0,               # Timeout for code execution
     custom_instructions="...",       # Additional instructions
+    grounded=True,                   # Return GroundedResponse with citations
 )
 ```
 
@@ -241,6 +268,11 @@ answer = await run_rlm_analysis(context, query, model="openai:gpt-5")
 
 # Sync
 answer = run_rlm_analysis_sync(context, query, model="openai:gpt-5")
+
+# With grounding (returns GroundedResponse)
+result = await run_rlm_analysis(context, query, grounded=True)
+print(result.info)       # Text with [N] markers
+print(result.grounding)  # {"1": "exact quote", ...}
 ```
 
 ### `RLMDependencies`

{pydantic_ai_rlm-0.1.1 → pydantic_ai_rlm-0.1.2}/README.md

@@ -24,6 +24,8 @@
    • 
   <b>Sub-Model Delegation</b>
   &nbsp;•&nbsp;
+  <b>Grounded Citations</b>
+  &nbsp;•&nbsp;
   <b>Fully Type-Safe</b>
 </p>
 
@@ -175,6 +177,30 @@ result = await agent.run(
 )
 ```
 
+### Grounded Responses with Citations
+
+Get answers with traceable citations back to the source:
+
+```python
+from pydantic_ai_rlm import run_rlm_analysis
+
+# Enable grounding for citation tracking
+result = await run_rlm_analysis(
+    context=financial_report,
+    query="What were the key revenue changes?",
+    model="openai:gpt-5",
+    grounded=True,  # Returns GroundedResponse instead of str
+)
+
+# Response contains citation markers
+print(result.info)
+# "Revenue increased [1] primarily due to [2]"
+
+# Grounding maps markers to exact quotes from the source
+print(result.grounding)
+# {"1": "by 45% year-over-year", "2": "expansion into Asian markets"}
+```
+
 ---
 
 ## API Reference
@@ -189,6 +215,7 @@ agent = create_rlm_agent(
     sub_model="openai:gpt-5-mini",  # Model for llm_query() (optional)
     code_timeout=60.0,               # Timeout for code execution
     custom_instructions="...",       # Additional instructions
+    grounded=True,                   # Return GroundedResponse with citations
 )
 ```
 
@@ -213,6 +240,11 @@ answer = await run_rlm_analysis(context, query, model="openai:gpt-5")
 
 # Sync
 answer = run_rlm_analysis_sync(context, query, model="openai:gpt-5")
+
+# With grounding (returns GroundedResponse)
+result = await run_rlm_analysis(context, query, grounded=True)
+print(result.info)       # Text with [N] markers
+print(result.grounding)  # {"1": "exact quote", ...}
 ```
 
 ### `RLMDependencies`

{pydantic_ai_rlm-0.1.1 → pydantic_ai_rlm-0.1.2}/examples/needle_in_haystack.py

@@ -65,6 +65,7 @@ def main():
         query=query,
         model="openai:gpt-5",
         sub_model="openai:gpt-5-mini",
+        grounded=True,
     )
 
     print(f"\nResult: {result}")

{pydantic_ai_rlm-0.1.1 → pydantic_ai_rlm-0.1.2}/examples/semantic_search.py

@@ -142,6 +142,7 @@ def main():
         query=query,
         model="openai:gpt-5",
         sub_model="openai:gpt-5-mini",
+        grounded=True,
     )
 
     print(f"\nResult: {result}")

{pydantic_ai_rlm-0.1.1 → pydantic_ai_rlm-0.1.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pydantic-ai-rlm"
-version = "0.1.1"
+version = "0.1.2"
 description = "Recursive Language Model (RLM) toolset for Pydantic AI - handle extremely large contexts"
 readme = "README.md"
 license = "MIT"

{pydantic_ai_rlm-0.1.1 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/__init__.py

@@ -1,7 +1,9 @@
 from .agent import create_rlm_agent, run_rlm_analysis, run_rlm_analysis_sync
 from .dependencies import ContextType, RLMConfig, RLMDependencies
 from .logging import configure_logging
+from .models import GroundedResponse
 from .prompts import (
+    GROUNDING_INSTRUCTIONS,
     LLM_QUERY_INSTRUCTIONS,
     RLM_INSTRUCTIONS,
     build_rlm_instructions,
@@ -13,9 +15,11 @@ from .toolset import (
 )
 
 __all__ = [
+    "GROUNDING_INSTRUCTIONS",
     "LLM_QUERY_INSTRUCTIONS",
     "RLM_INSTRUCTIONS",
     "ContextType",
+    "GroundedResponse",
     "REPLEnvironment",
     "REPLResult",
     "RLMConfig",

{pydantic_ai_rlm-0.1.1 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/agent.py

@@ -1,21 +1,45 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Literal, overload
 
 from pydantic_ai import Agent, UsageLimits
 
 from .dependencies import ContextType, RLMConfig, RLMDependencies
+from .models import GroundedResponse
 from .prompts import build_rlm_instructions
 from .toolset import create_rlm_toolset
 
 
+@overload
 def create_rlm_agent(
     model: str = "openai:gpt-5",
     sub_model: str | None = None,
     code_timeout: float = 60.0,
-    include_example_instructions: bool = True,
     custom_instructions: str | None = None,
-) -> Agent[RLMDependencies, str]:
+    *,
+    grounded: Literal[False] = False,
+) -> Agent[RLMDependencies, str]: ...
+
+
+@overload
+def create_rlm_agent(
+    model: str = "openai:gpt-5",
+    sub_model: str | None = None,
+    code_timeout: float = 60.0,
+    custom_instructions: str | None = None,
+    *,
+    grounded: Literal[True],
+) -> Agent[RLMDependencies, GroundedResponse]: ...
+
+
+def create_rlm_agent(
+    model: str = "openai:gpt-5",
+    sub_model: str | None = None,
+    code_timeout: float = 60.0,
+    custom_instructions: str | None = None,
+    *,
+    grounded: bool = False,
+) -> Agent[RLMDependencies, str] | Agent[RLMDependencies, GroundedResponse]:
     """
     Create a Pydantic AI agent with REPL code execution capabilities.
 
@@ -26,11 +50,12 @@ def create_rlm_agent(
             available in the REPL, allowing the agent to delegate sub-queries.
             Example: "openai:gpt-5-mini" or "anthropic:claude-3-haiku-20240307"
         code_timeout: Timeout for code execution in seconds
-        include_example_instructions: Include detailed examples in instructions
         custom_instructions: Additional instructions to append
+        grounded: If True, return a GroundedResponse with citation markers
 
     Returns:
-        Configured Agent instance
+        Configured Agent instance. Returns Agent[RLMDependencies, GroundedResponse]
+        when grounded=True, otherwise Agent[RLMDependencies, str].
 
     Example:
         ```python
@@ -48,19 +73,28 @@ def create_rlm_agent(
         )
         result = await agent.run("What are the main themes?", deps=deps)
         print(result.output)
+
+        # Create grounded agent
+        grounded_agent = create_rlm_agent(model="openai:gpt-5", grounded=True)
+        result = await grounded_agent.run("What happened?", deps=deps)
+        print(result.output.info)  # Response with [N] markers
+        print(result.output.grounding)  # {"1": "exact quote", ...}
         ```
     """
     toolset = create_rlm_toolset(code_timeout=code_timeout, sub_model=sub_model)
 
     instructions = build_rlm_instructions(
         include_llm_query=sub_model is not None,
+        include_grounding=grounded,
         custom_suffix=custom_instructions,
     )
 
-    agent: Agent[RLMDependencies, str] = Agent(
+    output_type: type[str] | type[GroundedResponse] = GroundedResponse if grounded else str
+
+    agent: Agent[RLMDependencies, Any] = Agent(
         model,
         deps_type=RLMDependencies,
-        output_type=str,
+        output_type=output_type,
         toolsets=[toolset],
         instructions=instructions,
     )
@@ -68,6 +102,34 @@ def create_rlm_agent(
     return agent
 
 
+@overload
+async def run_rlm_analysis(
+    context: ContextType,
+    query: str,
+    model: str = "openai:gpt-5",
+    sub_model: str | None = None,
+    config: RLMConfig | None = None,
+    max_tool_calls: int = 50,
+    *,
+    grounded: Literal[False] = False,
+    **agent_kwargs: Any,
+) -> str: ...
+
+
+@overload
+async def run_rlm_analysis(
+    context: ContextType,
+    query: str,
+    model: str = "openai:gpt-5",
+    sub_model: str | None = None,
+    config: RLMConfig | None = None,
+    max_tool_calls: int = 50,
+    *,
+    grounded: Literal[True],
+    **agent_kwargs: Any,
+) -> GroundedResponse: ...
+
+
 async def run_rlm_analysis(
     context: ContextType,
     query: str,
@@ -75,8 +137,10 @@ async def run_rlm_analysis(
     sub_model: str | None = None,
     config: RLMConfig | None = None,
     max_tool_calls: int = 50,
+    *,
+    grounded: bool = False,
     **agent_kwargs: Any,
-) -> str:
+) -> str | GroundedResponse:
     """
     Convenience function to run RLM analysis on a context.
 
@@ -89,25 +153,36 @@ async def run_rlm_analysis(
             available in the REPL, allowing the agent to delegate sub-queries.
         config: Optional RLMConfig for customization
         max_tool_calls: Maximum tool calls allowed
+        grounded: If True, return a GroundedResponse with citation markers
         **agent_kwargs: Additional arguments passed to create_rlm_agent()
 
     Returns:
-        The agent's final answer as a string
+        The agent's final answer. Returns GroundedResponse when grounded=True,
+        otherwise returns str.
 
     Example:
         ```python
         from pydantic_ai_rlm import run_rlm_analysis
 
-        # With sub-model for llm_query
+        # Standard string response
         answer = await run_rlm_analysis(
             context=huge_document,
             query="Find the magic number hidden in the text",
             sub_model="openai:gpt-5-mini",
         )
         print(answer)
+
+        # Grounded response with citations
+        result = await run_rlm_analysis(
+            context=document,
+            query="What was the revenue change?",
+            grounded=True,
+        )
+        print(result.info)  # "Revenue grew [1]..."
+        print(result.grounding)  # {"1": "increased by 45%", ...}
         ```
     """
-    agent = create_rlm_agent(model=model, sub_model=sub_model, **agent_kwargs)
+    agent = create_rlm_agent(model=model, sub_model=sub_model, grounded=grounded, **agent_kwargs)
 
     effective_config = config or RLMConfig()
     if sub_model and not effective_config.sub_model:
@@ -127,6 +202,7 @@ async def run_rlm_analysis(
     return result.output
 
 
+@overload
 def run_rlm_analysis_sync(
     context: ContextType,
     query: str,
@@ -134,14 +210,63 @@ def run_rlm_analysis_sync(
     sub_model: str | None = None,
     config: RLMConfig | None = None,
     max_tool_calls: int = 50,
+    *,
+    grounded: Literal[False] = False,
     **agent_kwargs: Any,
-) -> str:
+) -> str: ...
+
+
+@overload
+def run_rlm_analysis_sync(
+    context: ContextType,
+    query: str,
+    model: str = "openai:gpt-5",
+    sub_model: str | None = None,
+    config: RLMConfig | None = None,
+    max_tool_calls: int = 50,
+    *,
+    grounded: Literal[True],
+    **agent_kwargs: Any,
+) -> GroundedResponse: ...
+
+
+def run_rlm_analysis_sync(
+    context: ContextType,
+    query: str,
+    model: str = "openai:gpt-5",
+    sub_model: str | None = None,
+    config: RLMConfig | None = None,
+    max_tool_calls: int = 50,
+    *,
+    grounded: bool = False,
+    **agent_kwargs: Any,
+) -> str | GroundedResponse:
     """
     Synchronous version of run_rlm_analysis.
 
     See run_rlm_analysis() for full documentation.
+
+    Example:
+        ```python
+        from pydantic_ai_rlm import run_rlm_analysis_sync
+
+        # Standard string response
+        answer = run_rlm_analysis_sync(
+            context=document,
+            query="What happened?",
+        )
+
+        # Grounded response with citations
+        result = run_rlm_analysis_sync(
+            context=document,
+            query="What was the revenue change?",
+            grounded=True,
+        )
+        print(result.info)  # "Revenue grew [1]..."
+        print(result.grounding)  # {"1": "increased by 45%", ...}
+        ```
     """
-    agent = create_rlm_agent(model=model, sub_model=sub_model, **agent_kwargs)
+    agent = create_rlm_agent(model=model, sub_model=sub_model, grounded=grounded, **agent_kwargs)
 
     effective_config = config or RLMConfig()
     if sub_model and not effective_config.sub_model:

pydantic_ai_rlm-0.1.2/src/pydantic_ai_rlm/models.py

@@ -0,0 +1,23 @@
+"""Pydantic models for structured RLM outputs."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class GroundedResponse(BaseModel):
+    """A response with citation markers mapping to exact quotes from source documents.
+
+    Example:
+        ```python
+        response = GroundedResponse(
+            info="Revenue grew [1] driven by expansion [2]", grounding={"1": "increased by 45%", "2": "new markets in Asia"}
+        )
+        ```
+    """
+
+    info: str = Field(description="Response text with citation markers like [1]")
+    grounding: dict[str, str] = Field(
+        default_factory=dict,
+        description="Mapping from citation markers to exact quotes from the source",
+    )

{pydantic_ai_rlm-0.1.1 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/prompts.py

@@ -50,6 +50,61 @@ print(f"Final answer: {results}")
 4. **Be thorough** - For needle-in-haystack, search the entire context
 """
 
+GROUNDING_INSTRUCTIONS = """
+
+## Grounding Requirements
+
+Your response MUST include grounded citations. This means:
+
+1. **Citation Format**: Use markers like `[1]`, `[2]`, etc. in your response text
+2. **Exact Quotes**: Each marker must map to an EXACT quote from the source context (verbatim, no paraphrasing)
+3. **Quote Length**: Each quote should be 10-200 characters - enough to be meaningful but not too long
+4. **Consecutive Numbering**: Number citations consecutively starting from 1
+
+### Output Format
+
+Your final answer must be valid JSON with this structure:
+```json
+{
+   "info": "The document states that X Y Z [1]. Additionally, A B C [2]",
+   "grounding": {
+      "1": "exact quote from source",
+      "2": "another exact quote from source"
+   }
+}
+```
+
+### Example
+
+If the context contains: "The company's revenue increased by 45% in Q3 2024, driven by expansion into new markets in Asia."
+
+Your response should look like:
+```json
+{
+   "info": "Revenue showed strong growth [1] with geographic expansion being a key driver [2].",
+   "grounding": {
+      "1": "revenue increased by 45% in Q3 2024",
+      "2": "driven by expansion into new markets in Asia"
+   }
+}
+```
+
+### Finding Quotes in Code
+
+Use this approach to find and verify exact quotes:
+```python
+# Find a specific phrase in context
+search_term = "revenue"
+idx = context.lower().find(search_term)
+if idx != -1:
+    # Extract surrounding context for the quote
+    quote = context[max(0, idx):idx+100]
+    print(f"Found: {quote}")
+```
+
+**Important**: Every citation marker in your `info` field MUST have a corresponding entry in `grounding`. Only output the JSON object, no additional text.
+"""
+
 LLM_QUERY_INSTRUCTIONS = """
 
 ## Sub-LLM Queries
@@ -93,14 +148,15 @@ print(result)
 
 def build_rlm_instructions(
     include_llm_query: bool = False,
+    include_grounding: bool = False,
     custom_suffix: str | None = None,
 ) -> str:
     """
     Build RLM instructions with optional customization.
 
     Args:
-        include_examples: Whether to include detailed examples
         include_llm_query: Whether to include llm_query() documentation
+        include_grounding: Whether to include grounding/citation instructions
         custom_suffix: Additional instructions to append
 
     Returns:
@@ -109,8 +165,10 @@ def build_rlm_instructions(
     base = RLM_INSTRUCTIONS
 
     if include_llm_query:
-        llm_docs = LLM_QUERY_INSTRUCTIONS
-        base = f"{base}{llm_docs}"
+        base = f"{base}{LLM_QUERY_INSTRUCTIONS}"
+
+    if include_grounding:
+        base = f"{base}{GROUNDING_INSTRUCTIONS}"
 
     if custom_suffix:
         base = f"{base}\n\n## Additional Instructions\n\n{custom_suffix}"

{pydantic_ai_rlm-0.1.1 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/repl.py

@@ -7,6 +7,7 @@ import os
 import shutil
 import sys
 import tempfile
+import textwrap
 import threading
 import time
 from contextlib import contextmanager
@@ -328,6 +329,9 @@ with open(r'{context_path}', 'r', encoding='utf-8') as f:
         Returns:
             REPLResult with stdout, stderr, locals, and timing
         """
+        # Normalize code: remove common leading whitespace and strip
+        code = textwrap.dedent(code).strip()
+
         start_time = time.time()
         success = True
         stdout_content = ""