pydantic-ai-rlm 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

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/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/logging.py

@@ -47,7 +47,7 @@ class RLMLogger:
             )
             self.console.print(panel)
         else:
-            print(f"\n{'='*50}")
+            print(f"\n{'=' * 50}")
             print("CODE EXECUTION")
             print("=" * 50)
             print(code)
@@ -141,7 +141,7 @@ class RLMLogger:
     def _log_result_plain(self, result: REPLResult) -> None:
         """Log result using plain text."""
         status = "SUCCESS" if result.success else "ERROR"
-        print(f"\n{'='*50}")
+        print(f"\n{'=' * 50}")
         print(f"RESULT: {status} (executed in {result.execution_time:.3f}s)")
         print("=" * 50)
 
@@ -159,11 +159,7 @@ class RLMLogger:
                 stderr = stderr[:1000] + "\n... (truncated)"
             print(stderr)
 
-        user_vars = {
-            k: v
-            for k, v in result.locals.items()
-            if not k.startswith("_") and k not in ("context", "json", "re", "os")
-        }
+        user_vars = {k: v for k, v in result.locals.items() if not k.startswith("_") and k not in ("context", "json", "re", "os")}
         if user_vars:
             print("\nVariables:")
             for name, value in list(user_vars.items())[:10]:
@@ -198,7 +194,7 @@ class RLMLogger:
             )
             self.console.print(panel)
         else:
-            print(f"\n{'='*50}")
+            print(f"\n{'=' * 50}")
             print("LLM QUERY")
             print("=" * 50)
             display_prompt = prompt
@@ -226,7 +222,7 @@ class RLMLogger:
             )
             self.console.print(panel)
         else:
-            print(f"\n{'='*50}")
+            print(f"\n{'=' * 50}")
             print("LLM RESPONSE")
             print("=" * 50)
             display_response = response

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

{pydantic_ai_rlm-0.1.0.dist-info → pydantic_ai_rlm-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-rlm
-Version: 0.1.0
+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
@@ -16,7 +16,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Typing :: Typed
 Requires-Python: >=3.10
-Requires-Dist: pydantic-ai>=0.1.0
+Requires-Dist: pydantic-ai-slim[cli]>=0.1.0
 Provides-Extra: dev
 Requires-Dist: mypy>=1.0; extra == 'dev'
 Requires-Dist: pytest-asyncio>=0.21; extra == 'dev'
@@ -34,6 +34,7 @@ Description-Content-Type: text/markdown
 
 <p align="center">
   <a href="https://github.com/vstorm-co/pydantic-ai-rlm">GitHub</a> •
+  <a href="https://pypi.org/project/pydantic-ai-rlm/">PyPI</a> •
   <a href="https://github.com/vstorm-co/pydantic-ai-rlm#examples">Examples</a>
 </p>
 
@@ -41,6 +42,7 @@ Description-Content-Type: text/markdown
   <a href="https://www.python.org/downloads/"><img src="https://img.shields.io/badge/python-3.10+-blue.svg" alt="Python 3.10+"></a>
   <a href="https://opensource.org/licenses/MIT"><img src="https://img.shields.io/badge/License-MIT-yellow.svg" alt="License: MIT"></a>
   <a href="https://github.com/pydantic/pydantic-ai"><img src="https://img.shields.io/badge/Powered%20by-Pydantic%20AI-E92063?logo=pydantic&logoColor=white" alt="Pydantic AI"></a>
+  <a href="https://pypi.org/project/pydantic-ai-rlm/"><img src="https://img.shields.io/pypi/v/pydantic-ai-rlm.svg" alt="PyPI version"></a>
 </p>
 
 <p align="center">
@@ -50,6 +52,8 @@ Description-Content-Type: text/markdown
   &nbsp;•&nbsp;
   <b>Sub-Model Delegation</b>
   &nbsp;•&nbsp;
+  <b>Grounded Citations</b>
+  &nbsp;•&nbsp;
   <b>Fully Type-Safe</b>
 </p>
 
@@ -201,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
@@ -215,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
 )
 ```
 
@@ -239,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`
@@ -278,16 +312,19 @@ configure_logging(enabled=False)
 ```
 
 Install with rich logging support for syntax highlighting and styled output:
+
 ```bash
 pip install pydantic-ai-rlm[logging]
 ```
 
 Or install rich separately:
+
 ```bash
 pip install rich
 ```
 
 When enabled, you'll see:
+
 - Syntax-highlighted code being executed (with rich)
 - Execution results with status indicators (SUCCESS/ERROR)
 - Execution time for each code block

pydantic_ai_rlm-0.1.2.dist-info/RECORD

@@ -0,0 +1,14 @@
+pydantic_ai_rlm/__init__.py,sha256=iU_ToTPYqbgW3at4DI5fUVqcvsXvXtt2m61x-z3kVYM,913
+pydantic_ai_rlm/agent.py,sha256=9wZdl1POW6B9LOf4IquFCrHt34gBPGEw2Z9631GMmXQ,8379
+pydantic_ai_rlm/dependencies.py,sha256=bzPXhCRJ-0D5N0FS32LAETXxp6cbInU00qaMr0qd_iQ,1306
+pydantic_ai_rlm/logging.py,sha256=6AzSXBve5pRCfL1pGT5KkWwOZLvZGBYxS-USwieE6lY,9270
+pydantic_ai_rlm/models.py,sha256=9VMVybNBXU3kGZPzg6LPmxTN9qBhVFl7qY5RVd9Ch7o,716
+pydantic_ai_rlm/prompts.py,sha256=vof8tYJWG4ChLR5cEJeB8x0oeODdqyW7yoHPgS0veIY,5415
+pydantic_ai_rlm/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pydantic_ai_rlm/repl.py,sha256=9Bja7NJg5qObKwxvBOVKeGEOwDWgTNS2X2259r5tHgA,14995
+pydantic_ai_rlm/toolset.py,sha256=g1244qq2VzaQVhAQixLGPCWW7vJqQfjhUgf0kDxkkzo,5564
+pydantic_ai_rlm/utils.py,sha256=_tqdMcD-_UbRbdK6ECD-XlXfRHgeAKj7ulLaXrYOJx4,1500
+pydantic_ai_rlm-0.1.2.dist-info/METADATA,sha256=ntFi4yQHtRnc75ySjxg8jfpNw1CRNDw7dKsWiDyTirQ,12822
+pydantic_ai_rlm-0.1.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+pydantic_ai_rlm-0.1.2.dist-info/licenses/LICENSE,sha256=ktSo7DxQBRIZIhD4P53iVxq88CG3cvXRTBRqPT8zO1c,1085
+pydantic_ai_rlm-0.1.2.dist-info/RECORD,,

pydantic_ai_rlm-0.1.0.dist-info/RECORD

@@ -1,13 +0,0 @@
-pydantic_ai_rlm/__init__.py,sha256=94ev1VspkrVdB-PabCD6I791vZkBVASXv_qqneYQQ0I,794
-pydantic_ai_rlm/agent.py,sha256=_zHBLatfhquGw3vjHC06b0lBZRpLv2nrxi7w69f7lBY,4824
-pydantic_ai_rlm/dependencies.py,sha256=bzPXhCRJ-0D5N0FS32LAETXxp6cbInU00qaMr0qd_iQ,1306
-pydantic_ai_rlm/logging.py,sha256=DXJzfrfOnyR3VAbKF4_0Mc2MP7zqPmxyFDph4Q9kAb0,9308
-pydantic_ai_rlm/prompts.py,sha256=BpbugghZ_I_IjC_Il5pUClMG7bgWfzWgtsUsxwyjH10,3648
-pydantic_ai_rlm/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pydantic_ai_rlm/repl.py,sha256=twn6ArhaONYg49Feoi-sD6z0HkYHyHHc-prdBCsj0jo,14864
-pydantic_ai_rlm/toolset.py,sha256=g1244qq2VzaQVhAQixLGPCWW7vJqQfjhUgf0kDxkkzo,5564
-pydantic_ai_rlm/utils.py,sha256=_tqdMcD-_UbRbdK6ECD-XlXfRHgeAKj7ulLaXrYOJx4,1500
-pydantic_ai_rlm-0.1.0.dist-info/METADATA,sha256=2w1yJoJGl5HHCWcl6JRz5pDys0o4iBTLcWSsXuGt8nk,11613
-pydantic_ai_rlm-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-pydantic_ai_rlm-0.1.0.dist-info/licenses/LICENSE,sha256=ktSo7DxQBRIZIhD4P53iVxq88CG3cvXRTBRqPT8zO1c,1085
-pydantic_ai_rlm-0.1.0.dist-info/RECORD,,