pydantic-ai-rlm 0.1.0__tar.gz → 0.1.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- pydantic_ai_rlm-0.1.2/.github/workflows/publish.yml +44 -0
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/PKG-INFO +39 -2
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/README.md +37 -0
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/examples/needle_in_haystack.py +3 -4
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/examples/semantic_search.py +4 -8
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/pyproject.toml +2 -2
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/__init__.py +4 -0
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/agent.py +138 -13
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/logging.py +5 -9
- pydantic_ai_rlm-0.1.2/src/pydantic_ai_rlm/models.py +23 -0
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/prompts.py +61 -3
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/repl.py +4 -0
- pydantic_ai_rlm-0.1.0/.logfire/.gitignore +0 -1
- pydantic_ai_rlm-0.1.0/.logfire/logfire_credentials.json +0 -6
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/.gitignore +0 -0
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/.pre-commit-config.yaml +0 -0
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/LICENSE +0 -0
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/dependencies.py +0 -0
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/py.typed +0 -0
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/toolset.py +0 -0
- {pydantic_ai_rlm-0.1.0 → pydantic_ai_rlm-0.1.2}/src/pydantic_ai_rlm/utils.py +0 -0
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
name: Publish to PyPI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
release:
|
|
5
|
+
types: [published]
|
|
6
|
+
|
|
7
|
+
jobs:
|
|
8
|
+
build:
|
|
9
|
+
name: Build distribution
|
|
10
|
+
runs-on: ubuntu-latest
|
|
11
|
+
steps:
|
|
12
|
+
- uses: actions/checkout@v4
|
|
13
|
+
|
|
14
|
+
- name: Install uv
|
|
15
|
+
uses: astral-sh/setup-uv@v5
|
|
16
|
+
|
|
17
|
+
- name: Build package
|
|
18
|
+
run: uv build
|
|
19
|
+
|
|
20
|
+
- name: Upload distribution artifacts
|
|
21
|
+
uses: actions/upload-artifact@v4
|
|
22
|
+
with:
|
|
23
|
+
name: python-package-distributions
|
|
24
|
+
path: dist/
|
|
25
|
+
|
|
26
|
+
publish-pypi:
|
|
27
|
+
name: Publish to PyPI
|
|
28
|
+
needs: build
|
|
29
|
+
runs-on: ubuntu-latest
|
|
30
|
+
environment:
|
|
31
|
+
name: pypi
|
|
32
|
+
url: https://pypi.org/p/pydantic-ai-rlm
|
|
33
|
+
permissions:
|
|
34
|
+
id-token: write
|
|
35
|
+
|
|
36
|
+
steps:
|
|
37
|
+
- name: Download distribution artifacts
|
|
38
|
+
uses: actions/download-artifact@v4
|
|
39
|
+
with:
|
|
40
|
+
name: python-package-distributions
|
|
41
|
+
path: dist/
|
|
42
|
+
|
|
43
|
+
- name: Publish to PyPI
|
|
44
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: pydantic-ai-rlm
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.2
|
|
4
4
|
Summary: Recursive Language Model (RLM) toolset for Pydantic AI - handle extremely large contexts
|
|
5
5
|
Author: Pydantic AI RLM Contributors
|
|
6
6
|
License-Expression: MIT
|
|
@@ -16,7 +16,7 @@ Classifier: Programming Language :: Python :: 3.12
|
|
|
16
16
|
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
17
17
|
Classifier: Typing :: Typed
|
|
18
18
|
Requires-Python: >=3.10
|
|
19
|
-
Requires-Dist: pydantic-ai>=0.1.0
|
|
19
|
+
Requires-Dist: pydantic-ai-slim[cli]>=0.1.0
|
|
20
20
|
Provides-Extra: dev
|
|
21
21
|
Requires-Dist: mypy>=1.0; extra == 'dev'
|
|
22
22
|
Requires-Dist: pytest-asyncio>=0.21; extra == 'dev'
|
|
@@ -34,6 +34,7 @@ Description-Content-Type: text/markdown
|
|
|
34
34
|
|
|
35
35
|
<p align="center">
|
|
36
36
|
<a href="https://github.com/vstorm-co/pydantic-ai-rlm">GitHub</a> •
|
|
37
|
+
<a href="https://pypi.org/project/pydantic-ai-rlm/">PyPI</a> •
|
|
37
38
|
<a href="https://github.com/vstorm-co/pydantic-ai-rlm#examples">Examples</a>
|
|
38
39
|
</p>
|
|
39
40
|
|
|
@@ -41,6 +42,7 @@ Description-Content-Type: text/markdown
|
|
|
41
42
|
<a href="https://www.python.org/downloads/"><img src="https://img.shields.io/badge/python-3.10+-blue.svg" alt="Python 3.10+"></a>
|
|
42
43
|
<a href="https://opensource.org/licenses/MIT"><img src="https://img.shields.io/badge/License-MIT-yellow.svg" alt="License: MIT"></a>
|
|
43
44
|
<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>
|
|
45
|
+
<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>
|
|
44
46
|
</p>
|
|
45
47
|
|
|
46
48
|
<p align="center">
|
|
@@ -50,6 +52,8 @@ Description-Content-Type: text/markdown
|
|
|
50
52
|
•
|
|
51
53
|
<b>Sub-Model Delegation</b>
|
|
52
54
|
•
|
|
55
|
+
<b>Grounded Citations</b>
|
|
56
|
+
•
|
|
53
57
|
<b>Fully Type-Safe</b>
|
|
54
58
|
</p>
|
|
55
59
|
|
|
@@ -201,6 +205,30 @@ result = await agent.run(
|
|
|
201
205
|
)
|
|
202
206
|
```
|
|
203
207
|
|
|
208
|
+
### Grounded Responses with Citations
|
|
209
|
+
|
|
210
|
+
Get answers with traceable citations back to the source:
|
|
211
|
+
|
|
212
|
+
```python
|
|
213
|
+
from pydantic_ai_rlm import run_rlm_analysis
|
|
214
|
+
|
|
215
|
+
# Enable grounding for citation tracking
|
|
216
|
+
result = await run_rlm_analysis(
|
|
217
|
+
context=financial_report,
|
|
218
|
+
query="What were the key revenue changes?",
|
|
219
|
+
model="openai:gpt-5",
|
|
220
|
+
grounded=True, # Returns GroundedResponse instead of str
|
|
221
|
+
)
|
|
222
|
+
|
|
223
|
+
# Response contains citation markers
|
|
224
|
+
print(result.info)
|
|
225
|
+
# "Revenue increased [1] primarily due to [2]"
|
|
226
|
+
|
|
227
|
+
# Grounding maps markers to exact quotes from the source
|
|
228
|
+
print(result.grounding)
|
|
229
|
+
# {"1": "by 45% year-over-year", "2": "expansion into Asian markets"}
|
|
230
|
+
```
|
|
231
|
+
|
|
204
232
|
---
|
|
205
233
|
|
|
206
234
|
## API Reference
|
|
@@ -215,6 +243,7 @@ agent = create_rlm_agent(
|
|
|
215
243
|
sub_model="openai:gpt-5-mini", # Model for llm_query() (optional)
|
|
216
244
|
code_timeout=60.0, # Timeout for code execution
|
|
217
245
|
custom_instructions="...", # Additional instructions
|
|
246
|
+
grounded=True, # Return GroundedResponse with citations
|
|
218
247
|
)
|
|
219
248
|
```
|
|
220
249
|
|
|
@@ -239,6 +268,11 @@ answer = await run_rlm_analysis(context, query, model="openai:gpt-5")
|
|
|
239
268
|
|
|
240
269
|
# Sync
|
|
241
270
|
answer = run_rlm_analysis_sync(context, query, model="openai:gpt-5")
|
|
271
|
+
|
|
272
|
+
# With grounding (returns GroundedResponse)
|
|
273
|
+
result = await run_rlm_analysis(context, query, grounded=True)
|
|
274
|
+
print(result.info) # Text with [N] markers
|
|
275
|
+
print(result.grounding) # {"1": "exact quote", ...}
|
|
242
276
|
```
|
|
243
277
|
|
|
244
278
|
### `RLMDependencies`
|
|
@@ -278,16 +312,19 @@ configure_logging(enabled=False)
|
|
|
278
312
|
```
|
|
279
313
|
|
|
280
314
|
Install with rich logging support for syntax highlighting and styled output:
|
|
315
|
+
|
|
281
316
|
```bash
|
|
282
317
|
pip install pydantic-ai-rlm[logging]
|
|
283
318
|
```
|
|
284
319
|
|
|
285
320
|
Or install rich separately:
|
|
321
|
+
|
|
286
322
|
```bash
|
|
287
323
|
pip install rich
|
|
288
324
|
```
|
|
289
325
|
|
|
290
326
|
When enabled, you'll see:
|
|
327
|
+
|
|
291
328
|
- Syntax-highlighted code being executed (with rich)
|
|
292
329
|
- Execution results with status indicators (SUCCESS/ERROR)
|
|
293
330
|
- Execution time for each code block
|
|
@@ -6,6 +6,7 @@
|
|
|
6
6
|
|
|
7
7
|
<p align="center">
|
|
8
8
|
<a href="https://github.com/vstorm-co/pydantic-ai-rlm">GitHub</a> •
|
|
9
|
+
<a href="https://pypi.org/project/pydantic-ai-rlm/">PyPI</a> •
|
|
9
10
|
<a href="https://github.com/vstorm-co/pydantic-ai-rlm#examples">Examples</a>
|
|
10
11
|
</p>
|
|
11
12
|
|
|
@@ -13,6 +14,7 @@
|
|
|
13
14
|
<a href="https://www.python.org/downloads/"><img src="https://img.shields.io/badge/python-3.10+-blue.svg" alt="Python 3.10+"></a>
|
|
14
15
|
<a href="https://opensource.org/licenses/MIT"><img src="https://img.shields.io/badge/License-MIT-yellow.svg" alt="License: MIT"></a>
|
|
15
16
|
<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>
|
|
17
|
+
<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>
|
|
16
18
|
</p>
|
|
17
19
|
|
|
18
20
|
<p align="center">
|
|
@@ -22,6 +24,8 @@
|
|
|
22
24
|
•
|
|
23
25
|
<b>Sub-Model Delegation</b>
|
|
24
26
|
•
|
|
27
|
+
<b>Grounded Citations</b>
|
|
28
|
+
•
|
|
25
29
|
<b>Fully Type-Safe</b>
|
|
26
30
|
</p>
|
|
27
31
|
|
|
@@ -173,6 +177,30 @@ result = await agent.run(
|
|
|
173
177
|
)
|
|
174
178
|
```
|
|
175
179
|
|
|
180
|
+
### Grounded Responses with Citations
|
|
181
|
+
|
|
182
|
+
Get answers with traceable citations back to the source:
|
|
183
|
+
|
|
184
|
+
```python
|
|
185
|
+
from pydantic_ai_rlm import run_rlm_analysis
|
|
186
|
+
|
|
187
|
+
# Enable grounding for citation tracking
|
|
188
|
+
result = await run_rlm_analysis(
|
|
189
|
+
context=financial_report,
|
|
190
|
+
query="What were the key revenue changes?",
|
|
191
|
+
model="openai:gpt-5",
|
|
192
|
+
grounded=True, # Returns GroundedResponse instead of str
|
|
193
|
+
)
|
|
194
|
+
|
|
195
|
+
# Response contains citation markers
|
|
196
|
+
print(result.info)
|
|
197
|
+
# "Revenue increased [1] primarily due to [2]"
|
|
198
|
+
|
|
199
|
+
# Grounding maps markers to exact quotes from the source
|
|
200
|
+
print(result.grounding)
|
|
201
|
+
# {"1": "by 45% year-over-year", "2": "expansion into Asian markets"}
|
|
202
|
+
```
|
|
203
|
+
|
|
176
204
|
---
|
|
177
205
|
|
|
178
206
|
## API Reference
|
|
@@ -187,6 +215,7 @@ agent = create_rlm_agent(
|
|
|
187
215
|
sub_model="openai:gpt-5-mini", # Model for llm_query() (optional)
|
|
188
216
|
code_timeout=60.0, # Timeout for code execution
|
|
189
217
|
custom_instructions="...", # Additional instructions
|
|
218
|
+
grounded=True, # Return GroundedResponse with citations
|
|
190
219
|
)
|
|
191
220
|
```
|
|
192
221
|
|
|
@@ -211,6 +240,11 @@ answer = await run_rlm_analysis(context, query, model="openai:gpt-5")
|
|
|
211
240
|
|
|
212
241
|
# Sync
|
|
213
242
|
answer = run_rlm_analysis_sync(context, query, model="openai:gpt-5")
|
|
243
|
+
|
|
244
|
+
# With grounding (returns GroundedResponse)
|
|
245
|
+
result = await run_rlm_analysis(context, query, grounded=True)
|
|
246
|
+
print(result.info) # Text with [N] markers
|
|
247
|
+
print(result.grounding) # {"1": "exact quote", ...}
|
|
214
248
|
```
|
|
215
249
|
|
|
216
250
|
### `RLMDependencies`
|
|
@@ -250,16 +284,19 @@ configure_logging(enabled=False)
|
|
|
250
284
|
```
|
|
251
285
|
|
|
252
286
|
Install with rich logging support for syntax highlighting and styled output:
|
|
287
|
+
|
|
253
288
|
```bash
|
|
254
289
|
pip install pydantic-ai-rlm[logging]
|
|
255
290
|
```
|
|
256
291
|
|
|
257
292
|
Or install rich separately:
|
|
293
|
+
|
|
258
294
|
```bash
|
|
259
295
|
pip install rich
|
|
260
296
|
```
|
|
261
297
|
|
|
262
298
|
When enabled, you'll see:
|
|
299
|
+
|
|
263
300
|
- Syntax-highlighted code being executed (with rich)
|
|
264
301
|
- Execution results with status indicators (SUCCESS/ERROR)
|
|
265
302
|
- Execution time for each code block
|
|
@@ -10,12 +10,10 @@ import random
|
|
|
10
10
|
|
|
11
11
|
from dotenv import load_dotenv
|
|
12
12
|
|
|
13
|
-
from pydantic_ai_rlm import
|
|
13
|
+
from pydantic_ai_rlm import configure_logging, run_rlm_analysis_sync
|
|
14
14
|
|
|
15
15
|
|
|
16
|
-
def generate_massive_context(
|
|
17
|
-
num_lines: int = 1_000_000, answer: str = "1298418"
|
|
18
|
-
) -> str:
|
|
16
|
+
def generate_massive_context(num_lines: int = 1_000_000, answer: str = "1298418") -> str:
|
|
19
17
|
"""Generate a massive text context with a hidden magic number."""
|
|
20
18
|
print(f"Generating massive context with {num_lines:,} lines...")
|
|
21
19
|
|
|
@@ -67,6 +65,7 @@ def main():
|
|
|
67
65
|
query=query,
|
|
68
66
|
model="openai:gpt-5",
|
|
69
67
|
sub_model="openai:gpt-5-mini",
|
|
68
|
+
grounded=True,
|
|
70
69
|
)
|
|
71
70
|
|
|
72
71
|
print(f"\nResult: {result}")
|
|
@@ -10,7 +10,7 @@ import random
|
|
|
10
10
|
|
|
11
11
|
from dotenv import load_dotenv
|
|
12
12
|
|
|
13
|
-
from pydantic_ai_rlm import
|
|
13
|
+
from pydantic_ai_rlm import configure_logging, run_rlm_analysis_sync
|
|
14
14
|
|
|
15
15
|
|
|
16
16
|
def generate_semantic_context(num_documents: int = 500) -> tuple[str, str]:
|
|
@@ -67,10 +67,7 @@ def generate_semantic_context(num_documents: int = 500) -> tuple[str, str]:
|
|
|
67
67
|
"Ware",
|
|
68
68
|
"Hub",
|
|
69
69
|
]
|
|
70
|
-
companies = [
|
|
71
|
-
f"{random.choice(prefixes)}{random.choice(suffixes)}"
|
|
72
|
-
for _ in range(num_documents)
|
|
73
|
-
]
|
|
70
|
+
companies = [f"{random.choice(prefixes)}{random.choice(suffixes)}" for _ in range(num_documents)]
|
|
74
71
|
|
|
75
72
|
# Pick a random company to be the bankrupt one
|
|
76
73
|
bankrupt_idx = random.randint(100, num_documents - 100)
|
|
@@ -112,9 +109,7 @@ def generate_semantic_context(num_documents: int = 500) -> tuple[str, str]:
|
|
|
112
109
|
random.shuffle(documents)
|
|
113
110
|
|
|
114
111
|
context = "\n\n".join(documents)
|
|
115
|
-
print(
|
|
116
|
-
f"Bankrupt company: {bankrupt_company} (was at original index {bankrupt_idx})"
|
|
117
|
-
)
|
|
112
|
+
print(f"Bankrupt company: {bankrupt_company} (was at original index {bankrupt_idx})")
|
|
118
113
|
|
|
119
114
|
return context, bankrupt_company
|
|
120
115
|
|
|
@@ -147,6 +142,7 @@ def main():
|
|
|
147
142
|
query=query,
|
|
148
143
|
model="openai:gpt-5",
|
|
149
144
|
sub_model="openai:gpt-5-mini",
|
|
145
|
+
grounded=True,
|
|
150
146
|
)
|
|
151
147
|
|
|
152
148
|
print(f"\nResult: {result}")
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "pydantic-ai-rlm"
|
|
3
|
-
version = "0.1.
|
|
3
|
+
version = "0.1.2"
|
|
4
4
|
description = "Recursive Language Model (RLM) toolset for Pydantic AI - handle extremely large contexts"
|
|
5
5
|
readme = "README.md"
|
|
6
6
|
license = "MIT"
|
|
@@ -27,7 +27,7 @@ classifiers = [
|
|
|
27
27
|
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
28
28
|
"Typing :: Typed",
|
|
29
29
|
]
|
|
30
|
-
dependencies = ["pydantic-ai>=0.1.0"]
|
|
30
|
+
dependencies = ["pydantic-ai-slim[cli]>=0.1.0"]
|
|
31
31
|
|
|
32
32
|
[project.optional-dependencies]
|
|
33
33
|
dev = ["pytest>=7.0", "pytest-asyncio>=0.21", "mypy>=1.0", "ruff>=0.1"]
|
|
@@ -1,7 +1,9 @@
|
|
|
1
1
|
from .agent import create_rlm_agent, run_rlm_analysis, run_rlm_analysis_sync
|
|
2
2
|
from .dependencies import ContextType, RLMConfig, RLMDependencies
|
|
3
3
|
from .logging import configure_logging
|
|
4
|
+
from .models import GroundedResponse
|
|
4
5
|
from .prompts import (
|
|
6
|
+
GROUNDING_INSTRUCTIONS,
|
|
5
7
|
LLM_QUERY_INSTRUCTIONS,
|
|
6
8
|
RLM_INSTRUCTIONS,
|
|
7
9
|
build_rlm_instructions,
|
|
@@ -13,9 +15,11 @@ from .toolset import (
|
|
|
13
15
|
)
|
|
14
16
|
|
|
15
17
|
__all__ = [
|
|
18
|
+
"GROUNDING_INSTRUCTIONS",
|
|
16
19
|
"LLM_QUERY_INSTRUCTIONS",
|
|
17
20
|
"RLM_INSTRUCTIONS",
|
|
18
21
|
"ContextType",
|
|
22
|
+
"GroundedResponse",
|
|
19
23
|
"REPLEnvironment",
|
|
20
24
|
"REPLResult",
|
|
21
25
|
"RLMConfig",
|
|
@@ -1,21 +1,45 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
|
-
from typing import Any
|
|
3
|
+
from typing import Any, Literal, overload
|
|
4
4
|
|
|
5
5
|
from pydantic_ai import Agent, UsageLimits
|
|
6
6
|
|
|
7
7
|
from .dependencies import ContextType, RLMConfig, RLMDependencies
|
|
8
|
+
from .models import GroundedResponse
|
|
8
9
|
from .prompts import build_rlm_instructions
|
|
9
10
|
from .toolset import create_rlm_toolset
|
|
10
11
|
|
|
11
12
|
|
|
13
|
+
@overload
|
|
12
14
|
def create_rlm_agent(
|
|
13
15
|
model: str = "openai:gpt-5",
|
|
14
16
|
sub_model: str | None = None,
|
|
15
17
|
code_timeout: float = 60.0,
|
|
16
|
-
include_example_instructions: bool = True,
|
|
17
18
|
custom_instructions: str | None = None,
|
|
18
|
-
|
|
19
|
+
*,
|
|
20
|
+
grounded: Literal[False] = False,
|
|
21
|
+
) -> Agent[RLMDependencies, str]: ...
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
@overload
|
|
25
|
+
def create_rlm_agent(
|
|
26
|
+
model: str = "openai:gpt-5",
|
|
27
|
+
sub_model: str | None = None,
|
|
28
|
+
code_timeout: float = 60.0,
|
|
29
|
+
custom_instructions: str | None = None,
|
|
30
|
+
*,
|
|
31
|
+
grounded: Literal[True],
|
|
32
|
+
) -> Agent[RLMDependencies, GroundedResponse]: ...
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def create_rlm_agent(
|
|
36
|
+
model: str = "openai:gpt-5",
|
|
37
|
+
sub_model: str | None = None,
|
|
38
|
+
code_timeout: float = 60.0,
|
|
39
|
+
custom_instructions: str | None = None,
|
|
40
|
+
*,
|
|
41
|
+
grounded: bool = False,
|
|
42
|
+
) -> Agent[RLMDependencies, str] | Agent[RLMDependencies, GroundedResponse]:
|
|
19
43
|
"""
|
|
20
44
|
Create a Pydantic AI agent with REPL code execution capabilities.
|
|
21
45
|
|
|
@@ -26,11 +50,12 @@ def create_rlm_agent(
|
|
|
26
50
|
available in the REPL, allowing the agent to delegate sub-queries.
|
|
27
51
|
Example: "openai:gpt-5-mini" or "anthropic:claude-3-haiku-20240307"
|
|
28
52
|
code_timeout: Timeout for code execution in seconds
|
|
29
|
-
include_example_instructions: Include detailed examples in instructions
|
|
30
53
|
custom_instructions: Additional instructions to append
|
|
54
|
+
grounded: If True, return a GroundedResponse with citation markers
|
|
31
55
|
|
|
32
56
|
Returns:
|
|
33
|
-
Configured Agent instance
|
|
57
|
+
Configured Agent instance. Returns Agent[RLMDependencies, GroundedResponse]
|
|
58
|
+
when grounded=True, otherwise Agent[RLMDependencies, str].
|
|
34
59
|
|
|
35
60
|
Example:
|
|
36
61
|
```python
|
|
@@ -48,19 +73,28 @@ def create_rlm_agent(
|
|
|
48
73
|
)
|
|
49
74
|
result = await agent.run("What are the main themes?", deps=deps)
|
|
50
75
|
print(result.output)
|
|
76
|
+
|
|
77
|
+
# Create grounded agent
|
|
78
|
+
grounded_agent = create_rlm_agent(model="openai:gpt-5", grounded=True)
|
|
79
|
+
result = await grounded_agent.run("What happened?", deps=deps)
|
|
80
|
+
print(result.output.info) # Response with [N] markers
|
|
81
|
+
print(result.output.grounding) # {"1": "exact quote", ...}
|
|
51
82
|
```
|
|
52
83
|
"""
|
|
53
84
|
toolset = create_rlm_toolset(code_timeout=code_timeout, sub_model=sub_model)
|
|
54
85
|
|
|
55
86
|
instructions = build_rlm_instructions(
|
|
56
87
|
include_llm_query=sub_model is not None,
|
|
88
|
+
include_grounding=grounded,
|
|
57
89
|
custom_suffix=custom_instructions,
|
|
58
90
|
)
|
|
59
91
|
|
|
60
|
-
|
|
92
|
+
output_type: type[str] | type[GroundedResponse] = GroundedResponse if grounded else str
|
|
93
|
+
|
|
94
|
+
agent: Agent[RLMDependencies, Any] = Agent(
|
|
61
95
|
model,
|
|
62
96
|
deps_type=RLMDependencies,
|
|
63
|
-
output_type=
|
|
97
|
+
output_type=output_type,
|
|
64
98
|
toolsets=[toolset],
|
|
65
99
|
instructions=instructions,
|
|
66
100
|
)
|
|
@@ -68,6 +102,34 @@ def create_rlm_agent(
|
|
|
68
102
|
return agent
|
|
69
103
|
|
|
70
104
|
|
|
105
|
+
@overload
|
|
106
|
+
async def run_rlm_analysis(
|
|
107
|
+
context: ContextType,
|
|
108
|
+
query: str,
|
|
109
|
+
model: str = "openai:gpt-5",
|
|
110
|
+
sub_model: str | None = None,
|
|
111
|
+
config: RLMConfig | None = None,
|
|
112
|
+
max_tool_calls: int = 50,
|
|
113
|
+
*,
|
|
114
|
+
grounded: Literal[False] = False,
|
|
115
|
+
**agent_kwargs: Any,
|
|
116
|
+
) -> str: ...
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
@overload
|
|
120
|
+
async def run_rlm_analysis(
|
|
121
|
+
context: ContextType,
|
|
122
|
+
query: str,
|
|
123
|
+
model: str = "openai:gpt-5",
|
|
124
|
+
sub_model: str | None = None,
|
|
125
|
+
config: RLMConfig | None = None,
|
|
126
|
+
max_tool_calls: int = 50,
|
|
127
|
+
*,
|
|
128
|
+
grounded: Literal[True],
|
|
129
|
+
**agent_kwargs: Any,
|
|
130
|
+
) -> GroundedResponse: ...
|
|
131
|
+
|
|
132
|
+
|
|
71
133
|
async def run_rlm_analysis(
|
|
72
134
|
context: ContextType,
|
|
73
135
|
query: str,
|
|
@@ -75,8 +137,10 @@ async def run_rlm_analysis(
|
|
|
75
137
|
sub_model: str | None = None,
|
|
76
138
|
config: RLMConfig | None = None,
|
|
77
139
|
max_tool_calls: int = 50,
|
|
140
|
+
*,
|
|
141
|
+
grounded: bool = False,
|
|
78
142
|
**agent_kwargs: Any,
|
|
79
|
-
) -> str:
|
|
143
|
+
) -> str | GroundedResponse:
|
|
80
144
|
"""
|
|
81
145
|
Convenience function to run RLM analysis on a context.
|
|
82
146
|
|
|
@@ -89,25 +153,36 @@ async def run_rlm_analysis(
|
|
|
89
153
|
available in the REPL, allowing the agent to delegate sub-queries.
|
|
90
154
|
config: Optional RLMConfig for customization
|
|
91
155
|
max_tool_calls: Maximum tool calls allowed
|
|
156
|
+
grounded: If True, return a GroundedResponse with citation markers
|
|
92
157
|
**agent_kwargs: Additional arguments passed to create_rlm_agent()
|
|
93
158
|
|
|
94
159
|
Returns:
|
|
95
|
-
The agent's final answer
|
|
160
|
+
The agent's final answer. Returns GroundedResponse when grounded=True,
|
|
161
|
+
otherwise returns str.
|
|
96
162
|
|
|
97
163
|
Example:
|
|
98
164
|
```python
|
|
99
165
|
from pydantic_ai_rlm import run_rlm_analysis
|
|
100
166
|
|
|
101
|
-
#
|
|
167
|
+
# Standard string response
|
|
102
168
|
answer = await run_rlm_analysis(
|
|
103
169
|
context=huge_document,
|
|
104
170
|
query="Find the magic number hidden in the text",
|
|
105
171
|
sub_model="openai:gpt-5-mini",
|
|
106
172
|
)
|
|
107
173
|
print(answer)
|
|
174
|
+
|
|
175
|
+
# Grounded response with citations
|
|
176
|
+
result = await run_rlm_analysis(
|
|
177
|
+
context=document,
|
|
178
|
+
query="What was the revenue change?",
|
|
179
|
+
grounded=True,
|
|
180
|
+
)
|
|
181
|
+
print(result.info) # "Revenue grew [1]..."
|
|
182
|
+
print(result.grounding) # {"1": "increased by 45%", ...}
|
|
108
183
|
```
|
|
109
184
|
"""
|
|
110
|
-
agent = create_rlm_agent(model=model, sub_model=sub_model, **agent_kwargs)
|
|
185
|
+
agent = create_rlm_agent(model=model, sub_model=sub_model, grounded=grounded, **agent_kwargs)
|
|
111
186
|
|
|
112
187
|
effective_config = config or RLMConfig()
|
|
113
188
|
if sub_model and not effective_config.sub_model:
|
|
@@ -127,6 +202,7 @@ async def run_rlm_analysis(
|
|
|
127
202
|
return result.output
|
|
128
203
|
|
|
129
204
|
|
|
205
|
+
@overload
|
|
130
206
|
def run_rlm_analysis_sync(
|
|
131
207
|
context: ContextType,
|
|
132
208
|
query: str,
|
|
@@ -134,14 +210,63 @@ def run_rlm_analysis_sync(
|
|
|
134
210
|
sub_model: str | None = None,
|
|
135
211
|
config: RLMConfig | None = None,
|
|
136
212
|
max_tool_calls: int = 50,
|
|
213
|
+
*,
|
|
214
|
+
grounded: Literal[False] = False,
|
|
137
215
|
**agent_kwargs: Any,
|
|
138
|
-
) -> str:
|
|
216
|
+
) -> str: ...
|
|
217
|
+
|
|
218
|
+
|
|
219
|
+
@overload
|
|
220
|
+
def run_rlm_analysis_sync(
|
|
221
|
+
context: ContextType,
|
|
222
|
+
query: str,
|
|
223
|
+
model: str = "openai:gpt-5",
|
|
224
|
+
sub_model: str | None = None,
|
|
225
|
+
config: RLMConfig | None = None,
|
|
226
|
+
max_tool_calls: int = 50,
|
|
227
|
+
*,
|
|
228
|
+
grounded: Literal[True],
|
|
229
|
+
**agent_kwargs: Any,
|
|
230
|
+
) -> GroundedResponse: ...
|
|
231
|
+
|
|
232
|
+
|
|
233
|
+
def run_rlm_analysis_sync(
|
|
234
|
+
context: ContextType,
|
|
235
|
+
query: str,
|
|
236
|
+
model: str = "openai:gpt-5",
|
|
237
|
+
sub_model: str | None = None,
|
|
238
|
+
config: RLMConfig | None = None,
|
|
239
|
+
max_tool_calls: int = 50,
|
|
240
|
+
*,
|
|
241
|
+
grounded: bool = False,
|
|
242
|
+
**agent_kwargs: Any,
|
|
243
|
+
) -> str | GroundedResponse:
|
|
139
244
|
"""
|
|
140
245
|
Synchronous version of run_rlm_analysis.
|
|
141
246
|
|
|
142
247
|
See run_rlm_analysis() for full documentation.
|
|
248
|
+
|
|
249
|
+
Example:
|
|
250
|
+
```python
|
|
251
|
+
from pydantic_ai_rlm import run_rlm_analysis_sync
|
|
252
|
+
|
|
253
|
+
# Standard string response
|
|
254
|
+
answer = run_rlm_analysis_sync(
|
|
255
|
+
context=document,
|
|
256
|
+
query="What happened?",
|
|
257
|
+
)
|
|
258
|
+
|
|
259
|
+
# Grounded response with citations
|
|
260
|
+
result = run_rlm_analysis_sync(
|
|
261
|
+
context=document,
|
|
262
|
+
query="What was the revenue change?",
|
|
263
|
+
grounded=True,
|
|
264
|
+
)
|
|
265
|
+
print(result.info) # "Revenue grew [1]..."
|
|
266
|
+
print(result.grounding) # {"1": "increased by 45%", ...}
|
|
267
|
+
```
|
|
143
268
|
"""
|
|
144
|
-
agent = create_rlm_agent(model=model, sub_model=sub_model, **agent_kwargs)
|
|
269
|
+
agent = create_rlm_agent(model=model, sub_model=sub_model, grounded=grounded, **agent_kwargs)
|
|
145
270
|
|
|
146
271
|
effective_config = config or RLMConfig()
|
|
147
272
|
if sub_model and not effective_config.sub_model:
|
|
@@ -47,7 +47,7 @@ class RLMLogger:
|
|
|
47
47
|
)
|
|
48
48
|
self.console.print(panel)
|
|
49
49
|
else:
|
|
50
|
-
print(f"\n{'='*50}")
|
|
50
|
+
print(f"\n{'=' * 50}")
|
|
51
51
|
print("CODE EXECUTION")
|
|
52
52
|
print("=" * 50)
|
|
53
53
|
print(code)
|
|
@@ -141,7 +141,7 @@ class RLMLogger:
|
|
|
141
141
|
def _log_result_plain(self, result: REPLResult) -> None:
|
|
142
142
|
"""Log result using plain text."""
|
|
143
143
|
status = "SUCCESS" if result.success else "ERROR"
|
|
144
|
-
print(f"\n{'='*50}")
|
|
144
|
+
print(f"\n{'=' * 50}")
|
|
145
145
|
print(f"RESULT: {status} (executed in {result.execution_time:.3f}s)")
|
|
146
146
|
print("=" * 50)
|
|
147
147
|
|
|
@@ -159,11 +159,7 @@ class RLMLogger:
|
|
|
159
159
|
stderr = stderr[:1000] + "\n... (truncated)"
|
|
160
160
|
print(stderr)
|
|
161
161
|
|
|
162
|
-
user_vars = {
|
|
163
|
-
k: v
|
|
164
|
-
for k, v in result.locals.items()
|
|
165
|
-
if not k.startswith("_") and k not in ("context", "json", "re", "os")
|
|
166
|
-
}
|
|
162
|
+
user_vars = {k: v for k, v in result.locals.items() if not k.startswith("_") and k not in ("context", "json", "re", "os")}
|
|
167
163
|
if user_vars:
|
|
168
164
|
print("\nVariables:")
|
|
169
165
|
for name, value in list(user_vars.items())[:10]:
|
|
@@ -198,7 +194,7 @@ class RLMLogger:
|
|
|
198
194
|
)
|
|
199
195
|
self.console.print(panel)
|
|
200
196
|
else:
|
|
201
|
-
print(f"\n{'='*50}")
|
|
197
|
+
print(f"\n{'=' * 50}")
|
|
202
198
|
print("LLM QUERY")
|
|
203
199
|
print("=" * 50)
|
|
204
200
|
display_prompt = prompt
|
|
@@ -226,7 +222,7 @@ class RLMLogger:
|
|
|
226
222
|
)
|
|
227
223
|
self.console.print(panel)
|
|
228
224
|
else:
|
|
229
|
-
print(f"\n{'='*50}")
|
|
225
|
+
print(f"\n{'=' * 50}")
|
|
230
226
|
print("LLM RESPONSE")
|
|
231
227
|
print("=" * 50)
|
|
232
228
|
display_response = response
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
"""Pydantic models for structured RLM outputs."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from pydantic import BaseModel, Field
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class GroundedResponse(BaseModel):
|
|
9
|
+
"""A response with citation markers mapping to exact quotes from source documents.
|
|
10
|
+
|
|
11
|
+
Example:
|
|
12
|
+
```python
|
|
13
|
+
response = GroundedResponse(
|
|
14
|
+
info="Revenue grew [1] driven by expansion [2]", grounding={"1": "increased by 45%", "2": "new markets in Asia"}
|
|
15
|
+
)
|
|
16
|
+
```
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
info: str = Field(description="Response text with citation markers like [1]")
|
|
20
|
+
grounding: dict[str, str] = Field(
|
|
21
|
+
default_factory=dict,
|
|
22
|
+
description="Mapping from citation markers to exact quotes from the source",
|
|
23
|
+
)
|
|
@@ -50,6 +50,61 @@ print(f"Final answer: {results}")
|
|
|
50
50
|
4. **Be thorough** - For needle-in-haystack, search the entire context
|
|
51
51
|
"""
|
|
52
52
|
|
|
53
|
+
GROUNDING_INSTRUCTIONS = """
|
|
54
|
+
|
|
55
|
+
## Grounding Requirements
|
|
56
|
+
|
|
57
|
+
Your response MUST include grounded citations. This means:
|
|
58
|
+
|
|
59
|
+
1. **Citation Format**: Use markers like `[1]`, `[2]`, etc. in your response text
|
|
60
|
+
2. **Exact Quotes**: Each marker must map to an EXACT quote from the source context (verbatim, no paraphrasing)
|
|
61
|
+
3. **Quote Length**: Each quote should be 10-200 characters - enough to be meaningful but not too long
|
|
62
|
+
4. **Consecutive Numbering**: Number citations consecutively starting from 1
|
|
63
|
+
|
|
64
|
+
### Output Format
|
|
65
|
+
|
|
66
|
+
Your final answer must be valid JSON with this structure:
|
|
67
|
+
```json
|
|
68
|
+
{
|
|
69
|
+
"info": "The document states that X Y Z [1]. Additionally, A B C [2]",
|
|
70
|
+
"grounding": {
|
|
71
|
+
"1": "exact quote from source",
|
|
72
|
+
"2": "another exact quote from source"
|
|
73
|
+
}
|
|
74
|
+
}
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
### Example
|
|
78
|
+
|
|
79
|
+
If the context contains: "The company's revenue increased by 45% in Q3 2024, driven by expansion into new markets in Asia."
|
|
80
|
+
|
|
81
|
+
Your response should look like:
|
|
82
|
+
```json
|
|
83
|
+
{
|
|
84
|
+
"info": "Revenue showed strong growth [1] with geographic expansion being a key driver [2].",
|
|
85
|
+
"grounding": {
|
|
86
|
+
"1": "revenue increased by 45% in Q3 2024",
|
|
87
|
+
"2": "driven by expansion into new markets in Asia"
|
|
88
|
+
}
|
|
89
|
+
}
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
### Finding Quotes in Code
|
|
93
|
+
|
|
94
|
+
Use this approach to find and verify exact quotes:
|
|
95
|
+
```python
|
|
96
|
+
# Find a specific phrase in context
|
|
97
|
+
search_term = "revenue"
|
|
98
|
+
idx = context.lower().find(search_term)
|
|
99
|
+
if idx != -1:
|
|
100
|
+
# Extract surrounding context for the quote
|
|
101
|
+
quote = context[max(0, idx):idx+100]
|
|
102
|
+
print(f"Found: {quote}")
|
|
103
|
+
```
|
|
104
|
+
|
|
105
|
+
**Important**: Every citation marker in your `info` field MUST have a corresponding entry in `grounding`. Only output the JSON object, no additional text.
|
|
106
|
+
"""
|
|
107
|
+
|
|
53
108
|
LLM_QUERY_INSTRUCTIONS = """
|
|
54
109
|
|
|
55
110
|
## Sub-LLM Queries
|
|
@@ -93,14 +148,15 @@ print(result)
|
|
|
93
148
|
|
|
94
149
|
def build_rlm_instructions(
|
|
95
150
|
include_llm_query: bool = False,
|
|
151
|
+
include_grounding: bool = False,
|
|
96
152
|
custom_suffix: str | None = None,
|
|
97
153
|
) -> str:
|
|
98
154
|
"""
|
|
99
155
|
Build RLM instructions with optional customization.
|
|
100
156
|
|
|
101
157
|
Args:
|
|
102
|
-
include_examples: Whether to include detailed examples
|
|
103
158
|
include_llm_query: Whether to include llm_query() documentation
|
|
159
|
+
include_grounding: Whether to include grounding/citation instructions
|
|
104
160
|
custom_suffix: Additional instructions to append
|
|
105
161
|
|
|
106
162
|
Returns:
|
|
@@ -109,8 +165,10 @@ def build_rlm_instructions(
|
|
|
109
165
|
base = RLM_INSTRUCTIONS
|
|
110
166
|
|
|
111
167
|
if include_llm_query:
|
|
112
|
-
|
|
113
|
-
|
|
168
|
+
base = f"{base}{LLM_QUERY_INSTRUCTIONS}"
|
|
169
|
+
|
|
170
|
+
if include_grounding:
|
|
171
|
+
base = f"{base}{GROUNDING_INSTRUCTIONS}"
|
|
114
172
|
|
|
115
173
|
if custom_suffix:
|
|
116
174
|
base = f"{base}\n\n## Additional Instructions\n\n{custom_suffix}"
|
|
@@ -7,6 +7,7 @@ import os
|
|
|
7
7
|
import shutil
|
|
8
8
|
import sys
|
|
9
9
|
import tempfile
|
|
10
|
+
import textwrap
|
|
10
11
|
import threading
|
|
11
12
|
import time
|
|
12
13
|
from contextlib import contextmanager
|
|
@@ -328,6 +329,9 @@ with open(r'{context_path}', 'r', encoding='utf-8') as f:
|
|
|
328
329
|
Returns:
|
|
329
330
|
REPLResult with stdout, stderr, locals, and timing
|
|
330
331
|
"""
|
|
332
|
+
# Normalize code: remove common leading whitespace and strip
|
|
333
|
+
code = textwrap.dedent(code).strip()
|
|
334
|
+
|
|
331
335
|
start_time = time.time()
|
|
332
336
|
success = True
|
|
333
337
|
stdout_content = ""
|
|
@@ -1 +0,0 @@
|
|
|
1
|
-
*
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|