diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/prompts.py b/haiku_rag_slim/haiku/rag/agents/rlm/prompts.py index 90171fbd..f057b94f 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/prompts.py +++ b/haiku_rag_slim/haiku/rag/agents/rlm/prompts.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py b/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py index e10a5410..6db78141 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py +++ b/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py @@ -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, diff --git a/tests/agents/rlm/test_agent.py b/tests/agents/rlm/test_agent.py index f8bb2356..1e47d0ea 100644 --- a/tests/agents/rlm/test_agent.py +++ b/tests/agents/rlm/test_agent.py @@ -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('') + 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}" diff --git a/tests/agents/rlm/test_sandbox.py b/tests/agents/rlm/test_sandbox.py index bacb326b..03ba6917 100644 --- a/tests/agents/rlm/test_sandbox.py +++ b/tests/agents/rlm/test_sandbox.py @@ -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."""