feat(api): add system prompt support for LLM completion
Separate prompt generation into system and user prompts for better LLM instruction following. Backend now builds a detailed system prompt with constraints for math formatting, code block handling, boundary newlines, and OCR safety, while user prompt contains context and completion state flags. Added corresponding tests for both modules.
This commit is contained in:
@@ -0,0 +1,65 @@
|
||||
import asyncio
|
||||
import importlib
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
BACKEND_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(BACKEND_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(BACKEND_DIR))
|
||||
|
||||
try:
|
||||
llm = importlib.import_module("llm")
|
||||
except ModuleNotFoundError:
|
||||
pytest.skip("llm module dependencies are not available", allow_module_level=True)
|
||||
|
||||
|
||||
def test_call_ollama_messages_roles_with_system(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
async def fake_chat(**kwargs):
|
||||
captured["messages"] = kwargs["messages"]
|
||||
return {"message": {"content": "ok", "thinking": ""}}
|
||||
|
||||
monkeypatch.setattr(llm.client, "chat", fake_chat)
|
||||
|
||||
result = asyncio.run(
|
||||
llm.call_ollama(
|
||||
"user prompt body",
|
||||
system_prompt="system prompt body",
|
||||
tag="test",
|
||||
temperature=0.1,
|
||||
)
|
||||
)
|
||||
|
||||
assert result["content"] == "ok"
|
||||
assert captured["messages"][0]["role"] == "system"
|
||||
assert captured["messages"][0]["content"] == "system prompt body"
|
||||
assert captured["messages"][1]["role"] == "user"
|
||||
assert captured["messages"][1]["content"] == "user prompt body"
|
||||
|
||||
|
||||
def test_call_ollama_messages_roles_without_system(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
async def fake_chat(**kwargs):
|
||||
captured["messages"] = kwargs["messages"]
|
||||
return {"message": {"content": "ok", "thinking": ""}}
|
||||
|
||||
monkeypatch.setattr(llm.client, "chat", fake_chat)
|
||||
|
||||
result = asyncio.run(
|
||||
llm.call_ollama(
|
||||
"user prompt only",
|
||||
system_prompt="",
|
||||
tag="test-no-system",
|
||||
temperature=0.1,
|
||||
)
|
||||
)
|
||||
|
||||
assert result["content"] == "ok"
|
||||
assert len(captured["messages"]) == 1
|
||||
assert captured["messages"][0]["role"] == "user"
|
||||
assert captured["messages"][0]["content"] == "user prompt only"
|
||||
@@ -0,0 +1,56 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
BACKEND_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(BACKEND_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(BACKEND_DIR))
|
||||
|
||||
import prompt # noqa: E402
|
||||
|
||||
|
||||
def test_prompt_builds_system_and_user():
|
||||
system_prompt, user_prompt = prompt.build_completion_prompts(
|
||||
prefix="The result is ",
|
||||
suffix="for this dataset.",
|
||||
language_id="markdown",
|
||||
)
|
||||
|
||||
assert "Hard constraints you must follow" in system_prompt
|
||||
assert "strict KaTeX-compatible math" in system_prompt
|
||||
assert "$...$" in system_prompt
|
||||
assert "$$...$$" in system_prompt
|
||||
assert "```{language}" in system_prompt
|
||||
assert "CURSOR_IN_FENCED_CODE_BLOCK" in user_prompt
|
||||
assert "PREFIX_ENDS_WITH_NEWLINE" in user_prompt
|
||||
assert "SUFFIX_STARTS_WITH_NEWLINE" in user_prompt
|
||||
|
||||
|
||||
def test_cursor_in_fence_detection():
|
||||
assert prompt._cursor_in_fenced_code_block("") is False
|
||||
assert prompt._cursor_in_fenced_code_block("```python\nprint('x')\n") is True
|
||||
assert prompt._cursor_in_fenced_code_block("```python\nprint('x')\n```\n") is False
|
||||
assert prompt._cursor_in_fenced_code_block("text ```not-a-fence``` tail") is False
|
||||
|
||||
|
||||
def test_newline_flags():
|
||||
_, user_prompt_a = prompt.build_completion_prompts(
|
||||
prefix="Hello",
|
||||
suffix="World",
|
||||
)
|
||||
assert "CURSOR_IN_FENCED_CODE_BLOCK: false" in user_prompt_a
|
||||
assert "PREFIX_ENDS_WITH_NEWLINE: false" in user_prompt_a
|
||||
assert "SUFFIX_STARTS_WITH_NEWLINE: false" in user_prompt_a
|
||||
|
||||
_, user_prompt_b = prompt.build_completion_prompts(
|
||||
prefix="Hello\n",
|
||||
suffix="\nWorld",
|
||||
)
|
||||
assert "PREFIX_ENDS_WITH_NEWLINE: true" in user_prompt_b
|
||||
assert "SUFFIX_STARTS_WITH_NEWLINE: true" in user_prompt_b
|
||||
|
||||
|
||||
def test_examples_coverage():
|
||||
_, user_prompt = prompt.build_completion_prompts(prefix="", suffix="")
|
||||
for ex in range(1, 13):
|
||||
assert f"[EX{ex:02d}]" in user_prompt
|
||||
Reference in New Issue
Block a user