zotero-kb/tests/test_llm.py
2026-04-16 13:34:39 +08:00

79 lines
2.7 KiB
Python

import json
import pytest
from zotero_kb.llm import DeepSeekCardGenerationClient, create_card_generation_client
def test_create_card_generation_client_requires_deepseek_api_key() -> None:
with pytest.raises(ValueError, match="DEEPSEEK_API_KEY"):
create_card_generation_client("deepseek", "deepseek-chat", env={})
def test_deepseek_client_uses_env_api_key_and_parses_json() -> None:
captured = {}
def fake_request(url: str, payload: dict[str, object], headers: dict[str, str]) -> dict[str, object]:
captured["url"] = url
captured["payload"] = payload
captured["headers"] = headers
return {
"choices": [
{
"message": {
"content": json.dumps(
{
"summary": "DeepSeek summary",
"core_claims": ["Claim A"],
"methods": ["Method A"],
"evidence": ["Evidence A"],
"quotable_passages": ["Quote A"],
"writing_hints": ["Hint A"],
"keywords": ["kw-a"],
}
)
}
}
]
}
client = create_card_generation_client(
"deepseek",
"deepseek-chat",
env={"DEEPSEEK_API_KEY": "sk-test"},
request_fn=fake_request,
)
result = client.generate_card(
{
"title": "Scoped Retrieval for Drafting",
"abstract": "Merged metadata improves drafting.",
"notes": ["Project scoping helps."],
"attachment_texts": ["Attachment evidence."],
"tags": ["retrieval"],
}
)
assert result["summary"] == "DeepSeek summary"
assert captured["url"] == "https://api.deepseek.com/v1/chat/completions"
assert captured["payload"]["model"] == "deepseek-chat"
assert captured["headers"]["Authorization"] == "Bearer sk-test"
def test_deepseek_client_strips_markdown_fences() -> None:
client = DeepSeekCardGenerationClient(
api_key="sk-test",
request_fn=lambda *_args, **_kwargs: {
"choices": [
{
"message": {
"content": "```json\n{\"summary\":\"Fence summary\",\"core_claims\":[],\"methods\":[],\"evidence\":[],\"quotable_passages\":[],\"writing_hints\":[],\"keywords\":[]}\n```"
}
}
]
},
)
result = client.generate_card({"title": "Paper"})
assert result["summary"] == "Fence summary"