mirror of
https://github.com/Frooodle/Stirling-PDF.git
synced 2026-04-22 23:08:53 +02:00
Add tracking system to support optional PostHog tracking in AI engine (#6040)
Co-authored-by: ConnorYoh <40631091+ConnorYoh@users.noreply.github.com>
This commit is contained in:
@@ -3,7 +3,6 @@ from __future__ import annotations
|
||||
import pytest
|
||||
|
||||
from stirling.agents import PdfQuestionAgent
|
||||
from stirling.config import AppSettings
|
||||
from stirling.contracts import (
|
||||
ExtractedFileText,
|
||||
PdfQuestionAnswerResponse,
|
||||
@@ -12,12 +11,12 @@ from stirling.contracts import (
|
||||
PdfQuestionRequest,
|
||||
PdfTextSelection,
|
||||
)
|
||||
from stirling.services import build_runtime
|
||||
from stirling.services.runtime import AppRuntime
|
||||
|
||||
|
||||
class StubPdfQuestionAgent(PdfQuestionAgent):
|
||||
def __init__(self, response: PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse) -> None:
|
||||
super().__init__(build_runtime(build_test_settings()))
|
||||
def __init__(self, runtime: AppRuntime, response: PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse) -> None:
|
||||
super().__init__(runtime)
|
||||
self.response = response
|
||||
|
||||
async def _run_answer_agent(
|
||||
@@ -27,15 +26,6 @@ class StubPdfQuestionAgent(PdfQuestionAgent):
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
def invoice_page() -> ExtractedFileText:
|
||||
return ExtractedFileText(
|
||||
file_name="invoice.pdf",
|
||||
@@ -44,8 +34,8 @@ def invoice_page() -> ExtractedFileText:
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_question_agent_requires_extracted_text() -> None:
|
||||
agent = PdfQuestionAgent(build_runtime(build_test_settings()))
|
||||
async def test_pdf_question_agent_requires_extracted_text(runtime: AppRuntime) -> None:
|
||||
agent = PdfQuestionAgent(runtime)
|
||||
|
||||
response = await agent.handle(
|
||||
PdfQuestionRequest(question="What is the total?", page_text=[], file_names=["test.pdf"])
|
||||
@@ -55,12 +45,13 @@ async def test_pdf_question_agent_requires_extracted_text() -> None:
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_question_agent_returns_grounded_answer() -> None:
|
||||
async def test_pdf_question_agent_returns_grounded_answer(runtime: AppRuntime) -> None:
|
||||
agent = StubPdfQuestionAgent(
|
||||
runtime,
|
||||
PdfQuestionAnswerResponse(
|
||||
answer="The invoice total is 120.00.",
|
||||
evidence=[invoice_page()],
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
response = await agent.handle(
|
||||
@@ -76,8 +67,8 @@ async def test_pdf_question_agent_returns_grounded_answer() -> None:
|
||||
|
||||
|
||||
@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."))
|
||||
async def test_pdf_question_agent_returns_not_found_when_text_is_insufficient(runtime: AppRuntime) -> None:
|
||||
agent = StubPdfQuestionAgent(runtime, PdfQuestionNotFoundResponse(reason="The answer is not present in the text."))
|
||||
|
||||
response = await agent.handle(
|
||||
PdfQuestionRequest(
|
||||
|
||||
Reference in New Issue
Block a user