mirror of
https://github.com/Frooodle/Stirling-PDF.git
synced 2026-04-22 23:08:53 +02:00
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.
This commit is contained in:
James Brunton
GitHub
@@ -1,11 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
from pytest import Config
|
||||
|
||||
|
||||
def pytest_configure(config: Config) -> None:
|
||||
# Set required env vars in case there is no .env file
|
||||
os.environ.setdefault("STIRLING_OPENAI_API_KEY", "test")
|
||||
os.environ.setdefault("STIRLING_POSTHOG_API_KEY", "test")
|
||||
@@ -1,160 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pytest import MonkeyPatch
|
||||
|
||||
from chat_router import classify_chat_route
|
||||
from models import (
|
||||
ChatRouteRequest,
|
||||
ChatRouteResponse,
|
||||
CreateIntentHint,
|
||||
EditIntentHint,
|
||||
SmartFolderIntentHint,
|
||||
)
|
||||
|
||||
|
||||
def test_classify_chat_route_handles_smart_folder_intent(monkeypatch: MonkeyPatch):
|
||||
request = ChatRouteRequest(
|
||||
message="Create a workflow that batches PDFs overnight",
|
||||
history=[],
|
||||
has_files=False,
|
||||
has_create_session=False,
|
||||
has_edit_session=False,
|
||||
last_route="none",
|
||||
)
|
||||
expected_response = ChatRouteResponse(
|
||||
intent="smart_folder",
|
||||
smart_folder_intent=SmartFolderIntentHint(action="create"),
|
||||
reason="User wants to automate PDFs",
|
||||
)
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setattr("chat_router.run_ai", lambda *args, **kwargs: expected_response)
|
||||
response = classify_chat_route(request)
|
||||
|
||||
assert response.intent == "smart_folder"
|
||||
assert response.smart_folder_intent == expected_response.smart_folder_intent
|
||||
|
||||
|
||||
def test_greeting_without_files(monkeypatch: MonkeyPatch):
|
||||
"""Greetings should route to edit/info, not create"""
|
||||
request = ChatRouteRequest(
|
||||
message="Hello",
|
||||
has_files=False,
|
||||
has_create_session=False,
|
||||
has_edit_session=False,
|
||||
last_route="none",
|
||||
history=[],
|
||||
)
|
||||
expected_response = ChatRouteResponse(
|
||||
intent="edit",
|
||||
edit_intent=EditIntentHint(mode="info", requires_file_context=False),
|
||||
reason="Conversational greeting",
|
||||
)
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setattr("chat_router.run_ai", lambda *args, **kwargs: expected_response)
|
||||
response = classify_chat_route(request)
|
||||
|
||||
assert response.intent == "edit"
|
||||
assert response.edit_intent is not None
|
||||
assert response.edit_intent.mode == "info"
|
||||
|
||||
|
||||
def test_capability_question_without_files(monkeypatch: MonkeyPatch):
|
||||
"""'What can you do?' should route to edit/info"""
|
||||
request = ChatRouteRequest(
|
||||
message="What can you do?",
|
||||
has_files=False,
|
||||
has_create_session=False,
|
||||
has_edit_session=False,
|
||||
last_route="none",
|
||||
history=[],
|
||||
)
|
||||
expected_response = ChatRouteResponse(
|
||||
intent="edit",
|
||||
edit_intent=EditIntentHint(mode="info", requires_file_context=False),
|
||||
reason="User asking about capabilities",
|
||||
)
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setattr("chat_router.run_ai", lambda *args, **kwargs: expected_response)
|
||||
response = classify_chat_route(request)
|
||||
|
||||
assert response.intent == "edit"
|
||||
assert response.edit_intent is not None
|
||||
assert response.edit_intent.mode == "info"
|
||||
|
||||
|
||||
def test_help_request(monkeypatch: MonkeyPatch):
|
||||
"""Help requests should route to edit/info"""
|
||||
request = ChatRouteRequest(
|
||||
message="help",
|
||||
has_files=False,
|
||||
has_create_session=False,
|
||||
has_edit_session=False,
|
||||
last_route="none",
|
||||
history=[],
|
||||
)
|
||||
expected_response = ChatRouteResponse(
|
||||
intent="edit",
|
||||
edit_intent=EditIntentHint(mode="info", requires_file_context=False),
|
||||
reason="Help request",
|
||||
)
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setattr("chat_router.run_ai", lambda *args, **kwargs: expected_response)
|
||||
response = classify_chat_route(request)
|
||||
|
||||
assert response.intent == "edit"
|
||||
assert response.edit_intent is not None
|
||||
assert response.edit_intent.mode == "info"
|
||||
|
||||
|
||||
def test_actual_document_creation(monkeypatch: MonkeyPatch):
|
||||
"""Explicit creation requests should still route to create"""
|
||||
request = ChatRouteRequest(
|
||||
message="Create a business proposal document",
|
||||
has_files=False,
|
||||
has_create_session=False,
|
||||
has_edit_session=False,
|
||||
last_route="none",
|
||||
history=[],
|
||||
)
|
||||
expected_response = ChatRouteResponse(
|
||||
intent="create",
|
||||
create_intent=CreateIntentHint(action="start"),
|
||||
reason="User wants to create new document",
|
||||
)
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setattr("chat_router.run_ai", lambda *args, **kwargs: expected_response)
|
||||
response = classify_chat_route(request)
|
||||
|
||||
assert response.intent == "create"
|
||||
assert response.create_intent is not None
|
||||
assert response.create_intent.action == "start"
|
||||
|
||||
|
||||
def test_edit_command_with_files(monkeypatch: MonkeyPatch):
|
||||
"""Edit commands should still route to edit/command"""
|
||||
request = ChatRouteRequest(
|
||||
message="Compress this PDF",
|
||||
has_files=True,
|
||||
has_create_session=False,
|
||||
has_edit_session=False,
|
||||
last_route="none",
|
||||
history=[],
|
||||
)
|
||||
expected_response = ChatRouteResponse(
|
||||
intent="edit",
|
||||
edit_intent=EditIntentHint(mode="command", requires_file_context=False),
|
||||
reason="User wants to compress PDF",
|
||||
)
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setattr("chat_router.run_ai", lambda *args, **kwargs: expected_response)
|
||||
response = classify_chat_route(request)
|
||||
|
||||
assert response.intent == "edit"
|
||||
assert response.edit_intent is not None
|
||||
assert response.edit_intent.mode == "command"
|
||||
@@ -1,52 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pytest import MonkeyPatch
|
||||
|
||||
from editing.decisions import answer_conversational_info
|
||||
from file_processing_agent import ToolCatalogService
|
||||
|
||||
|
||||
def test_answer_conversational_info_greeting(monkeypatch: MonkeyPatch):
|
||||
"""Test handling of greeting without files"""
|
||||
|
||||
class MockResponse:
|
||||
message = "Hello! I can help you with PDF operations like compress, merge, split, and more."
|
||||
|
||||
def mock_run_ai(*args, **kwargs):
|
||||
return MockResponse()
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setattr("editing.decisions.run_ai", mock_run_ai)
|
||||
|
||||
tool_catalog = ToolCatalogService()
|
||||
result = answer_conversational_info(
|
||||
message="Hello",
|
||||
history=[],
|
||||
tool_catalog=tool_catalog,
|
||||
)
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
|
||||
|
||||
def test_answer_conversational_info_capabilities(monkeypatch: MonkeyPatch):
|
||||
"""Test handling of capability questions"""
|
||||
|
||||
class MockResponse:
|
||||
message = "I can help with compress, merge, split, rotate, watermark, OCR, and many other PDF operations."
|
||||
|
||||
def mock_run_ai(*args, **kwargs):
|
||||
return MockResponse()
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setattr("editing.decisions.run_ai", mock_run_ai)
|
||||
|
||||
tool_catalog = ToolCatalogService()
|
||||
result = answer_conversational_info(
|
||||
message="What can you do?",
|
||||
history=[],
|
||||
tool_catalog=tool_catalog,
|
||||
)
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
@@ -1,31 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import dotenv_values
|
||||
|
||||
ENGINE_ROOT = Path(__file__).parent.parent
|
||||
SRC_DIR = ENGINE_ROOT / "src"
|
||||
EXAMPLE_FILE = ENGINE_ROOT / "config" / ".env.example"
|
||||
|
||||
|
||||
def _parse_example_keys() -> set[str]:
|
||||
return set(dotenv_values(EXAMPLE_FILE).keys())
|
||||
|
||||
|
||||
def _find_stirling_env_vars() -> set[str]:
|
||||
env_vars: set[str] = set()
|
||||
for path in SRC_DIR.rglob("*.py"):
|
||||
for match in re.finditer(r"\b(STIRLING_\w+)\b", path.read_text()):
|
||||
env_vars.add(match.group(1))
|
||||
return env_vars
|
||||
|
||||
|
||||
def test_every_stirling_env_var_is_in_example_file():
|
||||
example_keys = _parse_example_keys()
|
||||
source_vars = _find_stirling_env_vars()
|
||||
missing = sorted(source_vars - example_keys)
|
||||
assert not missing, "env vars used in src/ but missing from config/.env.example:\n" + "\n".join(
|
||||
f" {v}" for v in missing
|
||||
)
|
||||
160
engine/tests/test_pdf_edit_agent.py
Normal file
160
engine/tests/test_pdf_edit_agent.py
Normal file
@@ -0,0 +1,160 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
|
||||
from stirling.agents import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection
|
||||
from stirling.config import AppSettings
|
||||
from stirling.contracts import (
|
||||
EditCannotDoResponse,
|
||||
EditClarificationRequest,
|
||||
EditPlanResponse,
|
||||
PdfEditRequest,
|
||||
ToolOperationStep,
|
||||
)
|
||||
from stirling.models.tool_models import CompressParams, OperationId, RotateParams
|
||||
from stirling.services import build_runtime
|
||||
|
||||
|
||||
def build_test_settings() -> AppSettings:
|
||||
return AppSettings(
|
||||
smart_model_name="test",
|
||||
fast_model_name="test",
|
||||
smart_model_max_tokens=8192,
|
||||
fast_model_max_tokens=2048,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ParameterSelectorCall:
|
||||
request: PdfEditRequest
|
||||
operation_plan: list[OperationId]
|
||||
operation_index: int
|
||||
generated_steps: list[ToolOperationStep]
|
||||
|
||||
|
||||
class RecordingParameterSelector:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[ParameterSelectorCall] = []
|
||||
|
||||
async def select(
|
||||
self,
|
||||
request: PdfEditRequest,
|
||||
operation_plan: list[OperationId],
|
||||
operation_index: int,
|
||||
generated_steps: list[ToolOperationStep],
|
||||
) -> RotateParams | CompressParams:
|
||||
self.calls.append(
|
||||
ParameterSelectorCall(
|
||||
request=request,
|
||||
operation_plan=operation_plan,
|
||||
operation_index=operation_index,
|
||||
generated_steps=list(generated_steps),
|
||||
)
|
||||
)
|
||||
if operation_index == 0:
|
||||
return RotateParams(angle=90)
|
||||
return CompressParams(compression_level=5)
|
||||
|
||||
|
||||
class StubPdfEditAgent(PdfEditAgent):
|
||||
def __init__(
|
||||
self,
|
||||
selection: PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse,
|
||||
parameter_selector: RecordingParameterSelector | PdfEditParameterSelector | None = None,
|
||||
) -> None:
|
||||
super().__init__(build_runtime(build_test_settings()))
|
||||
self.selection = selection
|
||||
if parameter_selector is not None:
|
||||
self.parameter_selector = parameter_selector
|
||||
|
||||
async def _select_plan(
|
||||
self,
|
||||
request: PdfEditRequest,
|
||||
) -> PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse:
|
||||
return self.selection
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_edit_agent_builds_multi_step_plan() -> None:
|
||||
parameter_selector = RecordingParameterSelector()
|
||||
agent = StubPdfEditAgent(
|
||||
PdfEditPlanSelection(
|
||||
operations=[OperationId.ROTATE, OperationId.COMPRESS],
|
||||
summary="Rotate the PDF, then compress it.",
|
||||
rationale="The pages need reorientation before reducing file size.",
|
||||
),
|
||||
parameter_selector=parameter_selector,
|
||||
)
|
||||
|
||||
response = await agent.handle(
|
||||
PdfEditRequest(
|
||||
user_message="Rotate the PDF clockwise and then compress it.",
|
||||
file_names=["scan.pdf"],
|
||||
)
|
||||
)
|
||||
|
||||
assert isinstance(response, EditPlanResponse)
|
||||
assert response.summary == "Rotate the PDF, then compress it."
|
||||
assert response.rationale == "The pages need reorientation before reducing file size."
|
||||
assert [step.tool for step in response.steps] == [OperationId.ROTATE, OperationId.COMPRESS]
|
||||
assert isinstance(response.steps[0].parameters, RotateParams)
|
||||
assert isinstance(response.steps[1].parameters, CompressParams)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_edit_agent_passes_previous_steps_to_parameter_selector() -> None:
|
||||
parameter_selector = RecordingParameterSelector()
|
||||
agent = StubPdfEditAgent(
|
||||
PdfEditPlanSelection(
|
||||
operations=[OperationId.ROTATE, OperationId.COMPRESS],
|
||||
summary="Rotate the PDF, then compress it.",
|
||||
),
|
||||
parameter_selector=parameter_selector,
|
||||
)
|
||||
|
||||
request = PdfEditRequest(
|
||||
user_message="Rotate the PDF clockwise and then compress it.",
|
||||
file_names=["scan.pdf"],
|
||||
)
|
||||
response = await agent.handle(request)
|
||||
|
||||
assert isinstance(response, EditPlanResponse)
|
||||
assert len(parameter_selector.calls) == 2
|
||||
assert parameter_selector.calls[0].operation_index == 0
|
||||
assert parameter_selector.calls[0].generated_steps == []
|
||||
assert parameter_selector.calls[1].operation_index == 1
|
||||
assert parameter_selector.calls[1].generated_steps == [
|
||||
ToolOperationStep(
|
||||
tool=OperationId.ROTATE,
|
||||
parameters=RotateParams(angle=90),
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_edit_agent_returns_clarification_without_partial_plan() -> None:
|
||||
agent = StubPdfEditAgent(
|
||||
EditClarificationRequest(
|
||||
question="Which pages should be rotated?",
|
||||
reason="The request does not say which pages to change.",
|
||||
)
|
||||
)
|
||||
|
||||
response = await agent.handle(PdfEditRequest(user_message="Rotate some pages."))
|
||||
|
||||
assert isinstance(response, EditClarificationRequest)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_edit_agent_returns_cannot_do_without_partial_plan() -> None:
|
||||
agent = StubPdfEditAgent(
|
||||
EditCannotDoResponse(
|
||||
reason="This request requires OCR, which is not part of PDF edit planning.",
|
||||
)
|
||||
)
|
||||
|
||||
response = await agent.handle(PdfEditRequest(user_message="Read this scan and summarize it."))
|
||||
|
||||
assert isinstance(response, EditCannotDoResponse)
|
||||
79
engine/tests/test_pdf_question_agent.py
Normal file
79
engine/tests/test_pdf_question_agent.py
Normal file
@@ -0,0 +1,79 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from stirling.agents import PdfQuestionAgent
|
||||
from stirling.config import AppSettings
|
||||
from stirling.contracts import (
|
||||
PdfQuestionAnswerResponse,
|
||||
PdfQuestionNeedTextResponse,
|
||||
PdfQuestionNotFoundResponse,
|
||||
PdfQuestionRequest,
|
||||
)
|
||||
from stirling.services import build_runtime
|
||||
|
||||
|
||||
class StubPdfQuestionAgent(PdfQuestionAgent):
|
||||
def __init__(self, response: PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse) -> None:
|
||||
super().__init__(build_runtime(build_test_settings()))
|
||||
self.response = response
|
||||
|
||||
async def _run_answer_agent(
|
||||
self,
|
||||
request: PdfQuestionRequest,
|
||||
) -> PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse:
|
||||
return self.response
|
||||
|
||||
|
||||
def build_test_settings() -> AppSettings:
|
||||
return AppSettings(
|
||||
smart_model_name="test",
|
||||
fast_model_name="test",
|
||||
smart_model_max_tokens=8192,
|
||||
fast_model_max_tokens=2048,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_question_agent_requires_extracted_text() -> None:
|
||||
agent = PdfQuestionAgent(build_runtime(build_test_settings()))
|
||||
|
||||
response = await agent.handle(PdfQuestionRequest(question="What is the total?", extracted_text=""))
|
||||
|
||||
assert isinstance(response, PdfQuestionNeedTextResponse)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_question_agent_returns_grounded_answer() -> None:
|
||||
agent = StubPdfQuestionAgent(
|
||||
PdfQuestionAnswerResponse(
|
||||
answer="The invoice total is 120.00.",
|
||||
evidence=["Invoice total: 120.00"],
|
||||
)
|
||||
)
|
||||
|
||||
response = await agent.handle(
|
||||
PdfQuestionRequest(
|
||||
question="What is the total?",
|
||||
extracted_text="Invoice total: 120.00",
|
||||
file_name="invoice.pdf",
|
||||
)
|
||||
)
|
||||
|
||||
assert isinstance(response, PdfQuestionAnswerResponse)
|
||||
assert response.answer == "The invoice total is 120.00."
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_question_agent_returns_not_found_when_text_is_insufficient() -> None:
|
||||
agent = StubPdfQuestionAgent(PdfQuestionNotFoundResponse(reason="The answer is not present in the text."))
|
||||
|
||||
response = await agent.handle(
|
||||
PdfQuestionRequest(
|
||||
question="What is the total?",
|
||||
extracted_text="This page contains only a shipping address.",
|
||||
file_name="invoice.pdf",
|
||||
)
|
||||
)
|
||||
|
||||
assert isinstance(response, PdfQuestionNotFoundResponse)
|
||||
@@ -1,45 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pytest import MonkeyPatch
|
||||
|
||||
from models import (
|
||||
AvailableTool,
|
||||
SmartFolderAutomation,
|
||||
SmartFolderConfig,
|
||||
SmartFolderCreateRequest,
|
||||
SmartFolderCreateResponse,
|
||||
SmartFolderOperation,
|
||||
)
|
||||
from smart_folder_creator import create_smart_folder_config
|
||||
|
||||
|
||||
def _build_sample_response() -> SmartFolderCreateResponse:
|
||||
return SmartFolderCreateResponse(
|
||||
assistant_message="I will build that folder for you.",
|
||||
smart_folder_config=SmartFolderConfig(
|
||||
name="Email Prep",
|
||||
description="Compress and split for email",
|
||||
automation=SmartFolderAutomation(
|
||||
name="Email Cleanup",
|
||||
description="Email prep steps",
|
||||
operations=[SmartFolderOperation(operation="compress-pdf", parameters='{"compressionLevel": 3}')],
|
||||
),
|
||||
icon="mail",
|
||||
accent_color="#0ea5e9",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_create_smart_folder_config_calls_ai_and_returns_response(monkeypatch: MonkeyPatch):
|
||||
request = SmartFolderCreateRequest(
|
||||
message="Create a folder that zips attachments",
|
||||
history=[],
|
||||
available_tools=[AvailableTool(id="compress-pdf", name="Compress PDFs")],
|
||||
)
|
||||
response_value = _build_sample_response()
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setattr("smart_folder_creator.run_ai", lambda *args, **kwargs: response_value)
|
||||
result = create_smart_folder_config(request)
|
||||
|
||||
assert result == response_value
|
||||
191
engine/tests/test_stirling_api.py
Normal file
191
engine/tests/test_stirling_api.py
Normal file
@@ -0,0 +1,191 @@
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from stirling.api import app
|
||||
from stirling.api.dependencies import (
|
||||
get_execution_planning_agent,
|
||||
get_orchestrator_agent,
|
||||
get_pdf_edit_agent,
|
||||
get_pdf_question_agent,
|
||||
get_user_spec_agent,
|
||||
)
|
||||
from stirling.config import AppSettings, load_settings
|
||||
from stirling.contracts import (
|
||||
AgentDraft,
|
||||
AgentDraftRequest,
|
||||
AgentDraftResponse,
|
||||
AgentExecutionRequest,
|
||||
AgentRevisionRequest,
|
||||
AgentRevisionResponse,
|
||||
CannotContinueExecutionAction,
|
||||
EditCannotDoResponse,
|
||||
OrchestratorRequest,
|
||||
PdfEditRequest,
|
||||
PdfQuestionNotFoundResponse,
|
||||
PdfQuestionRequest,
|
||||
UnsupportedCapabilityResponse,
|
||||
)
|
||||
from stirling.models.tool_models import RotateParams
|
||||
|
||||
|
||||
class StubSettingsProvider:
|
||||
def __call__(self) -> AppSettings:
|
||||
return AppSettings(
|
||||
smart_model_name="test",
|
||||
fast_model_name="test",
|
||||
smart_model_max_tokens=8192,
|
||||
fast_model_max_tokens=2048,
|
||||
)
|
||||
|
||||
|
||||
class StubOrchestratorAgent:
|
||||
async def handle(self, request: OrchestratorRequest) -> UnsupportedCapabilityResponse:
|
||||
return UnsupportedCapabilityResponse(capability="pdf_edit", message=request.user_message)
|
||||
|
||||
|
||||
class StubPdfEditAgent:
|
||||
async def handle(self, request: PdfEditRequest) -> EditCannotDoResponse:
|
||||
return EditCannotDoResponse(reason=request.user_message)
|
||||
|
||||
|
||||
class StubPdfQuestionAgent:
|
||||
async def handle(self, request: PdfQuestionRequest) -> PdfQuestionNotFoundResponse:
|
||||
return PdfQuestionNotFoundResponse(reason=request.question)
|
||||
|
||||
|
||||
class StubUserSpecAgent:
|
||||
async def draft(self, request: AgentDraftRequest) -> AgentDraftResponse:
|
||||
return AgentDraftResponse(
|
||||
draft=AgentDraft(
|
||||
name="Drafted",
|
||||
description="Route wiring test",
|
||||
objective=request.user_message,
|
||||
steps=[],
|
||||
)
|
||||
)
|
||||
|
||||
async def revise(self, request: AgentRevisionRequest) -> AgentRevisionResponse:
|
||||
return AgentRevisionResponse(draft=request.current_draft)
|
||||
|
||||
|
||||
class StubExecutionPlanningAgent:
|
||||
async def next_action(self, request: AgentExecutionRequest) -> CannotContinueExecutionAction:
|
||||
return CannotContinueExecutionAction(reason=str(request.current_step_index))
|
||||
|
||||
|
||||
client: TestClient = TestClient(app)
|
||||
|
||||
|
||||
def override_settings() -> AppSettings:
|
||||
return StubSettingsProvider()()
|
||||
|
||||
|
||||
def override_orchestrator_agent() -> StubOrchestratorAgent:
|
||||
return StubOrchestratorAgent()
|
||||
|
||||
|
||||
def override_pdf_edit_agent() -> StubPdfEditAgent:
|
||||
return StubPdfEditAgent()
|
||||
|
||||
|
||||
def override_pdf_question_agent() -> StubPdfQuestionAgent:
|
||||
return StubPdfQuestionAgent()
|
||||
|
||||
|
||||
def override_user_spec_agent() -> StubUserSpecAgent:
|
||||
return StubUserSpecAgent()
|
||||
|
||||
|
||||
def override_execution_agent() -> StubExecutionPlanningAgent:
|
||||
return StubExecutionPlanningAgent()
|
||||
|
||||
|
||||
app.dependency_overrides[load_settings] = override_settings
|
||||
app.dependency_overrides[get_orchestrator_agent] = override_orchestrator_agent
|
||||
app.dependency_overrides[get_pdf_edit_agent] = override_pdf_edit_agent
|
||||
app.dependency_overrides[get_pdf_question_agent] = override_pdf_question_agent
|
||||
app.dependency_overrides[get_user_spec_agent] = override_user_spec_agent
|
||||
app.dependency_overrides[get_execution_planning_agent] = override_execution_agent
|
||||
|
||||
|
||||
def test_health_route() -> None:
|
||||
response = client.get("/health")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["status"] == "ok"
|
||||
|
||||
|
||||
def test_orchestrator_route() -> None:
|
||||
response = client.post("/api/v1/orchestrator", json={"userMessage": "route this"})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["outcome"] == "unsupported_capability"
|
||||
|
||||
|
||||
def test_pdf_edit_route() -> None:
|
||||
response = client.post("/api/v1/pdf/edit", json={"userMessage": "rotate this"})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["outcome"] == "cannot_do"
|
||||
|
||||
|
||||
def test_pdf_questions_route() -> None:
|
||||
response = client.post("/api/v1/pdf/questions", json={"question": "what is this?"})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["outcome"] == "not_found"
|
||||
|
||||
|
||||
def test_agent_draft_route() -> None:
|
||||
response = client.post("/api/v1/agents/draft", json={"userMessage": "build me an agent"})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["outcome"] == "draft"
|
||||
|
||||
|
||||
def test_agent_revise_route() -> None:
|
||||
response = client.post(
|
||||
"/api/v1/agents/revise",
|
||||
json={
|
||||
"userMessage": "revise it",
|
||||
"currentDraft": {
|
||||
"name": "Drafted",
|
||||
"description": "Route wiring test",
|
||||
"objective": "build me an agent",
|
||||
"steps": [
|
||||
{
|
||||
"kind": "tool",
|
||||
"tool": "rotate",
|
||||
"parameters": RotateParams(angle=90).model_dump(by_alias=True),
|
||||
}
|
||||
],
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["outcome"] == "draft"
|
||||
|
||||
|
||||
def test_next_action_route() -> None:
|
||||
response = client.post(
|
||||
"/api/v1/agents/next-action",
|
||||
json={
|
||||
"agentSpec": {
|
||||
"name": "Drafted",
|
||||
"description": "Route wiring test",
|
||||
"objective": "build me an agent",
|
||||
"steps": [
|
||||
{
|
||||
"kind": "tool",
|
||||
"tool": "rotate",
|
||||
"parameters": RotateParams(angle=90).model_dump(by_alias=True),
|
||||
}
|
||||
],
|
||||
},
|
||||
"currentStepIndex": 0,
|
||||
"executionContext": {"inputFiles": ["input.pdf"], "metadata": {}},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["outcome"] == "cannot_continue"
|
||||
78
engine/tests/test_stirling_contracts.py
Normal file
78
engine/tests/test_stirling_contracts.py
Normal file
@@ -0,0 +1,78 @@
|
||||
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
|
||||
160
engine/tests/test_user_spec_agent.py
Normal file
160
engine/tests/test_user_spec_agent.py
Normal file
@@ -0,0 +1,160 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from stirling.agents import UserSpecAgent
|
||||
from stirling.config import AppSettings
|
||||
from stirling.contracts import (
|
||||
AgentDraft,
|
||||
AgentDraftRequest,
|
||||
AgentRevisionRequest,
|
||||
ConversationMessage,
|
||||
EditClarificationRequest,
|
||||
EditPlanResponse,
|
||||
ToolOperationStep,
|
||||
)
|
||||
from stirling.models.tool_models import CompressParams, OperationId, RotateParams
|
||||
from stirling.services import build_runtime
|
||||
|
||||
|
||||
def build_test_settings() -> AppSettings:
|
||||
return AppSettings(
|
||||
smart_model_name="test",
|
||||
fast_model_name="test",
|
||||
smart_model_max_tokens=8192,
|
||||
fast_model_max_tokens=2048,
|
||||
)
|
||||
|
||||
|
||||
class StubUserSpecAgent(UserSpecAgent):
|
||||
def __init__(self, draft_result: AgentDraft, revision_result: AgentDraft) -> None:
|
||||
super().__init__(build_runtime(build_test_settings()))
|
||||
self.draft_result = draft_result
|
||||
self.revision_result = revision_result
|
||||
self.edit_plan = EditPlanResponse(
|
||||
summary="Rotate the document.",
|
||||
steps=[
|
||||
ToolOperationStep(
|
||||
tool=OperationId.ROTATE,
|
||||
parameters=RotateParams(angle=90),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
async def _build_edit_plan(self, user_message: str) -> EditPlanResponse:
|
||||
return self.edit_plan
|
||||
|
||||
async def _run_draft_agent(self, request: AgentDraftRequest, edit_plan: EditPlanResponse) -> AgentDraft:
|
||||
return self.draft_result
|
||||
|
||||
async def _run_revision_agent(self, request: AgentRevisionRequest, edit_plan: EditPlanResponse) -> AgentDraft:
|
||||
return self.revision_result
|
||||
|
||||
|
||||
class ClarifyingUserSpecAgent(UserSpecAgent):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(build_runtime(build_test_settings()))
|
||||
|
||||
async def _build_edit_plan(self, user_message: str) -> EditClarificationRequest:
|
||||
return EditClarificationRequest(
|
||||
question="Which pages should be changed?",
|
||||
reason="The request does not specify the target pages.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_user_spec_agent_drafts_agent_spec() -> None:
|
||||
agent = StubUserSpecAgent(
|
||||
AgentDraft(
|
||||
name="Invoice Cleanup",
|
||||
description="Prepare invoices for review.",
|
||||
objective="Normalize invoices before accounting review.",
|
||||
steps=[
|
||||
ToolOperationStep(
|
||||
tool=OperationId.ROTATE,
|
||||
parameters=RotateParams(angle=90),
|
||||
)
|
||||
],
|
||||
),
|
||||
revision_result=AgentDraft(
|
||||
name="Unused",
|
||||
description="Unused",
|
||||
objective="Unused",
|
||||
steps=[],
|
||||
),
|
||||
)
|
||||
|
||||
response = await agent.draft(
|
||||
AgentDraftRequest(
|
||||
user_message="Build me an invoice cleanup agent.",
|
||||
conversation_history=[
|
||||
ConversationMessage(role="user", content="It should handle scanned PDFs."),
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
assert response.outcome == "draft"
|
||||
assert response.draft.name == "Invoice Cleanup"
|
||||
assert response.draft.steps[0].kind == "tool"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_user_spec_agent_revises_existing_draft() -> None:
|
||||
current_draft = AgentDraft(
|
||||
name="Invoice Cleanup",
|
||||
description="Prepare invoices for review.",
|
||||
objective="Normalize invoices before accounting review.",
|
||||
steps=[
|
||||
ToolOperationStep(
|
||||
tool=OperationId.ROTATE,
|
||||
parameters=RotateParams(angle=90),
|
||||
)
|
||||
],
|
||||
)
|
||||
agent = StubUserSpecAgent(
|
||||
draft_result=current_draft,
|
||||
revision_result=AgentDraft(
|
||||
name="Invoice Cleanup",
|
||||
description="Prepare invoices for review and reduce file size.",
|
||||
objective="Normalize invoices before accounting review.",
|
||||
steps=[
|
||||
ToolOperationStep(
|
||||
tool=OperationId.ROTATE,
|
||||
parameters=RotateParams(angle=90),
|
||||
),
|
||||
ToolOperationStep(
|
||||
tool=OperationId.COMPRESS,
|
||||
parameters=CompressParams(compression_level=5),
|
||||
),
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
response = await agent.revise(
|
||||
AgentRevisionRequest(
|
||||
user_message="Also compress the files before upload.",
|
||||
current_draft=current_draft,
|
||||
)
|
||||
)
|
||||
|
||||
assert response.outcome == "draft"
|
||||
assert len(response.draft.steps) == 2
|
||||
assert response.draft.steps[1].kind == "tool"
|
||||
|
||||
|
||||
def test_tool_operation_step_rejects_mismatched_parameters() -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
ToolOperationStep(
|
||||
tool=OperationId.ROTATE,
|
||||
parameters=CompressParams(compression_level=5),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_user_spec_agent_propagates_edit_clarification() -> None:
|
||||
agent = ClarifyingUserSpecAgent()
|
||||
|
||||
response = await agent.draft(AgentDraftRequest(user_message="Build an agent to rotate some pages."))
|
||||
|
||||
assert isinstance(response, EditClarificationRequest)
|
||||
Reference in New Issue
Block a user