mirror of
https://github.com/Frooodle/Stirling-PDF.git
synced 2026-04-06 03:19:39 +02:00
e10c5f6283
# 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.
80 lines
2.4 KiB
Python
80 lines
2.4 KiB
Python
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)
|