diff --git a/README.md b/README.md index 1555c77..791c13c 100644 --- a/README.md +++ b/README.md @@ -24,11 +24,12 @@ Traditional RAG rediscovers knowledge from scratch on every query. Nothing accum ### Features -- **Any format** — PDF, Word, PowerPoint, Excel, HTML, Markdown, text, CSV, and more via markitdown +- **Broad format support** — PDF, Word, Markdown, PowerPoint, HTML, Excel, CSV, text, and more via markitdown - **Scale to long documents** — Long and complex documents are handled via [PageIndex](https://github.com/VectifyAI/PageIndex) tree indexing, enabling accurate, vectorless long-context retrieval - **Native multi-modality** — Retrieves and understands figures, tables, and images, not just text -- **Auto wiki** — LLM generates summaries, concept pages, and cross-links. You curate sources; the LLM does the rest -- **Query** — Ask questions against your wiki. The LLM navigates your compiled knowledge to answer +- **Compiled Wiki** — LLM manages and compiles your documents into summaries, concept pages, and cross-links, all kept in sync +- **Query** — Ask questions (one-off) against your wiki. The LLM navigates your compiled knowledge to answer +- **Interactive Chat** — Multi-turn conversations with persisted sessions you can resume across runs - **Lint** — Health checks find contradictions, gaps, orphans, and stale content - **Watch mode** — Drop files into `raw/`, wiki updates automatically - **Obsidian compatible** — Wiki is plain `.md` files with `[[wikilinks]]`. Open in Obsidian for graph view and browsing @@ -55,11 +56,11 @@ openkb add paper.pdf openkb add ~/papers/ # Add a whole directory openkb add article.html -# 4. Ask questions +# 4. Ask a question openkb query "What are the main findings?" -# 5. Check wiki health -openkb lint +# 5. Or start an interactive chat session +openkb chat ``` ### Set up your LLM @@ -132,6 +133,7 @@ A single source might touch 10-15 wiki pages. Knowledge accumulates: each docume | `openkb add ` | Add documents and compile to wiki | | `openkb query "question"` | Ask a question against the knowledge base | | `openkb query "question" --save` | Ask and save the answer to `wiki/explorations/` | +| `openkb chat` | Start an interactive multi-turn chat (use `--resume`, `--list`, `--delete` to manage sessions) | | `openkb watch` | Watch `raw/` and auto-compile new files | | `openkb lint` | Run structural + knowledge health checks | | `openkb list` | List indexed documents and concepts | @@ -139,6 +141,20 @@ A single source might touch 10-15 wiki pages. Knowledge accumulates: each docume +### Interactive chat + +`openkb chat` opens an interactive chat session over your wiki knowledge base. Unlike the one-shot `openkb query`, each turn carries the conversation history, so you can dig into a topic without re-typing context. + +```bash +openkb chat # start a new session +openkb chat --resume # resume the most recent session +openkb chat --resume 20260411 # resume by id (unique prefix works) +openkb chat --list # list all sessions +openkb chat --delete # delete a session +``` + +`/help` lists all slash commands: e.g., `/save` exports the transcript, `/clear` starts a fresh session. + ### Configuration Settings are initialized by `openkb init`, and stored in `.openkb/config.yaml`: diff --git a/openkb/__init__.py b/openkb/__init__.py index 3dc1f76..8db482e 100644 --- a/openkb/__init__.py +++ b/openkb/__init__.py @@ -1 +1,7 @@ -__version__ = "0.1.0" +"""OpenKB package.""" +from importlib.metadata import PackageNotFoundError, version as _version + +try: + __version__ = _version("openkb") +except PackageNotFoundError: + __version__ = "0.0.0+unknown" diff --git a/openkb/agent/chat.py b/openkb/agent/chat.py new file mode 100644 index 0000000..42ac9f9 --- /dev/null +++ b/openkb/agent/chat.py @@ -0,0 +1,378 @@ +"""Interactive multi-turn chat REPL for the OpenKB knowledge base. + +Builds on the single-shot Q&A agent in ``openkb.agent.query`` and keeps +conversation state in ``ChatSession``. Uses prompt_toolkit for the input +line (history, editing, bottom toolbar) and streams responses directly to +stdout to preserve the existing ``query`` visual. +""" +from __future__ import annotations + +import os +import re +import sys +import time +from pathlib import Path +from typing import Any + +from prompt_toolkit import PromptSession +from prompt_toolkit.formatted_text import FormattedText +from prompt_toolkit.shortcuts import print_formatted_text +from prompt_toolkit.styles import Style + +from openkb.agent.chat_session import ChatSession +from openkb.agent.query import MAX_TURNS, build_query_agent +from openkb.log import append_log + + +_STYLE_DICT: dict[str, str] = { + "prompt": "bold #5fa0e0", + "bottom-toolbar": "noreverse nobold #8a8a8a bg:default", + "toolbar": "noreverse nobold #8a8a8a bg:default", + "toolbar.session": "noreverse #8a8a8a bg:default bold", + "header": "#8a8a8a", + "header.title": "bold #5fa0e0", + "tool": "#a8a8a8", + "tool.name": "#a8a8a8 bold", + "slash.ok": "ansigreen", + "slash.help": "#8a8a8a", + "error": "ansired bold", + "resume.turn": "#5fa0e0", + "resume.user": "bold", + "resume.assistant": "#8a8a8a", +} + +_HELP_TEXT = ( + "Commands:\n" + " /exit Exit (Ctrl-D also works)\n" + " /clear Start a fresh session (current one is kept on disk)\n" + " /save [name] Export transcript to wiki/explorations/\n" + " /help Show this" +) + +_SIGINT_EXIT_WINDOW = 2.0 + + +def _use_color(force_off: bool) -> bool: + if force_off: + return False + if os.environ.get("NO_COLOR", ""): + return False + if not sys.stdout.isatty(): + return False + return True + + +def _build_style(use_color: bool) -> Style: + return Style.from_dict(_STYLE_DICT if use_color else {}) + + +def _fmt(style: Style, *fragments: tuple[str, str]) -> None: + print_formatted_text(FormattedText(list(fragments)), style=style, end="") + + +def _format_tool_line(name: str, args: str, width: int = 78) -> str: + args = args or "" + args = args.replace("\n", " ") + base = f" \u00b7 {name}({args})" + if len(base) > width: + base = base[: width - 1] + "\u2026" + return base + + +def _extract_preview(text: str, limit: int = 150) -> str: + text = " ".join((text or "").strip().split()) + if len(text) <= limit: + return text + return text[: limit - 1] + "\u2026" + + +def _openkb_version() -> str: + from openkb import __version__ + return __version__ + + +def _display_kb_dir(kb_dir: Path) -> str: + home = str(Path.home()) + s = str(kb_dir) + if s == home: + return "~" + if s.startswith(home + "/"): + return "~" + s[len(home):] + return s + + +def _print_header(session: ChatSession, kb_dir: Path, style: Style) -> None: + disp_dir = _display_kb_dir(kb_dir) + version = _openkb_version() + version_suffix = f" v{version}\n" if version else "\n" + print() + _fmt( + style, + ("class:header.title", "OpenKB Chat"), + ("class:header", version_suffix), + ) + _fmt( + style, + ( + "class:header", + f"{disp_dir} \u00b7 {session.model} \u00b7 session {session.id}\n", + ), + ) + _fmt( + style, + ( + "class:header", + "Type /help for commands, Ctrl-D to exit, " + "Ctrl-C to abort current response.\n", + ), + ) + print() + + +def _print_resume_view(session: ChatSession, style: Style) -> None: + turns = list(zip(session.user_turns, session.assistant_texts)) + if not turns: + return + total = len(turns) + if total > 5: + omitted = total - 5 + _fmt( + style, + ("class:header", f"... {omitted} earlier turn(s) omitted\n"), + ) + turns = turns[-5:] + start = omitted + 1 + else: + start = 1 + + _fmt( + style, + ("class:header", f"Resumed session {total} turn(s)\n"), + ) + for i, (u, a) in enumerate(turns, start): + _fmt( + style, + ("class:resume.turn", f"[{i}] "), + ("class:resume.user", f">>> {u}\n"), + ) + if a: + preview = _extract_preview(a, 180) + extra = "" + if len(a) > len(preview): + extra = f" ({len(a)} chars)" + _fmt( + style, + ("class:resume.turn", f"[{i}] "), + ("class:resume.assistant", f" {preview}{extra}\n"), + ) + print() + + +def _bottom_toolbar(session: ChatSession) -> FormattedText: + return FormattedText( + [ + ("class:toolbar", " session "), + ("class:toolbar.session", session.id), + ( + "class:toolbar", + f" {session.turn_count} turn(s) {session.model} ", + ), + ] + ) + + +def _make_prompt_session(session: ChatSession, style: Style, use_color: bool) -> PromptSession: + return PromptSession( + message=FormattedText([("class:prompt", ">>> ")]), + style=style, + bottom_toolbar=(lambda: _bottom_toolbar(session)) if use_color else None, + ) + + +async def _run_turn(agent: Any, session: ChatSession, user_input: str, style: Style) -> None: + """Run one agent turn with streaming output and persist the new history.""" + from agents import ( + RawResponsesStreamEvent, + RunItemStreamEvent, + Runner, + ) + from openai.types.responses import ResponseTextDeltaEvent + + new_input = session.history + [{"role": "user", "content": user_input}] + + result = Runner.run_streamed(agent, new_input, max_turns=MAX_TURNS) + + sys.stdout.write("\n") + sys.stdout.flush() + collected: list[str] = [] + last_was_text = False + need_blank_before_text = False + try: + async for event in result.stream_events(): + if isinstance(event, RawResponsesStreamEvent): + if isinstance(event.data, ResponseTextDeltaEvent): + text = event.data.delta + if text: + if need_blank_before_text: + sys.stdout.write("\n") + need_blank_before_text = False + sys.stdout.write(text) + sys.stdout.flush() + collected.append(text) + last_was_text = True + elif isinstance(event, RunItemStreamEvent): + item = event.item + if item.type == "tool_call_item": + if last_was_text: + sys.stdout.write("\n") + sys.stdout.flush() + last_was_text = False + raw = item.raw_item + name = getattr(raw, "name", "?") + args = getattr(raw, "arguments", "") or "" + _fmt(style, ("class:tool", _format_tool_line(name, args) + "\n")) + need_blank_before_text = True + finally: + sys.stdout.write("\n\n") + sys.stdout.flush() + + answer = "".join(collected).strip() + if not answer: + answer = (result.final_output or "").strip() + session.record_turn(user_input, answer, result.to_input_list()) + + +def _save_transcript(kb_dir: Path, session: ChatSession, name: str | None) -> Path: + explore_dir = kb_dir / "wiki" / "explorations" + explore_dir.mkdir(parents=True, exist_ok=True) + + base = name or session.title or (session.user_turns[0] if session.user_turns else session.id) + slug = re.sub(r"[^a-z0-9]+", "-", base.lower()).strip("-")[:60] or session.id + date = session.created_at[:10].replace("-", "") + path = explore_dir / f"{slug}-{date}.md" + + lines: list[str] = [ + "---", + f'session: "{session.id}"', + f'model: "{session.model}"', + f'created: "{session.created_at}"', + "---", + "", + f"# Chat transcript {session.title or session.id}", + "", + ] + for i, (u, a) in enumerate(zip(session.user_turns, session.assistant_texts), 1): + lines.append(f"## [{i}] {u}") + lines.append("") + lines.append(a or "_(no response recorded)_") + lines.append("") + + path.write_text("\n".join(lines), encoding="utf-8") + return path + + +async def _handle_slash( + cmd: str, + kb_dir: Path, + session: ChatSession, + style: Style, +) -> str | None: + """Return ``"exit"`` to end the REPL, ``"new_session"`` to swap sessions, + or ``None`` to continue with the current session.""" + parts = cmd.split(maxsplit=1) + head = parts[0].lower() + arg = parts[1].strip() if len(parts) > 1 else "" + + if head in ("/exit", "/quit"): + _fmt(style, ("class:header", "Bye. Thanks for using OpenKB.\n\n")) + return "exit" + + if head == "/help": + _fmt(style, ("class:slash.help", _HELP_TEXT + "\n")) + return None + + if head == "/clear": + old_id = session.id + _fmt( + style, + ("class:slash.ok", f"Started new session (previous: {old_id})\n"), + ) + return "new_session" + + if head == "/save": + if not session.user_turns: + _fmt(style, ("class:error", "Nothing to save yet.\n")) + return None + path = _save_transcript(kb_dir, session, arg or None) + _fmt(style, ("class:slash.ok", f"Saved to {path}\n")) + return None + + _fmt( + style, + ("class:error", f"Unknown command: {head}. Try /help.\n"), + ) + return None + + +async def run_chat( + kb_dir: Path, + session: ChatSession, + *, + no_color: bool = False, +) -> None: + """Run the chat REPL against ``session`` until the user exits.""" + from openkb.config import load_config + + use_color = _use_color(force_off=no_color) + style = _build_style(use_color) + + config = load_config(kb_dir / ".openkb" / "config.yaml") + language = session.language or config.get("language", "en") + wiki_root = str(kb_dir / "wiki") + agent = build_query_agent(wiki_root, session.model, language=language) + + _print_header(session, kb_dir, style) + if session.turn_count > 0: + _print_resume_view(session, style) + + prompt_session = _make_prompt_session(session, style, use_color) + + last_sigint = 0.0 + + while True: + try: + user_input = await prompt_session.prompt_async() + last_sigint = 0.0 + except KeyboardInterrupt: + now = time.monotonic() + if last_sigint and (now - last_sigint) < _SIGINT_EXIT_WINDOW: + _fmt(style, ("class:header", "\nBye. Thanks for using OpenKB.\n\n")) + return + last_sigint = now + _fmt(style, ("class:header", "\n(Press Ctrl-C again to exit)\n")) + continue + except EOFError: + _fmt(style, ("class:header", "Bye. Thanks for using OpenKB.\n\n")) + return + + user_input = (user_input or "").strip() + if not user_input: + continue + + if user_input.startswith("/"): + action = await _handle_slash(user_input, kb_dir, session, style) + if action == "exit": + return + if action == "new_session": + session = ChatSession.new(kb_dir, session.model, session.language) + agent = build_query_agent(wiki_root, session.model, language=language) + prompt_session = _make_prompt_session(session, style, use_color) + continue + + append_log(kb_dir / "wiki", "query", user_input) + try: + await _run_turn(agent, session, user_input, style) + except KeyboardInterrupt: + _fmt(style, ("class:error", "\n[aborted]\n")) + except Exception as exc: + _fmt(style, ("class:error", f"[ERROR] {exc}\n")) diff --git a/openkb/agent/chat_session.py b/openkb/agent/chat_session.py new file mode 100644 index 0000000..01706ea --- /dev/null +++ b/openkb/agent/chat_session.py @@ -0,0 +1,280 @@ +"""Chat session persistence for `openkb chat`. + +Each session lives in ``/.openkb/chats/.json`` and stores a sanitized +agent-SDK history (from ``RunResult.to_input_list()``) alongside the user +messages and full assistant replies kept as plain strings for display and +export. Large tool-returned image payloads are replaced with lightweight +references before the history is reused or persisted. +""" +from __future__ import annotations + +import json +import os +import random +import string +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + + +_IMAGE_HISTORY_NOTE = ( + "Image output omitted from chat history to avoid persisting raw data URLs." +) + + +def _utcnow_iso() -> str: + return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _gen_id() -> str: + ts = datetime.now().strftime("%Y%m%d-%H%M%S") + rand = "".join(random.choices(string.ascii_lowercase + string.digits, k=3)) + return f"{ts}-{rand}" + + +def chats_dir(kb_dir: Path) -> Path: + return kb_dir / ".openkb" / "chats" + + +def _title_from(msg: str, limit: int = 60) -> str: + msg = " ".join(msg.strip().split()) + if len(msg) <= limit: + return msg + return msg[: limit - 1] + "\u2026" + + +def _image_history_placeholder(image_path: str | None) -> dict[str, str]: + text = _IMAGE_HISTORY_NOTE + if image_path: + text += f" Source path: {image_path}." + text += " Call get_image again if you need to inspect it." + return {"type": "input_text", "text": text} + + +def _extract_get_image_path(item: dict[str, Any]) -> str | None: + if item.get("type") != "function_call" or item.get("name") != "get_image": + return None + arguments = item.get("arguments") + if not isinstance(arguments, str): + return None + try: + payload = json.loads(arguments) + except json.JSONDecodeError: + return None + image_path = payload.get("image_path") + if isinstance(image_path, str) and image_path: + return image_path + return None + + +def _sanitize_history_value(value: Any, image_path: str | None = None) -> Any: + if isinstance(value, list): + return [_sanitize_history_value(item, image_path) for item in value] + if not isinstance(value, dict): + return value + + if value.get("type") == "input_image": + image_url = value.get("image_url") + if isinstance(image_url, str) and image_url.startswith("data:"): + return _image_history_placeholder(image_path) + + return { + key: _sanitize_history_value(item, image_path) + for key, item in value.items() + } + + +def sanitize_history(history: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Strip large image payloads from model history while keeping a re-fetch hint.""" + image_paths_by_call_id: dict[str, str] = {} + sanitized: list[dict[str, Any]] = [] + + for item in history: + if not isinstance(item, dict): + sanitized.append(item) + continue + + image_path = _extract_get_image_path(item) + call_id = item.get("call_id") + if image_path and isinstance(call_id, str): + image_paths_by_call_id[call_id] = image_path + + history_image_path = None + if item.get("type") == "function_call_output" and isinstance(call_id, str): + history_image_path = image_paths_by_call_id.get(call_id) + + sanitized.append(_sanitize_history_value(item, history_image_path)) + + return sanitized + + +@dataclass +class ChatSession: + id: str + created_at: str + updated_at: str + model: str + language: str + title: str + turn_count: int + history: list[dict[str, Any]] + user_turns: list[str] + assistant_texts: list[str] + path: Path + + @classmethod + def new(cls, kb_dir: Path, model: str, language: str) -> "ChatSession": + now = _utcnow_iso() + sid = _gen_id() + return cls( + id=sid, + created_at=now, + updated_at=now, + model=model, + language=language, + title="", + turn_count=0, + history=[], + user_turns=[], + assistant_texts=[], + path=chats_dir(kb_dir) / f"{sid}.json", + ) + + def to_dict(self) -> dict[str, Any]: + return { + "id": self.id, + "created_at": self.created_at, + "updated_at": self.updated_at, + "model": self.model, + "language": self.language, + "title": self.title, + "turn_count": self.turn_count, + "history": self.history, + "user_turns": self.user_turns, + "assistant_texts": self.assistant_texts, + } + + def save(self) -> None: + self.path.parent.mkdir(parents=True, exist_ok=True) + tmp = self.path.with_suffix(".json.tmp") + tmp.write_text( + json.dumps(self.to_dict(), ensure_ascii=False, indent=2, default=str), + encoding="utf-8", + ) + os.replace(tmp, self.path) + + def record_turn( + self, + user_message: str, + assistant_text: str, + new_history: list[dict[str, Any]], + ) -> None: + self.history = sanitize_history(new_history) + self.user_turns.append(user_message) + self.assistant_texts.append(assistant_text) + self.turn_count = len(self.user_turns) + if not self.title: + self.title = _title_from(user_message) + self.updated_at = _utcnow_iso() + self.save() + + +def load_session(kb_dir: Path, session_id: str) -> ChatSession: + path = chats_dir(kb_dir) / f"{session_id}.json" + data = json.loads(path.read_text(encoding="utf-8")) + return ChatSession( + id=data["id"], + created_at=data["created_at"], + updated_at=data["updated_at"], + model=data["model"], + language=data.get("language", "en"), + title=data.get("title", ""), + turn_count=data.get("turn_count", 0), + history=sanitize_history(data.get("history", [])), + user_turns=data.get("user_turns", []), + assistant_texts=data.get("assistant_texts", []), + path=path, + ) + + +def list_sessions(kb_dir: Path) -> list[dict[str, Any]]: + """Return session metadata dicts, most recently updated first.""" + d = chats_dir(kb_dir) + if not d.exists(): + return [] + out: list[dict[str, Any]] = [] + for p in d.glob("*.json"): + try: + data = json.loads(p.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + continue + out.append( + { + "id": data.get("id", p.stem), + "title": data.get("title", ""), + "turn_count": data.get("turn_count", 0), + "updated_at": data.get("updated_at", ""), + "model": data.get("model", ""), + } + ) + out.sort(key=lambda s: (s["updated_at"], s["id"]), reverse=True) + return out + + +def resolve_session_id(kb_dir: Path, query: str) -> str | None: + """Resolve a query to a full session id. + + ``query`` may be: + - ``"__latest__"`` — returns the most recently updated session id. + - A full session id — returned as-is if it exists. + - A unique prefix of a session id — expanded to the full id. + + Returns ``None`` if no session matches. Raises ``ValueError`` when a + prefix is ambiguous. + """ + sessions = list_sessions(kb_dir) + if not sessions: + return None + if query == "__latest__": + return sessions[0]["id"] + for s in sessions: + if s["id"] == query: + return s["id"] + matches = [s["id"] for s in sessions if s["id"].startswith(query)] + if len(matches) == 1: + return matches[0] + if len(matches) > 1: + raise ValueError( + f"Ambiguous session prefix '{query}' matches: {', '.join(matches)}" + ) + return None + + +def delete_session(kb_dir: Path, session_id: str) -> bool: + path = chats_dir(kb_dir) / f"{session_id}.json" + if path.exists(): + path.unlink() + return True + return False + + +def relative_time(iso_str: str) -> str: + """Render an ISO-8601 timestamp as a short relative string.""" + try: + t = datetime.strptime(iso_str, "%Y-%m-%dT%H:%M:%SZ").replace( + tzinfo=timezone.utc + ) + except (ValueError, TypeError): + return iso_str or "" + now = datetime.now(timezone.utc) + seconds = int((now - t).total_seconds()) + if seconds < 60: + return "just now" + if seconds < 3600: + return f"{seconds // 60}m ago" + if seconds < 86400: + return f"{seconds // 3600}h ago" + if seconds < 86400 * 7: + return f"{seconds // 86400}d ago" + return t.strftime("%Y-%m-%d") diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index d94a558..d202fc4 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -30,7 +30,7 @@ # --------------------------------------------------------------------------- _SYSTEM_TEMPLATE = """\ -You are a wiki compilation agent for a personal knowledge base. +You are OpenKB's wiki compilation agent for a personal knowledge base. {schema_md} diff --git a/openkb/agent/linter.py b/openkb/agent/linter.py index fb81da7..113a176 100644 --- a/openkb/agent/linter.py +++ b/openkb/agent/linter.py @@ -11,7 +11,7 @@ from openkb.schema import SCHEMA_MD, get_agents_md _LINTER_INSTRUCTIONS_TEMPLATE = """\ -You are a knowledge-base semantic lint agent. Your job is to audit the wiki +You are OpenKB's semantic lint agent. Your job is to audit the wiki for quality issues that structural tools cannot detect. {schema_md} @@ -50,7 +50,7 @@ def build_lint_agent(wiki_root: str, model: str, language: str = "en") -> Agent: """ schema_md = get_agents_md(Path(wiki_root)) instructions = _LINTER_INSTRUCTIONS_TEMPLATE.format(schema_md=schema_md) - instructions += f"\n\nIMPORTANT: Write all wiki content in {language} language." + instructions += f"\n\nIMPORTANT: Write the lint report in {language} language." @function_tool def list_files(directory: str) -> str: diff --git a/openkb/agent/query.py b/openkb/agent/query.py index d252ee6..39e0e40 100644 --- a/openkb/agent/query.py +++ b/openkb/agent/query.py @@ -6,13 +6,13 @@ from agents import Agent, Runner, function_tool from agents import ToolOutputImage, ToolOutputText -from openkb.agent.tools import read_wiki_file, read_wiki_image +from openkb.agent.tools import get_wiki_page_content, read_wiki_file, read_wiki_image MAX_TURNS = 50 from openkb.schema import get_agents_md _QUERY_INSTRUCTIONS_TEMPLATE = """\ -You are a knowledge-base Q&A agent. You answer questions by searching the wiki. +You are OpenKB, a knowledge-base Q&A agent. You answer questions by searching the wiki. {schema_md} @@ -20,7 +20,8 @@ 1. Read index.md to see all documents and concepts with brief summaries. Each document is marked (short) or (pageindex) to indicate its type. 2. Read relevant summary pages (summaries/) for document overviews. - Note: summaries may omit details. + Summaries may omit details — if you need more, follow the summary's + `full_text` frontmatter field to the source (see step 4). 3. Read concept pages (concepts/) for cross-document synthesis. 4. When you need detailed source document content, each summary page has a `full_text` frontmatter field with the path to the original document content: @@ -28,9 +29,8 @@ - PageIndex documents (doc_type: pageindex): use get_page_content(doc_name, pages) with tight page ranges. The summary shows document tree structure with page ranges to help you target. Never fetch the whole document. -5. When source content references images (e.g. ![image](sources/images/doc/file.png)), - use get_image to view them. Always view images when the question asks about - a figure, chart, diagram, or visual content. +5. Source content may reference images (e.g. ![image](sources/images/doc/file.png)). + Use the get_image tool to view them when needed. 6. Synthesize a clear, concise, well-cited answer grounded in wiki content. Answer based only on wiki content. Be concise. @@ -44,7 +44,7 @@ def build_query_agent(wiki_root: str, model: str, language: str = "en") -> Agent """Build and return the Q&A agent.""" schema_md = get_agents_md(Path(wiki_root)) instructions = _QUERY_INSTRUCTIONS_TEMPLATE.format(schema_md=schema_md) - instructions += f"\n\nIMPORTANT: Write all wiki content in {language} language." + instructions += f"\n\nIMPORTANT: Answer in {language} language." @function_tool def read_file(path: str) -> str: @@ -55,7 +55,7 @@ def read_file(path: str) -> str: return read_wiki_file(path, wiki_root) @function_tool - def get_page_content_tool(doc_name: str, pages: str) -> str: + def get_page_content(doc_name: str, pages: str) -> str: """Get text content of specific pages from a PageIndex (long) document. Only use for documents with doc_type: pageindex. For short documents, use read_file instead. @@ -63,13 +63,15 @@ def get_page_content_tool(doc_name: str, pages: str) -> str: doc_name: Document name (e.g. 'attention-is-all-you-need'). pages: Page specification (e.g. '3-5,7,10-12'). """ - from openkb.agent.tools import get_page_content - return get_page_content(doc_name, pages, wiki_root) + return get_wiki_page_content(doc_name, pages, wiki_root) @function_tool def get_image(image_path: str) -> ToolOutputImage | ToolOutputText: """View an image from the wiki. - Use when source content references images you need to see. + + Use when a question asks about a specific figure, chart, or diagram + you'd need to see to answer accurately. + Args: image_path: Image path relative to wiki root (e.g. 'sources/images/doc/p1_img1.png'). """ @@ -83,7 +85,7 @@ def get_image(image_path: str) -> ToolOutputImage | ToolOutputText: return Agent( name="wiki-query", instructions=instructions, - tools=[read_file, get_page_content_tool, get_image], + tools=[read_file, get_page_content, get_image], model=f"litellm/{model}", model_settings=ModelSettings(parallel_tool_calls=False), ) diff --git a/openkb/agent/tools.py b/openkb/agent/tools.py index 2fe930b..e905fda 100644 --- a/openkb/agent/tools.py +++ b/openkb/agent/tools.py @@ -89,7 +89,7 @@ def parse_pages(pages: str) -> list[int]: return sorted(n for n in result if n > 0) -def get_page_content(doc_name: str, pages: str, wiki_root: str) -> str: +def get_wiki_page_content(doc_name: str, pages: str, wiki_root: str) -> str: """Return formatted content for specified pages of a document. Reads ``{wiki_root}/sources/{doc_name}.json`` which must be a JSON array of diff --git a/openkb/cli.py b/openkb/cli.py index 550ee5c..028e546 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -395,6 +395,107 @@ def query(ctx, question, save): click.echo(f"\nSaved to {explore_path}") +@cli.command() +@click.option( + "--resume", "-r", "resume", + is_flag=False, flag_value="__latest__", default=None, metavar="[ID]", + help="Resume the latest chat session, or a specific one by id or prefix.", +) +@click.option( + "--list", "list_sessions_flag", + is_flag=True, default=False, + help="List chat sessions.", +) +@click.option( + "--delete", "delete_id", + default=None, metavar="ID", + help="Delete a chat session by id or prefix.", +) +@click.option( + "--no-color", "no_color", + is_flag=True, default=False, + help="Disable colored output.", +) +@click.pass_context +def chat(ctx, resume, list_sessions_flag, delete_id, no_color): + """Start an interactive chat with the knowledge base.""" + kb_dir = _find_kb_dir(ctx.obj.get("kb_dir_override")) + if kb_dir is None: + click.echo("No knowledge base found. Run `openkb init` first.") + return + + from openkb.agent.chat_session import ( + ChatSession, + delete_session, + list_sessions, + load_session, + relative_time, + resolve_session_id, + ) + + if list_sessions_flag: + sessions = list_sessions(kb_dir) + if not sessions: + click.echo("No chat sessions yet.") + return + click.echo(f" {'ID':<22} {'TURNS':<6} {'UPDATED':<12} TITLE") + click.echo(f" {'-'*22} {'-'*6} {'-'*12} {'-'*30}") + for s in sessions: + rel = relative_time(s.get("updated_at", "")) + title = s.get("title") or "(empty)" + click.echo( + f" {s['id']:<22} {s['turn_count']:<6} {rel:<12} {title}" + ) + click.echo( + f"\n{len(sessions)} session(s) in {kb_dir / '.openkb' / 'chats'}" + ) + return + + if delete_id is not None: + try: + resolved = resolve_session_id(kb_dir, delete_id) + except ValueError as exc: + click.echo(f"[ERROR] {exc}") + return + if not resolved: + click.echo(f"No matching session: {delete_id}") + return + if delete_session(kb_dir, resolved): + click.echo(f"Deleted session {resolved}") + else: + click.echo(f"Could not delete session: {resolved}") + return + + openkb_dir = kb_dir / ".openkb" + config = load_config(openkb_dir / "config.yaml") + _setup_llm_key(kb_dir) + + if resume is not None: + try: + resolved = resolve_session_id(kb_dir, resume) + except ValueError as exc: + click.echo(f"[ERROR] {exc}") + return + if not resolved: + if resume == "__latest__": + click.echo("No previous chat sessions to resume.") + else: + click.echo(f"No matching session: {resume}") + return + session = load_session(kb_dir, resolved) + else: + model: str = config.get("model", DEFAULT_CONFIG["model"]) + language: str = config.get("language", "en") + session = ChatSession.new(kb_dir, model, language) + + from openkb.agent.chat import run_chat + + try: + asyncio.run(run_chat(kb_dir, session, no_color=no_color)) + except Exception as exc: + click.echo(f"[ERROR] Chat failed: {exc}") + + @cli.command() @click.pass_context def watch(ctx): diff --git a/pyproject.toml b/pyproject.toml index 264ab9e..4af87be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "pyyaml", "python-dotenv", "json-repair", + "prompt_toolkit>=3.0", ] [project.urls] diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 3d95a88..196aae0 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -5,7 +5,7 @@ import pytest -from openkb.agent.tools import get_page_content, list_wiki_files, parse_pages, read_wiki_file, write_wiki_file +from openkb.agent.tools import get_wiki_page_content, list_wiki_files, parse_pages, read_wiki_file, write_wiki_file # --------------------------------------------------------------------------- @@ -159,11 +159,11 @@ def test_ignores_zero_and_negative(self): # --------------------------------------------------------------------------- -# get_page_content +# get_wiki_page_content # --------------------------------------------------------------------------- -class TestGetPageContent: +class TestGetWikiPageContent: def test_reads_pages_from_json(self, tmp_path): import json wiki_root = str(tmp_path) @@ -175,7 +175,7 @@ def test_reads_pages_from_json(self, tmp_path): {"page": 3, "content": "Page three text."}, ] (sources / "paper.json").write_text(json.dumps(pages), encoding="utf-8") - result = get_page_content("paper", "1,3", wiki_root) + result = get_wiki_page_content("paper", "1,3", wiki_root) assert "[Page 1]" in result assert "Page one text." in result assert "[Page 3]" in result @@ -185,7 +185,7 @@ def test_reads_pages_from_json(self, tmp_path): def test_returns_error_for_missing_file(self, tmp_path): wiki_root = str(tmp_path) (tmp_path / "sources").mkdir() - result = get_page_content("nonexistent", "1", wiki_root) + result = get_wiki_page_content("nonexistent", "1", wiki_root) assert "not found" in result.lower() def test_returns_error_for_no_matching_pages(self, tmp_path): @@ -195,7 +195,7 @@ def test_returns_error_for_no_matching_pages(self, tmp_path): sources.mkdir() pages = [{"page": 1, "content": "Only page."}] (sources / "paper.json").write_text(json.dumps(pages), encoding="utf-8") - result = get_page_content("paper", "99", wiki_root) + result = get_wiki_page_content("paper", "99", wiki_root) assert "no content" in result.lower() def test_includes_images_info(self, tmp_path): @@ -205,11 +205,11 @@ def test_includes_images_info(self, tmp_path): sources.mkdir() pages = [{"page": 1, "content": "Text.", "images": [{"path": "images/p/img.png", "width": 100, "height": 80}]}] (sources / "doc.json").write_text(json.dumps(pages), encoding="utf-8") - result = get_page_content("doc", "1", wiki_root) + result = get_wiki_page_content("doc", "1", wiki_root) assert "img.png" in result def test_path_escape_denied(self, tmp_path): wiki_root = str(tmp_path) (tmp_path / "sources").mkdir() - result = get_page_content("../../etc/passwd", "1", wiki_root) + result = get_wiki_page_content("../../etc/passwd", "1", wiki_root) assert "denied" in result.lower() or "not found" in result.lower() diff --git a/tests/test_chat_session.py b/tests/test_chat_session.py new file mode 100644 index 0000000..759e02a --- /dev/null +++ b/tests/test_chat_session.py @@ -0,0 +1,76 @@ +"""Tests for chat session persistence.""" +from __future__ import annotations + +import json + +from openkb.agent.chat_session import ChatSession, load_session + + +def _image_history() -> list[dict[str, object]]: + return [ + {"role": "user", "content": "Describe the diagram."}, + { + "type": "function_call", + "call_id": "call_123", + "name": "get_image", + "arguments": '{"image_path":"sources/images/doc/figure-1.png"}', + }, + { + "type": "function_call_output", + "call_id": "call_123", + "output": [ + { + "type": "input_image", + "image_url": "data:image/png;base64,AAAA", + } + ], + }, + ] + + +def test_record_turn_replaces_data_image_with_text_reference(tmp_path): + session = ChatSession.new(tmp_path, "gpt-4o-mini", "en") + + session.record_turn( + "Describe the diagram.", + "It is a flow chart.", + _image_history(), + ) + + saved = json.loads(session.path.read_text(encoding="utf-8")) + output_part = saved["history"][2]["output"][0] + + assert output_part["type"] == "input_text" + assert "data:image/png;base64,AAAA" not in session.path.read_text(encoding="utf-8") + assert "sources/images/doc/figure-1.png" in output_part["text"] + assert "Call get_image again" in output_part["text"] + + +def test_load_session_sanitizes_legacy_image_history(tmp_path): + session = ChatSession.new(tmp_path, "gpt-4o-mini", "en") + raw_history = _image_history() + session.path.parent.mkdir(parents=True, exist_ok=True) + session.path.write_text( + json.dumps( + { + "id": session.id, + "created_at": session.created_at, + "updated_at": session.updated_at, + "model": session.model, + "language": session.language, + "title": "", + "turn_count": 1, + "history": raw_history, + "user_turns": ["Describe the diagram."], + "assistant_texts": ["It is a flow chart."], + } + ), + encoding="utf-8", + ) + + loaded = load_session(tmp_path, session.id) + + output_part = loaded.history[2]["output"][0] + assert output_part["type"] == "input_text" + assert "data:image/png;base64,AAAA" not in output_part["text"] + assert "sources/images/doc/figure-1.png" in output_part["text"] diff --git a/tests/test_query.py b/tests/test_query.py index e00d2ea..d39fcaa 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -23,7 +23,7 @@ def test_agent_tool_names(self, tmp_path): agent = build_query_agent(str(tmp_path), "gpt-4o-mini") names = {t.name for t in agent.tools} assert "read_file" in names - assert "get_page_content_tool" in names + assert "get_page_content" in names assert "get_image" in names def test_instructions_mention_get_page_content(self, tmp_path):