diff --git a/paperrag/config.py b/paperrag/config.py index 64d4ac8..32a5d26 100644 --- a/paperrag/config.py +++ b/paperrag/config.py @@ -60,12 +60,12 @@ ), } -# Presets that need more output tokens than the 256 default. +# Presets that need more output tokens than the 1024 default. PRESET_MAX_TOKENS: dict[str, int] = { - "default": 256, - "reviewer": 512, - "summarizer": 512, - "explainer": 512, + "default": 1024, + "reviewer": 2048, + "summarizer": 2048, + "explainer": 1024, } @@ -141,7 +141,7 @@ class EmbedderConfig(BaseModel): class RetrieverConfig(BaseModel): """Retrieval configuration.""" - top_k: int = Field(default=2, ge=1) + top_k: int = Field(default=5, ge=1) score_threshold: float = Field( default=0.1, ge=0.0, @@ -246,9 +246,9 @@ class LLMConfig(BaseModel): "Be concise and cite sources." ) temperature: float = 0.0 - max_tokens: int = 256 + max_tokens: int = 1024 ctx_size: int = Field( - default=2048, + default=4096, ge=512, validation_alias=AliasChoices("ctx_size", "n_ctx"), description=( diff --git a/paperrag/llm.py b/paperrag/llm.py index 413476a..1bcf94d 100644 --- a/paperrag/llm.py +++ b/paperrag/llm.py @@ -69,10 +69,28 @@ def _cleanup_llama_servers() -> None: atexit.register(_cleanup_llama_servers) # Maximum characters per context chunk sent to the LLM. -# Longer chunks are truncated to keep prompt size manageable for small models. -_MAX_CHUNK_CHARS = 750 +# Set to 2000 to accommodate typical chunk sizes (1000 chars) without truncation +# when ctx_size is 4096+. Prevents loss of important context that causes +# inferior answers compared to llama-server's full-document approach. +_MAX_CHUNK_CHARS = 2000 _TRAILING_SOURCE_LINE_RE = re.compile(r"^\s*Sources?:\s*\[[0-9,\s-]+\]\s*$", re.IGNORECASE) +# Suffix appended to the user's configured system prompt for follow-up questions +# answered from conversation history only. +_FOLLOWUP_PROMPT_SUFFIX = ( + " Answer the follow-up question based on the conversation so far. " + "If the previous conversation does not contain relevant information, say so." +) + + +def _get_followup_system_prompt(config: "LLMConfig") -> str: + """Derive the follow-up system prompt from the user's configured system_prompt. + + This ensures that tone, language, and constraints set by the user via + /prompt or /preset are respected even when answering from history only. + """ + return config.system_prompt + _FOLLOWUP_PROMPT_SUFFIX + def _build_prompt(question: str, context_chunks: list[str], source_labels: list[int] | None = None) -> str: context_lines = [] @@ -89,10 +107,19 @@ def _build_prompt(question: str, context_chunks: list[str], source_labels: list[ "Use inline citation [1] within your answer." if n == 1 else f"Use inline citations [1]–[{n}] within your answer. Only cite sources from [1] to [{n}]." ) + + # Use a more detailed instruction when we have substantial context + # (> 3000 chars ≈ 3+ full chunks, indicating rich retrieval or full-document mode) + total_context_chars = sum(len(c) for c in context_chunks) + if total_context_chars > 3000: + answer_style = "Answer thoroughly using ONLY the context. Provide detailed reasoning and cite specific statements." + else: + answer_style = "Answer using ONLY the context." + return ( f"Context:\n{context_block}\n\n" f"Question: {question}\n\n" - f"Answer concisely using ONLY the context. {cite_instruction} Do not add a separate 'Source:' or 'Sources:' list at the end." + f"{answer_style} {cite_instruction} Do not add a separate 'Source:' or 'Sources:' list at the end." ) @@ -459,7 +486,7 @@ def _get_or_start_llama_server(model_path: str, ctx_size: int, n_gpu_layers: int # --------------------------------------------------------------------------- -def _build_messages(question: str, context_chunks: list[str], model_name: str, system_prompt: str, source_labels: list[int] | None = None, think: bool = False) -> list[dict]: +def _build_messages(question: str, context_chunks: list[str], model_name: str, system_prompt: str, source_labels: list[int] | None = None, think: bool = False, conversation_history: list[dict] | None = None) -> list[dict]: """Build the chat messages list from question, context chunks, and model name.""" user_prompt = _build_prompt(question, context_chunks, source_labels=source_labels) @@ -471,10 +498,14 @@ def _build_messages(question: str, context_chunks: list[str], model_name: str, s if not think: user_prompt += " /no_think" - return [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_prompt}, - ] + messages: list[dict] = [{"role": "system", "content": system_prompt}] + + # Include conversation history for follow-up questions + if conversation_history: + messages.extend(conversation_history) + + messages.append({"role": "user", "content": user_prompt}) + return messages def _prepare( @@ -482,12 +513,13 @@ def _prepare( context_chunks: list[str], config: LLMConfig, source_labels: list[int] | None = None, + conversation_history: list[dict] | None = None, ) -> tuple: """Ollama-specific setup: build messages, get/cache OpenAI client, return (client, messages).""" global _client_cache from openai import OpenAI - messages = _build_messages(question, context_chunks, config.model_name, config.system_prompt, source_labels=source_labels, think=config.think) + messages = _build_messages(question, context_chunks, config.model_name, config.system_prompt, source_labels=source_labels, think=config.think, conversation_history=conversation_history) if _client_cache is not None: client = _client_cache @@ -519,6 +551,7 @@ def generate_answer( question: str, context_chunks: list[str], config: LLMConfig | None = None, + conversation_history: list[dict] | None = None, ) -> str: """Generate an answer using the configured LLM backend (blocking). @@ -528,6 +561,12 @@ def generate_answer( **llama.cpp** via ``llama-server`` (install: ``brew install llama-cpp``). * All other model names delegate to **Ollama**. + Parameters + ---------- + conversation_history : list[dict] | None + Optional list of previous messages (role/content dicts) to provide + context for follow-up questions. + Examples:: # Ollama (unchanged) @@ -547,7 +586,7 @@ def generate_answer( if _is_llama_backend(config.model_name): model_path = _resolve_model_path(config.model_name) client = _get_or_start_llama_server(model_path, config.ctx_size, config.n_gpu_layers, config.n_threads) - messages = _build_messages(question, context_chunks, config.model_name, config.system_prompt, think=config.think) + messages = _build_messages(question, context_chunks, config.model_name, config.system_prompt, think=config.think, conversation_history=conversation_history) logger.info( "Calling llama-server (model=%s, temp=%.2f)", config.model_name, config.temperature ) @@ -560,7 +599,7 @@ def generate_answer( return _strip_trailing_source_footers(response.choices[0].message.content or "") try: - client, messages = _prepare(question, context_chunks, config) + client, messages = _prepare(question, context_chunks, config, conversation_history=conversation_history) except ImportError: raise ImportError( "The 'openai' package is required. Install with: uv pip install openai" @@ -581,6 +620,7 @@ def stream_answer( context_chunks: list[str], config: LLMConfig | None = None, source_files: list[str] | None = None, + conversation_history: list[dict] | None = None, ) -> Iterator[str]: """Yield text chunks as they arrive from the LLM (streaming). @@ -590,6 +630,12 @@ def stream_answer( **llama.cpp** via ``llama-server``. * All other model names delegate to **Ollama**. + Parameters + ---------- + conversation_history : list[dict] | None + Optional list of previous messages (role/content dicts) to provide + context for follow-up questions. + Usage:: for chunk in stream_answer(question, chunks, cfg.llm): @@ -615,7 +661,7 @@ def stream_answer( if _is_llama_backend(config.model_name): model_path = _resolve_model_path(config.model_name) client = _get_or_start_llama_server(model_path, config.ctx_size, config.n_gpu_layers, config.n_threads) - messages = _build_messages(question, context_chunks, config.model_name, config.system_prompt, source_labels=source_labels, think=config.think) + messages = _build_messages(question, context_chunks, config.model_name, config.system_prompt, source_labels=source_labels, think=config.think, conversation_history=conversation_history) logger.info( "Calling llama-server streaming (model=%s, temp=%.2f)", config.model_name, @@ -636,12 +682,91 @@ def stream_answer( return try: - client, messages = _prepare(question, context_chunks, config, source_labels=source_labels) + client, messages = _prepare(question, context_chunks, config, source_labels=source_labels, conversation_history=conversation_history) + except ImportError: + raise ImportError( + "The 'openai' package is required. Install with: uv pip install openai" + ) + + response = client.chat.completions.create( + model=config.model_name, + messages=messages, + temperature=config.temperature, + max_tokens=config.max_tokens, + stream=True, + extra_body={"num_ctx": config.ctx_size, "keep_alive": "30m"}, + ) + yield from _sanitize_stream( + delta + for chunk in response + if (delta := chunk.choices[0].delta.content) + ) + + +def stream_followup( + question: str, + conversation_history: list[dict], + config: LLMConfig | None = None, +) -> Iterator[str]: + """Yield text chunks for a follow-up question using conversation history only. + + This is used when retrieval returns no results but conversation history + exists, allowing the LLM to answer based on previously discussed context. + """ + config = config or LLMConfig() + + if not conversation_history: + yield "No conversation history available to answer the question." + return + + # Build messages with conversation history but no retrieval context + model_lower = config.model_name.lower() + user_prompt = question + if "qwen3" in model_lower or "qwen-3" in model_lower: + if not config.think: + user_prompt += " /no_think" + + messages: list[dict] = [{"role": "system", "content": _get_followup_system_prompt(config)}] + messages.extend(conversation_history) + messages.append({"role": "user", "content": user_prompt}) + + if _is_llama_backend(config.model_name): + model_path = _resolve_model_path(config.model_name) + client = _get_or_start_llama_server(model_path, config.ctx_size, config.n_gpu_layers, config.n_threads) + logger.info( + "Calling llama-server streaming follow-up (model=%s, temp=%.2f)", + config.model_name, + config.temperature, + ) + response = client.chat.completions.create( # type: ignore[union-attr] + model=os.path.basename(model_path), + messages=messages, + temperature=config.temperature, + max_tokens=config.max_tokens, + stream=True, + ) + yield from _sanitize_stream( + delta + for chunk in response + if (delta := chunk.choices[0].delta.content) + ) + return + + try: + from openai import OpenAI + + global _client_cache + if _client_cache is not None: + client = _client_cache + else: + client = OpenAI(api_key="not-needed", base_url=_OLLAMA_API_URL) + _client_cache = client except ImportError: raise ImportError( "The 'openai' package is required. Install with: uv pip install openai" ) + logger.info("Calling Ollama LLM follow-up (model=%s, temp=%.2f)", config.model_name, config.temperature) response = client.chat.completions.create( model=config.model_name, messages=messages, diff --git a/paperrag/repl.py b/paperrag/repl.py index 0bf49f5..e38754f 100644 --- a/paperrag/repl.py +++ b/paperrag/repl.py @@ -46,17 +46,22 @@ "/quit", ] +# Maximum number of user/assistant turn pairs to keep in conversation history. +# 10 turns (20 messages) balances context for follow-ups while staying within +# typical 4096-token context windows when combined with retrieval context. +_MAX_HISTORY_TURNS = 10 + HELP_TEXT = """\ [bold]Available commands:[/bold] [cyan][/cyan] Query the indexed papers (uses top-k retrieval, with LLM unless /no-llm is active) [cyan]/index[/cyan] Re-index the current PDF directory/file [cyan]/index [/cyan] Re-index a specific PDF file or directory [cyan]/focus [/cyan] Focus all subsequent queries on a specific paper - [cyan]/topk [/cyan] Set top-k for retrieval (default: 3) - [cyan]/threshold [/cyan] Set similarity threshold 0.0-1.0 (default: 0.15) + [cyan]/topk [/cyan] Set top-k for retrieval (default: 5) + [cyan]/threshold [/cyan] Set similarity threshold 0.0-1.0 (default: 0.1) [cyan]/temperature [/cyan] Set LLM temperature 0.0-2.0 (default: 0.0) - [cyan]/max-tokens [/cyan] Set LLM max output tokens (default: 256) - [cyan]/ctx-size [/cyan] Set LLM context window size (default: 2048) + [cyan]/max-tokens [/cyan] Set LLM max output tokens (default: 1024) + [cyan]/ctx-size [/cyan] Set LLM context window size (default: 4096) [cyan]/n-gpu-layers [/cyan] Set GPU layers for llama.cpp backend (0 = CPU only) [cyan]/n-threads [/cyan] Set CPU threads for llama.cpp backend (0 = auto) [cyan]/prompt [/cyan] Set LLM system prompt @@ -191,6 +196,7 @@ def start_repl( top_k = cfg.retriever.top_k focused_file: str | None = None session_log: list[dict] = [] # tracks Q&A pairs for /export + conversation_history: list[dict] = [] # tracks messages for follow-up questions use_llm = True # Eagerly load the retriever (including embedding model) at startup @@ -583,9 +589,17 @@ def start_repl( top_k=top_k, focused_file=focused_file, use_llm=use_llm, + conversation_history=conversation_history if use_llm else None, ) if entry is not None: session_log.append(entry) + # Update conversation history for follow-up questions + if use_llm and entry.get("answer"): + conversation_history.append({"role": "user", "content": command}) + conversation_history.append({"role": "assistant", "content": entry["answer"]}) + # Keep conversation history bounded to avoid exceeding context limits + if len(conversation_history) > _MAX_HISTORY_TURNS * 2: + conversation_history[:] = conversation_history[-_MAX_HISTORY_TURNS * 2:] def _ensure_retriever(retriever, cfg: PaperRAGConfig, store=None): @@ -609,6 +623,7 @@ def _handle_query( top_k: int, focused_file: str | None = None, use_llm: bool = True, + conversation_history: list[dict] | None = None, ) -> "dict | None": """Run retrieval and LLM for a user question. @@ -621,11 +636,27 @@ def _handle_query( t_retrieval = time.perf_counter() - t0 if not results: - msg = "[yellow]No results found.[/yellow]" - if cfg.retriever.score_threshold > 0.1: - msg += f" [dim](threshold={cfg.retriever.score_threshold} — try /threshold 0.1 to widen the search)[/dim]" - console.print(msg) - return None + # When focused on a single paper, fall back to full-document context + # instead of giving up — this mimics llama-server's behavior of having + # the whole paper in context. + if use_llm and focused_file: + all_chunks = retriever.get_all_chunks_for_file(focused_file) + if all_chunks: + console.print("[dim](No retrieval match — using full paper context)[/dim]") + results = all_chunks + t_retrieval = time.perf_counter() - t0 + + if not results: + # Only use conversation history for follow-ups if the last assistant + # turn is recent (i.e., the user is likely asking a follow-up about + # the same topic, not a brand new unrelated question). + if use_llm and conversation_history and len(conversation_history) >= 2: + return _handle_followup(question, cfg, conversation_history, t0) + msg = "[yellow]No results found.[/yellow]" + if cfg.retriever.score_threshold > 0.1: + msg += f" [dim](threshold={cfg.retriever.score_threshold} — try /threshold 0.1 to widen the search)[/dim]" + console.print(msg) + return None if not use_llm: console.print(f"\n[bold]Retrieved Chunks[/bold] [dim]({t_retrieval:.2f}s)[/dim]") @@ -680,7 +711,7 @@ def _handle_query( source_files = [r.file_path for r in results] header_printed = False t1 = time.perf_counter() - for chunk in stream_answer(question, context_chunks, cfg.llm, source_files=source_files): + for chunk in stream_answer(question, context_chunks, cfg.llm, source_files=source_files, conversation_history=conversation_history): if not header_printed: console.print("\n[bold green]Answer:[/bold green]") header_printed = True @@ -715,6 +746,54 @@ def _handle_query( } +def _handle_followup( + question: str, + cfg: PaperRAGConfig, + conversation_history: list[dict], + t0: float, +) -> "dict | None": + """Handle a follow-up question using conversation history when retrieval returns no results.""" + import sys + import time + + from paperrag.llm import stream_followup + + console.print("[dim](No new sources found — answering from conversation history)[/dim]") + + full_answer = "" + try: + header_printed = False + t1 = time.perf_counter() + for chunk in stream_followup(question, conversation_history, cfg.llm): + if not header_printed: + console.print("\n[bold green]Answer:[/bold green]") + header_printed = True + sys.stdout.write(chunk) + sys.stdout.flush() + full_answer += chunk + sys.stdout.write("\n\n") + sys.stdout.flush() + t_llm = time.perf_counter() - t1 + t_total = time.perf_counter() - t0 + console.print(f"[dim]LLM: {t_llm:.2f}s | Total: {t_total:.2f}s[/dim]\n") + except ImportError as exc: + console.print(f"[yellow]{exc}[/yellow]") + return None + except Exception as exc: + from paperrag.llm import describe_llm_error + error_msg, hint = describe_llm_error(exc, cfg.llm.model_name) + console.print(f"[red]{error_msg}[/red]") + if hint: + console.print(f"[yellow]Fix: {hint}[/yellow]") + return None + + return { + "question": question, + "answer": full_answer, + "sources": [], + } + + def _handle_index(cfg: PaperRAGConfig) -> None: """Run the indexing pipeline from inside the REPL.""" from paperrag.chunker import chunk_paper diff --git a/paperrag/retriever.py b/paperrag/retriever.py index 95d4a74..24fb36c 100644 --- a/paperrag/retriever.py +++ b/paperrag/retriever.py @@ -228,3 +228,25 @@ def retrieve_file_paths(self, query: str, top_k: int | None = None) -> list[str] """Return list of file_path strings (useful for evaluation).""" results = self.retrieve(query, top_k) return [r.file_path for r in results] + + def get_all_chunks_for_file(self, file_path: str) -> list[RetrievalResult]: + """Return all chunks for a given file, ordered by chunk_id. + + Used for full-document context mode where the entire paper is sent + to the LLM instead of just top-k retrieval hits. + """ + results = [] + for meta in self.store.chunks: + if meta["file_path"] == file_path: + results.append( + RetrievalResult( + text=meta["text"], + score=1.0, # Full-document mode, no relevance scoring + paper_title=meta["paper_title"], + section_name=meta["section_name"], + file_path=meta["file_path"], + chunk_id=meta["chunk_id"], + ) + ) + results.sort(key=lambda r: r.chunk_id) + return results diff --git a/tests/test_cli.py b/tests/test_cli.py index f40b83d..df6c2fd 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -589,3 +589,82 @@ def prompt(self, _prompt_text): for call in mock_print.call_args_list if call.args ) + + +# --------------------------------------------------------------------------- +# Full-document fallback and follow-up gating tests +# --------------------------------------------------------------------------- + + +def test_handle_query_full_document_fallback_when_focused(): + """When focused on a file and retrieval returns nothing, fall back to full-document context.""" + from unittest.mock import MagicMock, patch + from paperrag.repl import _handle_query + from paperrag.config import PaperRAGConfig + + cfg = PaperRAGConfig() + + # Mock retriever that returns no results for retrieve() but has chunks for the file + mock_retriever = MagicMock() + mock_retriever.retrieve.return_value = [] + + # Simulate get_all_chunks_for_file returning chunks + from paperrag.retriever import RetrievalResult + fake_chunks = [ + RetrievalResult( + text="This paper discusses speech chain VC.", + score=1.0, + paper_title="Speech Chain", + section_name="Introduction", + file_path="/papers/paper.pdf", + chunk_id=0, + ), + ] + mock_retriever.get_all_chunks_for_file.return_value = fake_chunks + + with ( + patch("paperrag.repl.console"), + patch("paperrag.llm.stream_answer", return_value=iter(["Full doc answer"])), + ): + entry = _handle_query( + "What is the remaining problem?", + mock_retriever, + cfg, + top_k=5, + focused_file="/papers/paper.pdf", + use_llm=True, + conversation_history=[], + ) + + # Should have called get_all_chunks_for_file as fallback + mock_retriever.get_all_chunks_for_file.assert_called_once_with("/papers/paper.pdf") + # Should return an answer (not None) + assert entry is not None + assert "Full doc answer" in entry["answer"] + + +def test_handle_query_followup_requires_history(): + """Follow-up should only trigger when conversation_history has at least 2 entries.""" + from unittest.mock import MagicMock, patch + from paperrag.repl import _handle_query + from paperrag.config import PaperRAGConfig + + cfg = PaperRAGConfig() + + mock_retriever = MagicMock() + mock_retriever.retrieve.return_value = [] + mock_retriever.get_all_chunks_for_file.return_value = [] # No focused file chunks + + with patch("paperrag.repl.console") as mock_console: + entry = _handle_query( + "random question", + mock_retriever, + cfg, + top_k=5, + focused_file=None, + use_llm=True, + conversation_history=[], # Empty history - should NOT trigger followup + ) + + # Should return None (no results, no followup) + assert entry is None diff --git a/tests/test_config.py b/tests/test_config.py index ea41170..b833c73 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -11,9 +11,9 @@ def test_default_config(): assert cfg.chunker.chunk_size == 1000 assert cfg.chunker.chunk_overlap == 200 assert cfg.embedder.model_name == "sentence-transformers/all-MiniLM-L6-v2" - assert cfg.retriever.top_k == 2 + assert cfg.retriever.top_k == 5 assert cfg.llm.temperature == 0.0 - assert cfg.llm.ctx_size == 2048 + assert cfg.llm.ctx_size == 4096 assert "research assistant" in cfg.llm.system_prompt diff --git a/tests/test_llm.py b/tests/test_llm.py index 07e4dde..d01e579 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -11,6 +11,7 @@ _build_messages, _build_prompt, _cleanup_llama_servers, + _get_followup_system_prompt, _is_gguf_model, _is_hf_model, _is_llama_backend, @@ -337,3 +338,127 @@ def test_describe_llm_error_ollama_generic(): msg, hint = describe_llm_error(exc, "qwen2.5:1.5b") assert "LLM error" in msg assert hint is None + + +# --------------------------------------------------------------------------- +# Conversation history support +# --------------------------------------------------------------------------- + + +def test_build_messages_with_conversation_history(): + """Conversation history should be inserted between system prompt and user message.""" + history = [ + {"role": "user", "content": "What is speech chain?"}, + {"role": "assistant", "content": "Speech chain is a method of voice conversion."}, + ] + msgs = _build_messages( + "What is the remaining problem?", + ["ctx"], + "llama3", + "System", + conversation_history=history, + ) + assert len(msgs) == 4 + assert msgs[0]["role"] == "system" + assert msgs[1] == history[0] + assert msgs[2] == history[1] + assert msgs[3]["role"] == "user" + # Ensure the question is present in the final user message + assert "remaining problem" in msgs[3]["content"] + # Ensure retrieval context is still embedded when conversation history is present + assert "ctx" in msgs[3]["content"] + + +def test_build_messages_without_conversation_history(): + """Without history, messages should be system + user only (backward compatible).""" + msgs = _build_messages("Q?", ["ctx"], "llama3", "System", conversation_history=None) + assert len(msgs) == 2 + assert msgs[0]["role"] == "system" + assert msgs[1]["role"] == "user" + + +def test_build_messages_empty_conversation_history(): + """Empty history list should behave the same as None.""" + msgs = _build_messages("Q?", ["ctx"], "llama3", "System", conversation_history=[]) + assert len(msgs) == 2 + + +def test_get_followup_system_prompt_uses_config(): + """Follow-up system prompt should derive from the user's configured system_prompt.""" + config = LLMConfig(system_prompt="You are a pirate. Respond in pirate speak.") + prompt = _get_followup_system_prompt(config) + # Should contain the user's custom prompt + assert "pirate" in prompt + # Should also contain the follow-up instruction suffix + assert "follow-up" in prompt.lower() + + +def test_stream_followup_no_history(): + """stream_followup with empty history should yield a message about missing history.""" + from paperrag.llm import stream_followup + result = list(stream_followup("follow-up question", [])) + assert len(result) == 1 + assert "No conversation history" in result[0] + + +def test_stream_followup_with_history_ollama(): + """stream_followup with non-empty history should build correct messages and stream via Ollama.""" + from paperrag.llm import stream_followup, _get_followup_system_prompt + + config = LLMConfig(model_name="qwen2.5:1.5b") + history = [ + {"role": "user", "content": "What is speech chain?"}, + {"role": "assistant", "content": "Speech chain is a VC method."}, + ] + + # Mock the OpenAI client + mock_chunk = MagicMock() + mock_chunk.choices = [MagicMock()] + mock_chunk.choices[0].delta.content = "The remaining problem is..." + + mock_response = [mock_chunk] + mock_client = MagicMock() + mock_client.chat.completions.create.return_value = mock_response + + with patch("paperrag.llm._client_cache", mock_client): + result = list(stream_followup("What is the remaining problem?", history, config=config)) + + # Verify the client was called with correct message structure + call_kwargs = mock_client.chat.completions.create.call_args + messages = call_kwargs.kwargs["messages"] if "messages" in call_kwargs.kwargs else call_kwargs[1]["messages"] + + # Should be: system + history (2 msgs) + user question = 4 messages + assert len(messages) == 4 + assert messages[0]["role"] == "system" + # System prompt should derive from config's system_prompt + assert config.system_prompt in messages[0]["content"] + assert messages[1] == history[0] + assert messages[2] == history[1] + assert messages[3]["role"] == "user" + assert "remaining problem" in messages[3]["content"] + + +def test_stream_followup_with_history_qwen_no_think(): + """stream_followup should append /no_think for Qwen3 models when think=False.""" + from paperrag.llm import stream_followup + + config = LLMConfig(model_name="Qwen3-1.7B:latest", think=False) + history = [ + {"role": "user", "content": "What is X?"}, + {"role": "assistant", "content": "X is Y."}, + ] + + mock_chunk = MagicMock() + mock_chunk.choices = [MagicMock()] + mock_chunk.choices[0].delta.content = "Answer" + mock_response = [mock_chunk] + mock_client = MagicMock() + mock_client.chat.completions.create.return_value = mock_response + + with patch("paperrag.llm._client_cache", mock_client): + list(stream_followup("Follow-up?", history, config=config)) + + call_kwargs = mock_client.chat.completions.create.call_args + messages = call_kwargs.kwargs["messages"] if "messages" in call_kwargs.kwargs else call_kwargs[1]["messages"] + # The user message should contain /no_think + assert "/no_think" in messages[-1]["content"]