diff --git a/hyperextract/cli/cli.py b/hyperextract/cli/cli.py index b62c29e..c59905e 100644 --- a/hyperextract/cli/cli.py +++ b/hyperextract/cli/cli.py @@ -647,7 +647,7 @@ def search( ) -def chat_loop(ka, ka_path: str): +def chat_loop(ka, ka_path: str, top_k: int = 3): """Interactive chat loop.""" console.print( "\n[bold green]Entering interactive mode. Type 'exit' or 'quit' to stop.[/bold green]\n" @@ -669,7 +669,7 @@ def chat_loop(ka, ka_path: str): break if not query.strip(): continue - response = ka.chat(query) + response = ka.chat(query, top_k=top_k) console.print() console.print(response.content) console.print() @@ -740,7 +740,7 @@ def talk( raise typer.Exit(1) if interactive: - chat_loop(ka, ka_path) + chat_loop(ka, ka_path, top_k=top_k) else: with console.status("[bold blue]Thinking..."): try: diff --git a/tests/cli/test_talk.py b/tests/cli/test_talk.py new file mode 100644 index 0000000..eb19cc6 --- /dev/null +++ b/tests/cli/test_talk.py @@ -0,0 +1,20 @@ +"""Tests for the interactive `he talk` chat loop.""" + +import hyperextract.cli.cli as climod + + +def test_chat_loop_passes_top_k(monkeypatch): + """Interactive chat must forward top_k to ka.chat, not use the default.""" + recorded = {} + + class _StubKA: + def chat(self, query, top_k=3): + recorded["top_k"] = top_k + return type("_Resp", (), {"content": "ok"})() + + queries = iter(["hello", "exit"]) + monkeypatch.setattr(climod.console, "input", lambda *a, **k: next(queries)) + + climod.chat_loop(_StubKA(), "some/ka", top_k=10) + + assert recorded["top_k"] == 10