ydy0615
e28125079c
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.
66 lines
1.8 KiB
Python
66 lines
1.8 KiB
Python
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"
|