Joseph Doherty
25 lines
767 B
Python
25 lines
767 B
Python
import pytest
|
|
from pydantic import BaseModel
|
|
from chat.llm.mock import MockLLMClient
|
|
from chat.llm.classify import classify
|
|
|
|
|
|
class Verdict(BaseModel):
|
|
score: int
|
|
reason: str
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_classify_parses_valid_json():
|
|
mock = MockLLMClient(canned=['{"score": 2, "reason": "notable"}'])
|
|
result = await classify(mock, model="m", system="x", user="y", schema=Verdict)
|
|
assert result.score == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_classify_falls_back_on_unparseable_after_retry():
|
|
mock = MockLLMClient(canned=["nope", "still nope"])
|
|
default = Verdict(score=1, reason="fallback")
|
|
result = await classify(mock, model="m", system="x", user="y", schema=Verdict, default=default)
|
|
assert result.reason == "fallback"
|