diff --git a/python/ebook_search/answer.py b/python/ebook_search/answer.py index c39eb2e..d7371a8 100644 --- a/python/ebook_search/answer.py +++ b/python/ebook_search/answer.py @@ -6,6 +6,7 @@ import logging from typing import TYPE_CHECKING from python.ebook_search.llm_interface import request_chat_completion +from python.ebook_search.prompts import load_prompt if TYPE_CHECKING: import httpx @@ -42,16 +43,7 @@ async def answer_query( content = await request_chat_completion( client, config, - [ - { - "role": "system", - "content": ( - "Answer only from the provided context. Cite sources with bracketed numbers like [1]. " - "If the context is insufficient, say so." - ), - }, - {"role": "user", "content": f"Question:\n{query}\n\nContext:\n{context}"}, - ], + load_prompt("answer").messages(query=query, context=context), ) logger.info(f"ebook_answer_request_complete {config.chat_model=} answer_length={len(content)}") diff --git a/python/ebook_search/prompts/__init__.py b/python/ebook_search/prompts/__init__.py new file mode 100644 index 0000000..c9c53aa --- /dev/null +++ b/python/ebook_search/prompts/__init__.py @@ -0,0 +1,8 @@ +"""LLM prompt templates for EPUB search.""" + +from python.ebook_search.prompts.lib import ( + Prompt, + load_prompt, +) + +__all__ = ["Prompt", "load_prompt"] diff --git a/python/ebook_search/prompts/answer.toml b/python/ebook_search/prompts/answer.toml new file mode 100644 index 0000000..963f5b9 --- /dev/null +++ b/python/ebook_search/prompts/answer.toml @@ -0,0 +1,9 @@ +system = """\ +Answer only from the provided context. Cite sources with bracketed numbers like [1]. \ +If the context is insufficient, say so.""" +user = """\ +Question: +{query} + +Context: +{context}""" diff --git a/python/ebook_search/prompts/lib.py b/python/ebook_search/prompts/lib.py new file mode 100644 index 0000000..86d903c --- /dev/null +++ b/python/ebook_search/prompts/lib.py @@ -0,0 +1,44 @@ +"""Load and render TOML-backed LLM prompt templates.""" + +from __future__ import annotations + +import tomllib +from dataclasses import dataclass +from functools import cache +from pathlib import Path + + +@dataclass(frozen=True) +class Prompt: + """A system and user prompt pair loaded from TOML.""" + + system: str + user: str + + def messages(self, **values: str) -> list[dict[str, str]]: + """Render this prompt as OpenAI-style chat messages.""" + return [ + {"role": "system", "content": self.system.format(**values)}, + {"role": "user", "content": self.user.format(**values)}, + ] + + +@cache +def _get_prompt_dir() -> Path: + """Return the directory containing prompt template files.""" + return Path(__file__).resolve().parent + + +@cache +def load_prompt(name: str) -> Prompt: + """Load and validate a named system and user prompt pair from TOML.""" + path = _get_prompt_dir() / f"{name}.toml" + with path.open("rb") as file: + body = tomllib.load(file) + + system = body.get("system") + user = body.get("user") + if not isinstance(system, str) or not isinstance(user, str): + msg = f"{path} must define string system and user prompts" + raise TypeError(msg) + return Prompt(system=system, user=user) diff --git a/python/ebook_search/prompts/phrase_judge.toml b/python/ebook_search/prompts/phrase_judge.toml new file mode 100644 index 0000000..fd75b96 --- /dev/null +++ b/python/ebook_search/prompts/phrase_judge.toml @@ -0,0 +1,10 @@ +system = """\ +Judge whether a candidate phrase from a book should be protected for RAG retrieval. \ +Do not extract new phrases. Reject common grammar fragments, ordinary nonspecific \ +phrases, unstable fragments, and phrases kept only because they are frequent. Keep \ +people, places, organizations, factions, events, technologies, fictional conditions, \ +magic systems, formal titles, named concepts, and recurring world-specific terms. \ +Return only a JSON object with keys: keep, canonical, category, aliases, confidence, \ +importance, allow_nested, suppress_children, reason.""" + +user = "{candidate_json}" diff --git a/python/ebook_search/protected_phrases/judge_ngrams.py b/python/ebook_search/protected_phrases/judge_ngrams.py index 54019dc..b0b49e6 100644 --- a/python/ebook_search/protected_phrases/judge_ngrams.py +++ b/python/ebook_search/protected_phrases/judge_ngrams.py @@ -14,6 +14,7 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from python.ebook_search.llm_interface import request_chat_completion +from python.ebook_search.prompts import load_prompt from python.ebook_search.protected_phrases.extraction import ( candidate_source_names, get_sample_contexts, @@ -369,20 +370,7 @@ def build_judge_messages(candidate: PhraseCandidate) -> list[dict[str, str]]: "chapter_count": candidate.chapter_count, "contexts": candidate.sample_contexts, } - return [ - { - "role": "system", - "content": ( - "Judge whether a candidate phrase from a book should be protected for RAG retrieval. " - "Do not extract new phrases. Reject common grammar fragments, ordinary nonspecific phrases, " - "unstable fragments, and phrases kept only because they are frequent. Keep people, places, " - "organizations, factions, events, technologies, fictional conditions, magic systems, formal titles, " - "named concepts, and recurring world-specific terms. Return only a JSON object with keys: keep, " - "canonical, category, aliases, confidence, importance, allow_nested, suppress_children, reason." - ), - }, - {"role": "user", "content": json.dumps(payload, ensure_ascii=True)}, - ] + return load_prompt("phrase_judge").messages(candidate_json=json.dumps(payload, ensure_ascii=True)) def parse_llm_judgment(content: str, config: EbookSearchConfig) -> LLMJudgment: diff --git a/tests/ebook_search/test_http.py b/tests/ebook_search/test_http.py index eae0dc1..f7d195f 100644 --- a/tests/ebook_search/test_http.py +++ b/tests/ebook_search/test_http.py @@ -93,6 +93,14 @@ async def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) - payload = kwargs["json"] assert isinstance(payload, dict) assert payload["model"] == "deepseek-v4-flash" + assert payload["messages"] == [ + { + "role": "system", + "content": "Answer only from the provided context. Cite sources with bracketed numbers like [1]. " + "If the context is insufficient, say so.", + }, + {"role": "user", "content": "Question:\nquestion\n\nContext:\n[1] Book\nsource"}, + ] async def test_embed_texts_uses_httpx_embeddings(mocker: MockerFixture) -> None: