114 lines
3.5 KiB
Python
114 lines
3.5 KiB
Python
import pytest
|
|
from types import SimpleNamespace
|
|
|
|
from genesis.writer.writer_agent import (
|
|
WriterAgent,
|
|
WRITER_PROMPT_TEMPLATE,
|
|
CHAPTER_OUTPUT_SCHEMA,
|
|
)
|
|
from genesis.writer.models import GenerationContext, ChapterSpec
|
|
from genesis.writer.writer_state import WriterState
|
|
from genesis.writer.exceptions import WriterGenerationError
|
|
|
|
|
|
class FakeEngine:
|
|
def __init__(self):
|
|
self.calls = 0
|
|
|
|
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
|
self.calls += 1
|
|
return SimpleNamespace(
|
|
data={"title": variables["title"], "blocks": [{"type": "paragraph", "text": "ok"}]},
|
|
status="ok",
|
|
)
|
|
|
|
|
|
class FakePromptRegistry:
|
|
@staticmethod
|
|
def get_or_create(name, template):
|
|
return SimpleNamespace(name=name, version="1", template=template)
|
|
|
|
|
|
def _ctx(cid, title):
|
|
return GenerationContext(
|
|
chapter_id=cid,
|
|
title=title,
|
|
template_marker=ChapterSpec(chapter_id=cid, title=title, section_placeholder=None),
|
|
structured_source="source text",
|
|
write_rules=["W1"],
|
|
design_rules=["D1"],
|
|
template_styles={"Heading1"},
|
|
prior_state=None,
|
|
)
|
|
|
|
|
|
def test_generate_chapter_success():
|
|
agent = WriterAgent(
|
|
session_id="s",
|
|
engine=FakeEngine(),
|
|
prompt_registry=FakePromptRegistry(),
|
|
state=WriterState(["db_design"]),
|
|
)
|
|
content = agent.generate_chapter(_ctx("db_design", "DB 設計"))
|
|
assert content.chapter_id == "db_design"
|
|
assert content.blocks and content.blocks[0].type == "paragraph"
|
|
|
|
|
|
def test_generate_chapter_retries():
|
|
class Boom(FakeEngine):
|
|
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
|
raise RuntimeError("boom")
|
|
|
|
agent = WriterAgent(
|
|
session_id="s",
|
|
engine=Boom(),
|
|
prompt_registry=FakePromptRegistry(),
|
|
state=WriterState(["db_design"]),
|
|
max_retries=2,
|
|
)
|
|
with pytest.raises(WriterGenerationError):
|
|
agent.generate_chapter(_ctx("db_design", "DB 設計"))
|
|
|
|
|
|
class AsyncFakeEngine:
|
|
async def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
|
return SimpleNamespace(
|
|
data={"title": variables["title"], "blocks": [{"type": "paragraph", "text": "ok"}]},
|
|
status="ok",
|
|
)
|
|
|
|
|
|
def test_generate_chapter_with_async_engine():
|
|
agent = WriterAgent(
|
|
session_id="s",
|
|
engine=AsyncFakeEngine(),
|
|
prompt_registry=FakePromptRegistry(),
|
|
state=WriterState(["db_design"]),
|
|
)
|
|
content = agent.generate_chapter(_ctx("db_design", "DB 設計"))
|
|
assert content.chapter_id == "db_design"
|
|
assert content.blocks and content.blocks[0].type == "paragraph"
|
|
|
|
|
|
def test_generate_chapter_propagates_engine_error_detail():
|
|
class FailingEngine:
|
|
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
|
return SimpleNamespace(
|
|
data={},
|
|
status="failed",
|
|
error="LLM HTTP 401: invalid key",
|
|
error_code="LLM_NETWORK_ERROR",
|
|
)
|
|
|
|
agent = WriterAgent(
|
|
session_id="s",
|
|
engine=FailingEngine(),
|
|
prompt_registry=FakePromptRegistry(),
|
|
state=WriterState(["db_design"]),
|
|
max_retries=1,
|
|
)
|
|
with pytest.raises(WriterGenerationError) as exc:
|
|
agent.generate_chapter(_ctx("db_design", "DB 設計"))
|
|
assert "LLM HTTP 401" in str(exc.value)
|
|
assert "LLM_NETWORK_ERROR" in str(exc.value)
|