diff --git a/docs/configuration.md b/docs/configuration.md index c1507e49..1302ade6 100644 --- a/docs/configuration.md +++ b/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 ``` diff --git a/docs/python.md b/docs/python.md index a498c22a..fb46f0ad 100644 --- a/docs/python.md +++ b/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: diff --git a/src/haiku/rag/qa/anthropic.py b/src/haiku/rag/qa/anthropic.py index 6c516aff..c3138e1f 100644 --- a/src/haiku/rag/qa/anthropic.py +++ b/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", diff --git a/src/haiku/rag/qa/base.py b/src/haiku/rag/qa/base.py index eea670df..92a53060 100644 --- a/src/haiku/rag/qa/base.py +++ b/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 = [] diff --git a/src/haiku/rag/qa/ollama.py b/src/haiku/rag/qa/ollama.py index c993b510..6b363542 100644 --- a/src/haiku/rag/qa/ollama.py +++ b/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 = self._format_search_results(search_results) + context = await self._search_and_expand(query, limit=limit) messages.append( { "role": "tool", diff --git a/src/haiku/rag/qa/openai.py b/src/haiku/rag/qa/openai.py index 27396d7a..0ab45793 100644 --- a/src/haiku/rag/qa/openai.py +++ b/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( - query, limit=limit - ) - - context = self._format_search_results(search_results) + context = await self._search_and_expand(query, limit=limit) messages.append( ChatCompletionToolMessageParam(