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:
@@ -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."))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user