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

@@ -4,7 +4,6 @@ import pytest
from pydantic import ValidationError
from stirling.agents import UserSpecAgent
from stirling.config import AppSettings
from stirling.contracts import (
AgentDraft,
AgentDraftRequest,
@@ -15,21 +14,12 @@ from stirling.contracts import (
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,
)
from stirling.services.runtime import AppRuntime
class StubUserSpecAgent(UserSpecAgent):
def __init__(self, draft_result: AgentDraft, revision_result: AgentDraft) -> None:
super().__init__(build_runtime(build_test_settings()))
def __init__(self, runtime: AppRuntime, draft_result: AgentDraft, revision_result: AgentDraft) -> None:
super().__init__(runtime)
self.draft_result = draft_result
self.revision_result = revision_result
self.edit_plan = EditPlanResponse(
@@ -53,8 +43,8 @@ class StubUserSpecAgent(UserSpecAgent):
class ClarifyingUserSpecAgent(UserSpecAgent):
def __init__(self) -> None:
super().__init__(build_runtime(build_test_settings()))
def __init__(self, runtime: AppRuntime) -> None:
super().__init__(runtime)
async def _build_edit_plan(self, user_message: str) -> EditClarificationRequest:
return EditClarificationRequest(
@@ -64,8 +54,9 @@ class ClarifyingUserSpecAgent(UserSpecAgent):
@pytest.mark.anyio
async def test_user_spec_agent_drafts_agent_spec() -> None:
async def test_user_spec_agent_drafts_agent_spec(runtime: AppRuntime) -> None:
agent = StubUserSpecAgent(
runtime,
AgentDraft(
name="Invoice Cleanup",
description="Prepare invoices for review.",
@@ -100,7 +91,7 @@ async def test_user_spec_agent_drafts_agent_spec() -> None:
@pytest.mark.anyio
async def test_user_spec_agent_revises_existing_draft() -> None:
async def test_user_spec_agent_revises_existing_draft(runtime: AppRuntime) -> None:
current_draft = AgentDraft(
name="Invoice Cleanup",
description="Prepare invoices for review.",
@@ -113,6 +104,7 @@ async def test_user_spec_agent_revises_existing_draft() -> None:
],
)
agent = StubUserSpecAgent(
runtime,
draft_result=current_draft,
revision_result=AgentDraft(
name="Invoice Cleanup",
@@ -152,8 +144,8 @@ def test_tool_operation_step_rejects_mismatched_parameters() -> None:
@pytest.mark.anyio
async def test_user_spec_agent_propagates_edit_clarification() -> None:
agent = ClarifyingUserSpecAgent()
async def test_user_spec_agent_propagates_edit_clarification(runtime: AppRuntime) -> None:
agent = ClarifyingUserSpecAgent(runtime)
response = await agent.draft(AgentDraftRequest(user_message="Build an agent to rotate some pages."))