Use the expand_context when searching in the qa agent

docs/configuration.md

@@ -172,4 +172,8 @@ DEFAULT_DATA_DIR="/path/to/data"
 ```bash
 # Chunk size for document processing
 CHUNK_SIZE=256
+
+# Number of adjacent chunks to include before/after retrieved chunks for context
+# 0 = no expansion (default), 1 = include 1 chunk before and after, etc.
+CONTEXT_CHUNK_RADIUS=0
 ```

docs/python.md

@@ -130,6 +130,24 @@ for chunk, relevance_score in results:
     print(f"Document metadata: {chunk.document_meta}")
 ```
 
+### Expanding Search Context
+
+Expand search results with adjacent chunks for more complete context:
+
+```python
+# Get initial search results
+search_results = await client.search("machine learning", limit=3)
+
+# Expand with adjacent chunks based on CONTEXT_CHUNK_RADIUS setting
+expanded_results = await client.expand_context(search_results)
+
+# The expanded results contain chunks with combined content from adjacent chunks
+for chunk, score in expanded_results:
+    print(f"Expanded content: {chunk.content}")  # Now includes before/after chunks
+```
+
+This is automatically used by the QA system when `CONTEXT_CHUNK_RADIUS > 0` to provide better answers with more complete context.
+
 ## Question Answering
 
 Ask questions about your documents:

src/haiku/rag/qa/anthropic.py

@@ -79,12 +79,10 @@ try:
                                     else 3
                                 )
 
-                                search_results = await self._client.search(
+                                context = await self._search_and_expand(
                                     query, limit=limit
                                 )
 
-                                context = self._format_search_results(search_results)
-
                                 tool_results.append(
                                     {
                                         "type": "tool_result",

src/haiku/rag/qa/base.py

@@ -20,6 +20,12 @@ class QuestionAnswerAgentBase:
             "QABase is an abstract class. Please implement the answer method in a subclass."
         )
 
+    async def _search_and_expand(self, query: str, limit: int = 3) -> str:
+        """Search for documents and expand context, then format as JSON"""
+        search_results = await self._client.search(query, limit=limit)
+        expanded_results = await self._client.expand_context(search_results)
+        return self._format_search_results(expanded_results)
+
     def _format_search_results(self, search_results) -> str:
         """Format search results as JSON list of {content, score, document_uri}"""
         formatted_results = []

src/haiku/rag/qa/ollama.py

@@ -44,9 +44,7 @@ class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase):
                         query = args.get("query", question)
                         limit = int(args.get("limit", 3))
 
-                        search_results = await self._client.search(query, limit=limit)
+                        context = await self._search_and_expand(query, limit=limit)
-
-                        context = self._format_search_results(search_results)
                         messages.append(
                             {
                                 "role": "tool",

src/haiku/rag/qa/openai.py

@@ -77,11 +77,7 @@ try:
                             query = args.get("query", question)
                             limit = int(args.get("limit", 3))
 
-                            search_results = await self._client.search(
+                            context = await self._search_and_expand(query, limit=limit)
-                                query, limit=limit
-                            )
-
-                            context = self._format_search_results(search_results)
 
                             messages.append(
                                 ChatCompletionToolMessageParam(