Session memory: condense-then-retrieve

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})

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()

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"