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

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
+### llm(prompt) -> str
-Ask a question using the QA agent with RAG. Returns the answer as a string.
+Call an LLM directly with the given prompt. Returns the response as a string.
-Use this for semantic analysis that benefits from LLM reasoning.
+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
+# Get document content
-results = search("machine learning approaches")
+content = get_document("Q1 Report")
-# Then use ask() to synthesize an answer
+# Use llm() to classify sentiment
-summary = ask("What are the main machine learning approaches discussed?")
+sentiment = llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
-print(summary)
+print(sentiment)
 ```
 
 ## Workflow

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):
+    def _make_llm(self):
-        """Create sync ask function that uses QA agent."""
+        """Create sync llm function for plain LLM calls without RAG."""
 
-        def ask(question: str) -> str:
+        def llm(prompt: str) -> str:
-            async def _ask():
+            async def _llm():
-                answer, citations = await self.client.ask(
+                from pydantic_ai import Agent
-                    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
 
-                        self.context.search_results.append(
+                from haiku.rag.utils import get_model
-                            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
 
-            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,

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('<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}"

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):
+    async def test_llm(self, repl_env_empty):
-        """Test ask function calls client with correct args."""
+        """Test llm function is available in sandbox."""
-        from unittest.mock import AsyncMock
+        result = await repl_env_empty.execute_async("print(callable(llm))")
-
-        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())"
-        )
         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."""