Files
Stirling-PDF/engine/tests/test_stirling_contracts.py
James Brunton e10c5f6283 Redesign Python AI engine (#5991)
# Description of Changes
Redesign the Python AI engine to be properly agentic and make use of
`pydantic-ai` instead of `langchain` for correctness and ergonomics.
This should be a good foundation for us to build our AI engine on going
forwards.
2026-03-26 10:35:47 +00:00

79 lines
2.2 KiB
Python

from collections.abc import Iterator
import pytest
from stirling.config import AppSettings, load_settings
from stirling.contracts import (
AgentExecutionRequest,
AgentSpec,
AgentSpecStep,
EditPlanResponse,
ExecutionContext,
OrchestratorRequest,
PdfQuestionAnswerResponse,
ToolOperationStep,
)
from stirling.models.tool_models import OperationId, RotateParams
def test_orchestrator_request_accepts_user_message() -> None:
request = OrchestratorRequest(user_message="Rotate the PDF")
assert request.user_message == "Rotate the PDF"
def test_agent_execution_request_uses_typed_agent_spec() -> None:
steps: list[AgentSpecStep] = [
ToolOperationStep(
tool=OperationId.ROTATE,
parameters=RotateParams(angle=90),
)
]
request = AgentExecutionRequest(
agent_spec=AgentSpec(
name="Invoice cleanup",
description="Normalise inbound invoices",
objective="Prepare uploads for accounting review",
steps=steps,
),
current_step_index=0,
execution_context=ExecutionContext(input_files=["invoice.pdf"]),
)
assert request.agent_spec.steps[0].kind == "tool"
def test_edit_plan_response_has_typed_steps() -> None:
steps = [ToolOperationStep(tool=OperationId.ROTATE, parameters=RotateParams(angle=90))]
response = EditPlanResponse(
summary="Rotate the input PDF by 90 degrees.",
steps=steps,
)
assert response.steps[0].tool == OperationId.ROTATE
def test_pdf_question_answer_defaults_evidence_list() -> None:
response = PdfQuestionAnswerResponse(answer="The invoice total is 120.00")
assert response.evidence == []
@pytest.fixture(autouse=True)
def clear_settings_cache() -> Iterator[None]:
load_settings.cache_clear()
yield
load_settings.cache_clear()
def test_app_settings_accepts_model_configuration() -> None:
settings = AppSettings(
smart_model_name="claude-sonnet-4-5-20250929",
fast_model_name="claude-haiku-4-5-20251001",
smart_model_max_tokens=8192,
fast_model_max_tokens=2048,
)
assert settings.smart_model_name
assert settings.fast_model_max_tokens == 2048