danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Replace ask() with llm(). Add and document the programs written in tests

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-29 19:15:41 +02:00
parent 5e7832a6ea
commit fa4d7731d7
No known key found for this signature in database
4 changed files with 178 additions and 76 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
haiku_rag_slim/haiku/rag/agents/rlm/prompts.py
View file

@ -27,9 +27,10 @@ Get the structured DoclingDocument object for advanced analysis.
Returns a DoclingDocument object, or None if not found.
See "DoclingDocument API" section below for how to use it.
### ask(question) -> str
Ask a question using the QA agent with RAG. Returns the answer as a string.
Use this for semantic analysis that benefits from LLM reasoning.
### llm(prompt) -> str
Call an LLM directly with the given prompt. Returns the response as a string.
Use this for classification, summarization, extraction, or any task where you
already have the content and just need LLM reasoning.
## Standard Library Modules
You can import: json, re, collections, math, statistics, itertools, functools, datetime, typing
@ -41,7 +42,7 @@ You can import: json, re, collections, math, statistics, itertools, functools, d
3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
4. **Use print() Liberally**: The REPL captures stdout - print intermediate results to see what you're working with.
5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and use collections.
6. **Use ask() for Semantic Analysis**: When you need to understand meaning or interpret content, use the ask() function.
6. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use llm().
7. **Cite Your Sources**: Track which documents/chunks informed your answer for citation.
## DoclingDocument API
@ -112,13 +113,13 @@ for r in results:
print(f"Average: ${sum(numbers)/len(numbers):,.2f}")
```
### Using ask() for semantic analysis
### Using llm() for classification
```python
# First search to find relevant content
results = search("machine learning approaches")
# Then use ask() to synthesize an answer
summary = ask("What are the main machine learning approaches discussed?")
print(summary)
# Get document content
content = get_document("Q1 Report")
# Use llm() to classify sentiment
sentiment = llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
print(sentiment)
```
## Workflow

43
haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py
View file

@ -139,7 +139,7 @@ class REPLEnvironment:
"list_documents": self._make_list_documents(),
"get_document": self._make_get_document(),
"get_docling_document": self._make_get_docling_document(),
"ask": self._make_ask(),
"llm": self._make_llm(),
}
self.locals: dict[str, Any] = {}
@ -246,38 +246,23 @@ class REPLEnvironment:
return get_docling_document
def _make_ask(self):
"""Create sync ask function that uses QA agent."""
def _make_llm(self):
"""Create sync llm function for plain LLM calls without RAG."""
def ask(question: str) -> str:
async def _ask():
answer, citations = await self.client.ask(
question, filter=self.context.filter
)
for c in citations:
for sr in self.context.search_results:
if sr.chunk_id == c.chunk_id:
break
else:
from haiku.rag.store.models import SearchResult
def llm(prompt: str) -> str:
async def _llm():
from pydantic_ai import Agent
self.context.search_results.append(
SearchResult(
chunk_id=c.chunk_id,
document_id=c.document_id,
document_title=c.document_title or "",
document_uri=c.document_uri,
content=c.content,
score=1.0,
page_numbers=c.page_numbers,
headings=c.headings or [],
)
)
return answer
from haiku.rag.utils import get_model
return self._run_async_from_thread(_ask())
model = get_model(self.config.model)
agent: Agent[None, str] = Agent(model, output_type=str)
result = await agent.run(prompt)
return result.output
return ask
return self._run_async_from_thread(_llm())
return llm
def _safe_import(
self,

154
tests/agents/rlm/test_agent.py
View file

@ -119,7 +119,12 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_count_documents(self, allow_model_requests, temp_db_path):
"""Test RLM agent can count documents."""
"""Test RLM agent can count documents.
Agent program:
docs = list_documents(limit=1000)
print(len(docs))
"""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as client:
@ -134,7 +139,24 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_aggregation(self, allow_model_requests, temp_db_path):
"""Test RLM agent can perform aggregation across documents."""
"""Test RLM agent can perform aggregation across documents.
Agent program:
import re
revs = {}
for d in ['Q1 Report', 'Q2 Report', 'Q3 Report']:
content = get_document(d)
if content:
vals = re.findall(r'\\$([\\d,]+)', content)
if vals:
rev = sum(int(v.replace(',', '')) for v in vals)
else:
rev = None
else:
rev = None
revs[d] = rev
print(revs)
"""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as client:
@ -157,7 +179,15 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_with_filter(self, allow_model_requests, temp_db_path):
"""Test RLM agent respects filter parameter."""
"""Test RLM agent respects filter parameter.
Agent program:
docs = list_documents(limit=1000)
print(len(docs))
print(docs[:5])
The filter is applied via context, so list_documents() only sees "Cats".
"""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as client:
@ -177,7 +207,17 @@ class TestClientRLMIntegration:
async def test_rlm_docling_document_structure(
self, allow_model_requests, temp_db_path
):
"""Test RLM agent can analyze document structure using DoclingDocument."""
"""Test RLM agent can analyze document structure using DoclingDocument.
Agent program:
docs = list_documents(limit=20)
print(docs)
doc = get_docling_document('<doc_id>')
print(doc.name)
print('tables:', len(doc.tables))
print('pictures:', len(doc.pictures))
"""
from pathlib import Path
from haiku.rag.client import HaikuRAG
@ -197,3 +237,109 @@ class TestClientRLMIntegration:
# The doclaynet.pdf has 1 table and 1 picture
assert "1" in answer
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_semantic_analysis_with_llm(
self, allow_model_requests, temp_db_path
):
"""Test RLM agent can use llm() for semantic analysis combined with computation.
Agent program:
docs = list_documents(limit=100)
print(len(docs))
print([d['title'] for d in docs[:20]])
sentiments = {}
for title in ['Q1 Update', 'Q2 Update', 'Q3 Update']:
content = get_document(title)
if content:
result = llm(f"Classify sentiment as positive/negative/mixed: {content}")
sentiments[title] = result
print(sentiments)
"""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
"The new product launch exceeded expectations. Sales grew 40% "
"and customer feedback has been overwhelmingly positive. "
"Team morale is at an all-time high.",
title="Q1 Update",
)
await client.create_document(
"We faced significant challenges this quarter. Supply chain issues "
"caused delays, and we missed our revenue target by 15%. "
"Several key employees left the company.",
title="Q2 Update",
)
await client.create_document(
"Mixed results this quarter. While product quality improved, "
"marketing campaigns underperformed. Revenue was flat compared "
"to last year but customer retention increased.",
title="Q3 Update",
)
answer = await client.rlm(
"Analyze the sentiment of each quarterly update. "
"How many quarters were positive, negative, and mixed?"
)
# Should identify: Q1=positive, Q2=negative, Q3=mixed
assert "positive" in answer.lower()
assert "negative" in answer.lower()
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_search_and_extract(self, allow_model_requests, temp_db_path):
"""Test RLM agent can use search() to find content and extract information.
Agent program:
results = search("document element types", limit=20)
print(len(results))
for r in results[:5]:
print(r['document_title'], r['chunk_id'], r['score'])
print(r['content'][:200])
results = search("DocBank element types", limit=10)
...
"""
from pathlib import Path
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig
pdf_path = Path("tests/data/doclaynet.pdf")
config = AppConfig()
config.processing.conversion_options.do_ocr = False
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document_from_source(pdf_path)
answer = await client.rlm(
"Search for content about document element types or labels. "
"What are all the different document element types mentioned? "
"List them all."
)
# The doclaynet.pdf defines exactly 11 class labels for document elements
# Normalize Unicode hyphens (U+2011 non-breaking hyphen) to regular hyphens
answer_lower = answer.lower().replace("\u2011", "-")
expected_labels = [
"caption",
"footnote",
"formula",
"list-item",
"page-footer",
"page-header",
"picture",
"section-header",
"table",
"text",
"title",
]
for label in expected_labels:
# Allow for hyphen or space variants
assert (
label in answer_lower or label.replace("-", " ") in answer_lower
), f"Missing label: {label}"

36
tests/agents/rlm/test_sandbox.py
View file

@ -345,20 +345,11 @@ class TestHaikuRAGBridgeFunctions:
assert "True" in result.stdout
@pytest.mark.asyncio
async def test_ask(self, repl_env_empty):
"""Test ask function calls client with correct args."""
from unittest.mock import AsyncMock
repl_env_empty.client.ask = AsyncMock(return_value=("The fox is brown.", []))
result = await repl_env_empty.execute_async(
"answer = ask('What color is the fox?')\nprint('fox' in answer.lower())"
)
async def test_llm(self, repl_env_empty):
"""Test llm function is available in sandbox."""
result = await repl_env_empty.execute_async("print(callable(llm))")
assert result.success
assert "True" in result.stdout
repl_env_empty.client.ask.assert_called_once_with(
"What color is the fox?", filter=None
)
class TestSandboxExecution:
@ -490,27 +481,6 @@ class TestContextFilter:
limit=10, offset=0, filter="title = 'Report'"
)
@pytest.mark.asyncio
async def test_context_filter_applied_to_ask(self, temp_db_path):
"""ask applies context filter automatically."""
from unittest.mock import AsyncMock
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
async with HaikuRAG(temp_db_path, create=True) as client:
context = RLMContext(filter="metadata->>'category' = 'finance'")
repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
client.ask = AsyncMock(return_value=("Answer", []))
await repl.execute_async("ask('What is the revenue?')")
client.ask.assert_called_once_with(
"What is the revenue?", filter="metadata->>'category' = 'finance'"
)
class TestSecurityEscapes:
"""Test that common security escape attempts are blocked."""