From ed8e72847d4fb0e73f6ff98679452a80375d72f4 Mon Sep 17 00:00:00 2001 From: markgewhite Date: Sat, 11 Apr 2026 23:02:51 +0100 Subject: [PATCH 1/2] Added condenser module for rewriting follow-up questions with tests Co-Authored-By: Claude Opus 4.6 (1M context) --- src/generation/condenser.py | 50 ++++++++++++++++++++++ tests/test_condenser.py | 82 +++++++++++++++++++++++++++++++++++++ 2 files changed, 132 insertions(+) create mode 100644 src/generation/condenser.py create mode 100644 tests/test_condenser.py diff --git a/src/generation/condenser.py b/src/generation/condenser.py new file mode 100644 index 0000000..b37633f --- /dev/null +++ b/src/generation/condenser.py @@ -0,0 +1,50 @@ +"""Condense follow-up questions into standalone questions using chat history. + +Not via LangChain: uses direct ollama.chat() for the condensation LLM call, +keeping full control over the prompt format and avoiding framework overhead. +""" + +import ollama + +CONDENSE_PROMPT = ( + "Given the following conversation history and a follow-up question, " + "rewrite the follow-up question as a standalone question that captures " + "the full context. Return ONLY the rewritten question, nothing else." +) + + +class Condenser: + """Rewrites follow-up questions into standalone questions via LLM.""" + + def __init__(self, model: str) -> None: + self._model = model + + def condense( + self, + question: str, + chat_history: list[dict[str, str]], + ) -> str: + """Rewrite a follow-up question into a standalone question. + + If chat_history is empty, returns the question unchanged (no LLM call). + """ + if not chat_history: + return question + + history_text = "\n".join( + f"{msg['role'].title()}: {msg['content']}" for msg in chat_history + ) + + messages = [ + {"role": "system", "content": CONDENSE_PROMPT}, + { + "role": "user", + "content": ( + f"Conversation history:\n{history_text}\n\n" + f"Follow-up question: {question}" + ), + }, + ] + + response = ollama.chat(model=self._model, messages=messages) + return response["message"]["content"].strip() diff --git a/tests/test_condenser.py b/tests/test_condenser.py new file mode 100644 index 0000000..154f747 --- /dev/null +++ b/tests/test_condenser.py @@ -0,0 +1,82 @@ +"""Tests for the condenser module.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from src.generation.condenser import Condenser + + +class TestCondenser: + """Tests for the Condenser class.""" + + @patch("src.generation.condenser.ollama") + def test_condense_skips_when_no_history(self, mock_ollama): + """First question (no history) returns the original question unchanged.""" + condenser = Condenser(model="llama3.1:8b") + result = condenser.condense("What are filing deadlines?", []) + + assert result == "What are filing deadlines?" + mock_ollama.chat.assert_not_called() + + @patch("src.generation.condenser.ollama") + def test_condense_rewrites_with_history(self, mock_ollama): + """Follow-up question is rewritten using chat history.""" + mock_ollama.chat.return_value = { + "message": {"content": "What are the penalties for missing tax filing deadlines?"} + } + + condenser = Condenser(model="llama3.1:8b") + history = [ + {"role": "user", "content": "What are filing deadlines?"}, + {"role": "assistant", "content": "Filing deadlines are in April."}, + ] + result = condenser.condense("And what about penalties?", history) + + assert result == "What are the penalties for missing tax filing deadlines?" + mock_ollama.chat.assert_called_once() + + @patch("src.generation.condenser.ollama") + def test_condense_passes_correct_model(self, mock_ollama): + """Condenser uses the configured model for the LLM call.""" + mock_ollama.chat.return_value = { + "message": {"content": "standalone question"} + } + + condenser = Condenser(model="custom-model") + condenser.condense("follow up", [{"role": "user", "content": "first"}]) + + call_kwargs = mock_ollama.chat.call_args + assert call_kwargs[1]["model"] == "custom-model" + + @patch("src.generation.condenser.ollama") + def test_condense_includes_history_in_prompt(self, mock_ollama): + """The condensation prompt includes chat history.""" + mock_ollama.chat.return_value = { + "message": {"content": "standalone question"} + } + + condenser = Condenser(model="llama3.1:8b") + history = [ + {"role": "user", "content": "What is Python?"}, + {"role": "assistant", "content": "A programming language."}, + ] + condenser.condense("Tell me more", history) + + messages = mock_ollama.chat.call_args[1]["messages"] + prompt_text = " ".join(m["content"] for m in messages) + assert "What is Python?" in prompt_text + assert "A programming language." in prompt_text + assert "Tell me more" in prompt_text + + @patch("src.generation.condenser.ollama") + def test_condense_strips_whitespace(self, mock_ollama): + """Condensed output has leading/trailing whitespace stripped.""" + mock_ollama.chat.return_value = { + "message": {"content": " standalone question \n"} + } + + condenser = Condenser(model="llama3.1:8b") + result = condenser.condense("follow up", [{"role": "user", "content": "first"}]) + + assert result == "standalone question" From 29e7dcc4f11a1cf523ce647d427090fce51b8be6 Mon Sep 17 00:00:00 2001 From: markgewhite Date: Sat, 11 Apr 2026 23:03:27 +0100 Subject: [PATCH 2/2] Integrated condenser into Chainlit pipeline with per-session chat history Co-Authored-By: Claude Opus 4.6 (1M context) --- app.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/app.py b/app.py index 0a322ff..4cd050a 100644 --- a/app.py +++ b/app.py @@ -6,6 +6,7 @@ from src.config import ConfigError, load_config from src.generation.answerer import answer, build_source_elements +from src.generation.condenser import Condenser from src.health_check import check_models, check_ollama from src.ingestion.chunker import chunk_documents from src.ingestion.loader import load_folder @@ -88,6 +89,11 @@ async def on_chat_start(): reranker = Reranker(model_name=config.retrieval.reranker_model) cl.user_session.set("reranker", reranker) + # Create condenser for follow-up questions and init chat history + condenser = Condenser(model=config.models.llm) + cl.user_session.set("condenser", condenser) + cl.user_session.set("chat_history", []) + if doc_count > 0: await cl.Message( content=f"{doc_count} chunks indexed. Ask me anything!" @@ -114,18 +120,23 @@ async def on_message(message: cl.Message): ).send() return - candidates = retriever.retrieve(message.content) + # Condense follow-up questions into standalone queries + condenser = cl.user_session.get("condenser") + chat_history = cl.user_session.get("chat_history") + query = condenser.condense(message.content, chat_history) + + candidates = retriever.retrieve(query) # Rerank candidates with cross-encoder reranker = cl.user_session.get("reranker") results = reranker.rerank( - message.content, candidates, top_k=config.retrieval.rerank_top_k + query, candidates, top_k=config.retrieval.rerank_top_k ) # Stream the answer token by token msg = cl.Message(content="") async for token in answer( - message.content, + query, results, model=config.models.llm, ): @@ -140,3 +151,7 @@ async def on_message(message: cl.Message): display=el_data["display"], ) await element.send(for_id=msg.id) + + # Update chat history (kept separate from answer context window) + chat_history.append({"role": "user", "content": message.content}) + chat_history.append({"role": "assistant", "content": msg.content})