Skip to content

Commit 620f714

Browse files
committed
Address code review feedback for slash commands
- Rename _add_single_file → add_single_file (public API, consistent with print_status/print_list/run_lint already being public) - Add progress counter [i/total] for /add when processing directories - Move `import asyncio` to module-level imports in chat.py - Add 12 tests covering all new slash commands in _handle_slash
1 parent 542ad7a commit 620f714

4 files changed

Lines changed: 203 additions & 12 deletions

File tree

openkb/agent/chat.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
"""
88
from __future__ import annotations
99

10+
import asyncio
1011
import os
1112
import re
1213
import sys
@@ -277,8 +278,7 @@ def _save_transcript(kb_dir: Path, session: ChatSession, name: str | None) -> Pa
277278

278279
async def _run_add(arg: str, kb_dir: Path, style: Style) -> None:
279280
"""Add a document or directory to the knowledge base from the chat REPL."""
280-
import asyncio
281-
from openkb.cli import _add_single_file, SUPPORTED_EXTENSIONS
281+
from openkb.cli import add_single_file, SUPPORTED_EXTENSIONS
282282

283283
target = Path(arg).expanduser()
284284
if not target.is_absolute():
@@ -297,13 +297,16 @@ async def _run_add(arg: str, kb_dir: Path, style: Style) -> None:
297297
if not files:
298298
_fmt(style, ("class:error", f"No supported files found in {arg}.\n"))
299299
return
300-
for f in files:
301-
await asyncio.to_thread(_add_single_file, f, kb_dir)
300+
total = len(files)
301+
_fmt(style, ("class:slash.help", f"Found {total} supported file(s) in {arg}.\n"))
302+
for i, f in enumerate(files, 1):
303+
_fmt(style, ("class:slash.help", f"\n[{i}/{total}] "))
304+
await asyncio.to_thread(add_single_file, f, kb_dir)
302305
else:
303306
if target.suffix.lower() not in SUPPORTED_EXTENSIONS:
304307
_fmt(style, ("class:error", f"Unsupported file type: {target.suffix}\n"))
305308
return
306-
await asyncio.to_thread(_add_single_file, target, kb_dir)
309+
await asyncio.to_thread(add_single_file, target, kb_dir)
307310

308311

309312
async def _handle_slash(

openkb/cli.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ def _find_kb_dir(override: Path | None = None) -> Path | None:
128128
return None
129129

130130

131-
def _add_single_file(file_path: Path, kb_dir: Path) -> None:
131+
def add_single_file(file_path: Path, kb_dir: Path) -> None:
132132
"""Convert, index, and compile a single document into the knowledge base.
133133
134134
Steps:
@@ -346,15 +346,15 @@ def add(ctx, path):
346346
click.echo(f"Found {total} supported file(s) in {path}.")
347347
for i, f in enumerate(files, 1):
348348
click.echo(f"\n[{i}/{total}] ", nl=False)
349-
_add_single_file(f, kb_dir)
349+
add_single_file(f, kb_dir)
350350
else:
351351
if target.suffix.lower() not in SUPPORTED_EXTENSIONS:
352352
click.echo(
353353
f"Unsupported file type: {target.suffix}. "
354354
f"Supported: {', '.join(sorted(SUPPORTED_EXTENSIONS))}"
355355
)
356356
return
357-
_add_single_file(target, kb_dir)
357+
add_single_file(target, kb_dir)
358358

359359

360360
@cli.command()
@@ -519,7 +519,7 @@ def on_new_files(paths):
519519
f"Supported: {', '.join(sorted(SUPPORTED_EXTENSIONS))}"
520520
)
521521
continue
522-
_add_single_file(fp, kb_dir)
522+
add_single_file(fp, kb_dir)
523523

524524
click.echo(f"Watching {raw_dir} for new documents. Press Ctrl+C to stop.")
525525
watch_directory(raw_dir, on_new_files)

tests/test_add_command.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,13 +63,13 @@ def test_add_missing_init(self, tmp_path):
6363
result = runner.invoke(cli, ["add", "somefile.pdf"])
6464
assert "No knowledge base found" in result.output
6565

66-
def test_add_single_file_calls_helper(self, tmp_path):
66+
def testadd_single_file_calls_helper(self, tmp_path):
6767
kb_dir = self._setup_kb(tmp_path)
6868
doc = tmp_path / "test.md"
6969
doc.write_text("# Hello")
7070

7171
runner = CliRunner()
72-
with patch("openkb.cli._add_single_file") as mock_add, \
72+
with patch("openkb.cli.add_single_file") as mock_add, \
7373
patch("openkb.cli._find_kb_dir", return_value=kb_dir):
7474
result = runner.invoke(cli, ["add", str(doc)])
7575
mock_add.assert_called_once_with(doc, kb_dir)
@@ -83,7 +83,7 @@ def test_add_directory_calls_helper_for_each_file(self, tmp_path):
8383
(docs_dir / "ignore.xyz").write_text("skip me")
8484

8585
runner = CliRunner()
86-
with patch("openkb.cli._add_single_file") as mock_add, \
86+
with patch("openkb.cli.add_single_file") as mock_add, \
8787
patch("openkb.cli._find_kb_dir", return_value=kb_dir):
8888
result = runner.invoke(cli, ["add", str(docs_dir)])
8989
# Should be called for .md and .txt but not .xyz

tests/test_chat_slash_commands.py

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
"""Tests for slash commands in the chat REPL."""
2+
from __future__ import annotations
3+
4+
import json
5+
from pathlib import Path
6+
from unittest.mock import AsyncMock, patch
7+
8+
import pytest
9+
10+
from prompt_toolkit.styles import Style
11+
12+
from openkb.agent.chat import _handle_slash, _run_add
13+
from openkb.agent.chat_session import ChatSession
14+
15+
16+
def _setup_kb(tmp_path: Path) -> Path:
17+
"""Create a minimal KB structure and return kb_dir."""
18+
kb_dir = tmp_path
19+
(kb_dir / "raw").mkdir()
20+
(kb_dir / "wiki" / "sources" / "images").mkdir(parents=True)
21+
(kb_dir / "wiki" / "summaries").mkdir(parents=True)
22+
(kb_dir / "wiki" / "concepts").mkdir(parents=True)
23+
(kb_dir / "wiki" / "reports").mkdir(parents=True)
24+
openkb_dir = kb_dir / ".openkb"
25+
openkb_dir.mkdir()
26+
(openkb_dir / "config.yaml").write_text("model: gpt-4o-mini\n")
27+
(openkb_dir / "hashes.json").write_text(json.dumps({}))
28+
return kb_dir
29+
30+
31+
def _make_session(kb_dir: Path) -> ChatSession:
32+
return ChatSession.new(kb_dir, "gpt-4o-mini", "en")
33+
34+
35+
_STYLE = Style.from_dict({})
36+
37+
38+
def _collect_fmt():
39+
"""Return (patch, collected) where collected is a list of printed strings."""
40+
collected: list[str] = []
41+
42+
def _fake_fmt(_style, *fragments):
43+
for _cls, text in fragments:
44+
collected.append(text)
45+
46+
return patch("openkb.agent.chat._fmt", _fake_fmt), collected
47+
48+
49+
# --- /status and /list use click.echo, captured by capsys ---
50+
51+
52+
@pytest.mark.asyncio
53+
async def test_slash_status(tmp_path, capsys):
54+
kb_dir = _setup_kb(tmp_path)
55+
session = _make_session(kb_dir)
56+
result = await _handle_slash("/status", kb_dir, session, _STYLE)
57+
assert result is None
58+
output = capsys.readouterr().out
59+
assert "Knowledge Base Status" in output
60+
61+
62+
@pytest.mark.asyncio
63+
async def test_slash_list_empty(tmp_path, capsys):
64+
kb_dir = _setup_kb(tmp_path)
65+
session = _make_session(kb_dir)
66+
result = await _handle_slash("/list", kb_dir, session, _STYLE)
67+
assert result is None
68+
output = capsys.readouterr().out
69+
assert "No documents indexed yet" in output
70+
71+
72+
@pytest.mark.asyncio
73+
async def test_slash_list_with_docs(tmp_path, capsys):
74+
kb_dir = _setup_kb(tmp_path)
75+
hashes = {"abc": {"name": "paper.pdf", "type": "pdf"}}
76+
(kb_dir / ".openkb" / "hashes.json").write_text(json.dumps(hashes))
77+
session = _make_session(kb_dir)
78+
result = await _handle_slash("/list", kb_dir, session, _STYLE)
79+
assert result is None
80+
output = capsys.readouterr().out
81+
assert "paper.pdf" in output
82+
83+
84+
# --- /add, /exit, /clear, /help, /unknown use _fmt → need patching ---
85+
86+
87+
@pytest.mark.asyncio
88+
async def test_slash_add_missing_arg(tmp_path):
89+
kb_dir = _setup_kb(tmp_path)
90+
session = _make_session(kb_dir)
91+
p, collected = _collect_fmt()
92+
with p:
93+
result = await _handle_slash("/add", kb_dir, session, _STYLE)
94+
assert result is None
95+
assert any("Usage: /add <path>" in s for s in collected)
96+
97+
98+
@pytest.mark.asyncio
99+
async def test_slash_add_nonexistent_path(tmp_path):
100+
kb_dir = _setup_kb(tmp_path)
101+
session = _make_session(kb_dir)
102+
p, collected = _collect_fmt()
103+
with p:
104+
result = await _handle_slash("/add /no/such/path", kb_dir, session, _STYLE)
105+
assert result is None
106+
assert any("Path does not exist" in s for s in collected)
107+
108+
109+
@pytest.mark.asyncio
110+
async def test_slash_add_unsupported_type(tmp_path):
111+
kb_dir = _setup_kb(tmp_path)
112+
bad_file = tmp_path / "file.xyz"
113+
bad_file.write_text("data")
114+
session = _make_session(kb_dir)
115+
p, collected = _collect_fmt()
116+
with p:
117+
result = await _handle_slash(f"/add {bad_file}", kb_dir, session, _STYLE)
118+
assert result is None
119+
assert any("Unsupported file type" in s for s in collected)
120+
121+
122+
@pytest.mark.asyncio
123+
async def test_slash_add_single_file(tmp_path):
124+
kb_dir = _setup_kb(tmp_path)
125+
doc = tmp_path / "test.md"
126+
doc.write_text("# Hello")
127+
p, _collected = _collect_fmt()
128+
with p, patch("openkb.cli.add_single_file") as mock_add:
129+
await _run_add(str(doc), kb_dir, _STYLE)
130+
mock_add.assert_called_once_with(doc, kb_dir)
131+
132+
133+
@pytest.mark.asyncio
134+
async def test_slash_add_directory_with_progress(tmp_path):
135+
kb_dir = _setup_kb(tmp_path)
136+
docs_dir = tmp_path / "docs"
137+
docs_dir.mkdir()
138+
(docs_dir / "a.md").write_text("# A")
139+
(docs_dir / "b.txt").write_text("B")
140+
(docs_dir / "skip.xyz").write_text("skip")
141+
p, collected = _collect_fmt()
142+
with p, patch("openkb.cli.add_single_file") as mock_add:
143+
await _run_add(str(docs_dir), kb_dir, _STYLE)
144+
assert mock_add.call_count == 2
145+
output = "".join(collected)
146+
assert "Found 2 supported file(s)" in output
147+
assert "[1/2]" in output
148+
assert "[2/2]" in output
149+
150+
151+
@pytest.mark.asyncio
152+
async def test_slash_lint(tmp_path):
153+
kb_dir = _setup_kb(tmp_path)
154+
session = _make_session(kb_dir)
155+
with patch("openkb.cli.run_lint", new_callable=AsyncMock, return_value=tmp_path / "report.md"):
156+
result = await _handle_slash("/lint", kb_dir, session, _STYLE)
157+
assert result is None
158+
159+
160+
@pytest.mark.asyncio
161+
async def test_slash_unknown(tmp_path):
162+
kb_dir = _setup_kb(tmp_path)
163+
session = _make_session(kb_dir)
164+
p, collected = _collect_fmt()
165+
with p:
166+
result = await _handle_slash("/foobar", kb_dir, session, _STYLE)
167+
assert result is None
168+
assert any("Unknown command" in s for s in collected)
169+
170+
171+
@pytest.mark.asyncio
172+
async def test_slash_exit(tmp_path):
173+
kb_dir = _setup_kb(tmp_path)
174+
session = _make_session(kb_dir)
175+
p, _collected = _collect_fmt()
176+
with p:
177+
result = await _handle_slash("/exit", kb_dir, session, _STYLE)
178+
assert result == "exit"
179+
180+
181+
@pytest.mark.asyncio
182+
async def test_slash_clear(tmp_path):
183+
kb_dir = _setup_kb(tmp_path)
184+
session = _make_session(kb_dir)
185+
p, _collected = _collect_fmt()
186+
with p:
187+
result = await _handle_slash("/clear", kb_dir, session, _STYLE)
188+
assert result == "new_session"

0 commit comments

Comments
 (0)