Enhancement: AI LLM chunk size and context window config (#12891)

src/paperless_ai/client.py

@@ -59,6 +59,7 @@ class AIClient:
             return Ollama(
                 model=self.settings.llm_model or "llama3.1",
                 base_url=endpoint,
+                context_window=self.settings.llm_context_size,
                 request_timeout=120,
                 system_prompt=LLM_SYSTEM_PROMPT,
                 client=Client(

src/paperless_ai/embedding.py

@@ -74,6 +74,7 @@ def get_embedding_model() -> "BaseEmbedding":
             embedding = OllamaEmbedding(
                 model_name=config.llm_embedding_model or "embeddinggemma",
                 base_url=endpoint,
+                ollama_additional_kwargs={"num_ctx": config.llm_context_size},
             )
             embedding._client = Client(
                 host=endpoint,

src/paperless_ai/indexing.py

@@ -12,6 +12,7 @@ from documents.models import Document
 from documents.models import PaperlessTask
 from documents.utils import IterWrapper
 from documents.utils import identity
+from paperless.config import AIConfig
 from paperless_ai.embedding import build_llm_index_text
 from paperless_ai.embedding import get_embedding_dim
 from paperless_ai.embedding import get_embedding_model
@@ -23,9 +24,7 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger("paperless_ai.indexing")
 
-RAG_CONTEXT_WINDOW = 8192
 RAG_NUM_OUTPUT = 512
-RAG_CHUNK_SIZE = 1024
 RAG_CHUNK_OVERLAP = 200
 
 
@@ -95,7 +94,11 @@ def get_or_create_storage_context(*, rebuild=False):
     )
 
 
-def build_document_node(document: Document) -> list["BaseNode"]:
+def build_document_node(
+    document: Document,
+    *,
+    chunk_size: int | None = None,
+) -> list["BaseNode"]:
     """
     Given a Document, returns parsed Nodes ready for indexing.
     """
@@ -126,9 +129,10 @@ def build_document_node(document: Document) -> list["BaseNode"]:
         metadata=metadata,
         excluded_embed_metadata_keys=list(metadata.keys()),
     )
+    chunk_size = chunk_size or get_rag_chunk_size()
     parser = SimpleNodeParser(
-        chunk_size=RAG_CHUNK_SIZE,
-        chunk_overlap=get_rag_chunk_overlap(),
+        chunk_size=chunk_size,
+        chunk_overlap=get_rag_chunk_overlap(chunk_size),
     )
     return parser.get_nodes_from_documents([doc])
 
@@ -186,18 +190,36 @@ def vector_store_file_exists():
     return Path(settings.LLM_INDEX_DIR / "default__vector_store.json").exists()
 
 
-def get_rag_chunk_overlap() -> int:
-    return min(RAG_CHUNK_OVERLAP, RAG_CHUNK_SIZE - 1)
+def get_rag_chunk_size() -> int:
+    return AIConfig().llm_embedding_chunk_size
 
 
-def get_rag_prompt_helper():
+def get_rag_context_size() -> int:
+    return AIConfig().llm_context_size
+
+
+def get_rag_chunk_overlap(chunk_size: int | None = None) -> int:
+    chunk_size = chunk_size or get_rag_chunk_size()
+    return min(RAG_CHUNK_OVERLAP, chunk_size - 1)
+
+
+def get_rag_prompt_helper(
+    *,
+    chunk_size: int | None = None,
+    context_size: int | None = None,
+):
     from llama_index.core.indices.prompt_helper import PromptHelper
 
+    if chunk_size is None or context_size is None:
+        config = AIConfig()
+        chunk_size = chunk_size or config.llm_embedding_chunk_size
+        context_size = context_size or config.llm_context_size
+
     return PromptHelper(
-        context_window=RAG_CONTEXT_WINDOW,
+        context_window=context_size,
         num_output=RAG_NUM_OUTPUT,
         chunk_overlap_ratio=0.1,
-        chunk_size_limit=RAG_CHUNK_SIZE,
+        chunk_size_limit=chunk_size,
     )
 
 
@@ -219,6 +241,9 @@ def update_llm_index(
         logger.warning(msg)
         return msg
 
+    config = AIConfig()
+    chunk_size = config.llm_embedding_chunk_size
+
     if rebuild or not vector_store_file_exists():
         # remove meta.json to force re-detection of embedding dim
         (settings.LLM_INDEX_DIR / "meta.json").unlink(missing_ok=True)
@@ -230,7 +255,7 @@ def update_llm_index(
         llama_settings.Settings.embed_model = embed_model
         storage_context = get_or_create_storage_context(rebuild=True)
         for document in iter_wrapper(documents):
-            document_nodes = build_document_node(document)
+            document_nodes = build_document_node(document, chunk_size=chunk_size)
             nodes.extend(document_nodes)
 
         index = VectorStoreIndex(
@@ -262,10 +287,10 @@ def update_llm_index(
 
                 # Again, delete from docstore, FAISS IndexFlatL2 are append-only
                 index.docstore.delete_document(node.node_id)
-                nodes.extend(build_document_node(document))
+                nodes.extend(build_document_node(document, chunk_size=chunk_size))
             else:
                 # New document, add it
-                nodes.extend(build_document_node(document))
+                nodes.extend(build_document_node(document, chunk_size=chunk_size))
 
         if nodes:
             msg = "LLM index updated successfully."
@@ -287,7 +312,7 @@ def llm_index_add_or_update_document(document: Document):
     Adds or updates a document in the LLM index.
     If the document already exists, it will be replaced.
     """
-    new_nodes = build_document_node(document)
+    new_nodes = build_document_node(document, chunk_size=get_rag_chunk_size())
 
     index = load_or_build_index(nodes=new_nodes)
 
@@ -309,15 +334,27 @@ def llm_index_remove_document(document: Document):
     index.storage_context.persist(persist_dir=settings.LLM_INDEX_DIR)
 
 
-def truncate_content(content: str) -> str:
+def truncate_content(
+    content: str,
+    *,
+    chunk_size: int | None = None,
+    context_size: int | None = None,
+) -> str:
     from llama_index.core.prompts import PromptTemplate
     from llama_index.core.text_splitter import TokenTextSplitter
 
-    prompt_helper = get_rag_prompt_helper()
+    if chunk_size is None or context_size is None:
+        config = AIConfig()
+        chunk_size = chunk_size or config.llm_embedding_chunk_size
+        context_size = context_size or config.llm_context_size
+    prompt_helper = get_rag_prompt_helper(
+        chunk_size=chunk_size,
+        context_size=context_size,
+    )
     splitter = TokenTextSplitter(
         separator=" ",
-        chunk_size=RAG_CHUNK_SIZE,
-        chunk_overlap=get_rag_chunk_overlap(),
+        chunk_size=chunk_size,
+        chunk_overlap=get_rag_chunk_overlap(chunk_size),
     )
     content_chunks = splitter.split_text(content)
     truncated_chunks = prompt_helper.truncate(
@@ -376,8 +413,11 @@ def query_similar_documents(
         doc_ids=doc_node_ids,
     )
 
+    config = AIConfig()
     query_text = truncate_content(
         (document.title or "") + "\n" + (document.content or ""),
+        chunk_size=config.llm_embedding_chunk_size,
+        context_size=config.llm_context_size,
     )
     results = retriever.retrieve(query_text)
 

src/paperless_ai/tests/test_ai_indexing.py

@@ -11,6 +11,7 @@ from llama_index.core.base.embeddings.base import BaseEmbedding
 from documents.models import Document
 from documents.models import PaperlessTask
 from documents.tests.factories import PaperlessTaskFactory
+from paperless.models import ApplicationConfiguration
 from paperless_ai import indexing
 
 
@@ -81,20 +82,32 @@ def test_build_document_node_excludes_metadata_from_embedding(real_document) ->
 
 @pytest.mark.django_db
 def test_build_document_node_uses_rag_chunk_settings(real_document) -> None:
+    app_config, _ = ApplicationConfiguration.objects.get_or_create()
+    app_config.llm_embedding_chunk_size = 512
+    app_config.save()
+
     with patch("llama_index.core.node_parser.SimpleNodeParser") as mock_parser:
         mock_parser.return_value.get_nodes_from_documents.return_value = []
 
         indexing.build_document_node(real_document)
 
-        mock_parser.assert_called_once_with(chunk_size=1024, chunk_overlap=200)
+        mock_parser.assert_called_once_with(chunk_size=512, chunk_overlap=200)
 
 
 def test_get_rag_chunk_overlap_clamps_to_chunk_size() -> None:
-    with (
-        patch("paperless_ai.indexing.RAG_CHUNK_SIZE", 64),
-        patch("paperless_ai.indexing.RAG_CHUNK_OVERLAP", 128),
-    ):
-        assert indexing.get_rag_chunk_overlap() == 63
+    with patch("paperless_ai.indexing.RAG_CHUNK_OVERLAP", 128):
+        assert indexing.get_rag_chunk_overlap(64) == 63
+
+
+@pytest.mark.django_db
+def test_get_rag_prompt_helper_uses_context_setting() -> None:
+    app_config, _ = ApplicationConfiguration.objects.get_or_create()
+    app_config.llm_context_size = 4096
+    app_config.save()
+
+    prompt_helper = indexing.get_rag_prompt_helper()
+
+    assert prompt_helper.context_window == 4096
 
 
 @pytest.mark.django_db
@@ -103,13 +116,22 @@ def test_update_llm_index(
     real_document,
     mock_embed_model,
 ) -> None:
-    with patch("documents.models.Document.objects.all") as mock_all:
+    mock_config = MagicMock()
+    mock_config.llm_embedding_chunk_size = 512
+    with (
+        patch("documents.models.Document.objects.all") as mock_all,
+        patch("paperless_ai.indexing.AIConfig", return_value=mock_config) as ai_config,
+        patch("paperless_ai.indexing.build_document_node") as build_document_node,
+    ):
         mock_queryset = MagicMock()
         mock_queryset.exists.return_value = True
         mock_queryset.__iter__.return_value = iter([real_document])
         mock_all.return_value = mock_queryset
+        build_document_node.return_value = []
         indexing.update_llm_index(rebuild=True)
 
+        ai_config.assert_called_once()
+        build_document_node.assert_called_once_with(real_document, chunk_size=512)
         assert any(temp_llm_index_dir.glob("*.json"))
 
 

src/paperless_ai/tests/test_chat.py

@@ -143,6 +143,7 @@ def test_document_filtered_retriever_handles_empty_faiss_index() -> None:
     mock_index.vector_store.query.assert_not_called()
 
 
+@pytest.mark.django_db
 def test_stream_chat_with_one_document_retrieval(
     mock_document,
     patch_embed_nodes,
@@ -186,6 +187,7 @@ def test_stream_chat_with_one_document_retrieval(
         )
 
 
+@pytest.mark.django_db
 def test_stream_chat_with_multiple_documents_retrieval(patch_embed_nodes) -> None:
     with (
         patch("paperless_ai.chat.AIClient") as mock_client_cls,

src/paperless_ai/tests/test_client.py

@@ -15,6 +15,7 @@ def mock_ai_config():
     with patch("paperless_ai.client.AIConfig") as MockAIConfig:
         mock_config = MagicMock()
         mock_config.llm_allow_internal_endpoints = True
+        mock_config.llm_context_size = 8192
         MockAIConfig.return_value = mock_config
         yield mock_config
 
@@ -41,6 +42,7 @@ def test_get_llm_ollama(mock_ai_config, mock_ollama_llm):
     mock_ollama_llm.assert_called_once_with(
         model="test_model",
         base_url="http://test-url",
+        context_window=8192,
         request_timeout=120,
         system_prompt=LLM_SYSTEM_PROMPT,
         client=ANY,

src/paperless_ai/tests/test_embedding.py

@@ -19,6 +19,7 @@ def mock_ai_config():
     with patch("paperless_ai.embedding.AIConfig") as MockAIConfig:
         MockAIConfig.return_value.llm_embedding_endpoint = None
         MockAIConfig.return_value.llm_allow_internal_endpoints = True
+        MockAIConfig.return_value.llm_context_size = 8192
         yield MockAIConfig
 
 
@@ -140,6 +141,7 @@ def test_get_embedding_model_ollama(mock_ai_config):
         MockOllamaEmbedding.assert_called_once_with(
             model_name="embeddinggemma",
             base_url="http://test-url",
+            ollama_additional_kwargs={"num_ctx": 8192},
         )
         assert model == MockOllamaEmbedding.return_value
 
@@ -157,6 +159,7 @@ def test_get_embedding_model_ollama_prefers_embedding_endpoint(mock_ai_config):
         MockOllamaEmbedding.assert_called_once_with(
             model_name="embeddinggemma",
             base_url="http://embedding-url",
+            ollama_additional_kwargs={"num_ctx": 8192},
         )
         assert model == MockOllamaEmbedding.return_value