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:
James Brunton
2026-04-14 18:45:47 +01:00
committed by GitHub
parent 4ada46ca56
commit 2bf5f0b18e
15 changed files with 397 additions and 112 deletions

View File

@@ -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(