ai-contract-platform/backend/tests/test_prompt_runtime_service.py

import unittest
from datetime import datetime, timezone

from app.models.prompt_asset import PromptAsset
from app.models.prompt_asset_run_link import PromptAssetRunLink
from app.models.prompt_asset_version import PromptAssetVersion
from app.services.prompt_runtime_service import PromptRuntimeService


class StubPromptAssetRepository:
    def __init__(self, assets):
        self.assets = assets

    def get_published_by_key(self, tenant_id, prompt_key):
        for item in self.assets:
            if (
                item.tenant_id == tenant_id
                and item.prompt_key == prompt_key
                and item.status == "published"
                and item.published_version_id
            ):
                return item
        return None

    def list_published_by_type(self, tenant_id, prompt_type, limit=20):
        return [
            item
            for item in self.assets
            if item.tenant_id == tenant_id
            and item.prompt_type == prompt_type
            and item.status == "published"
            and item.published_version_id
        ][:limit]


class StubPromptAssetVersionRepository:
    def __init__(self, versions):
        self.versions = {item.id: item for item in versions}

    def get_by_id(self, version_id):
        return self.versions.get(version_id)


class StubPromptAssetRunLinkRepository:
    def __init__(self):
        self.items = []

    def create(self, run_link: PromptAssetRunLink):
        self.items.append(run_link)
        return run_link


class PromptRuntimeServiceTestCase(unittest.TestCase):
    def setUp(self):
        now = datetime.now(timezone.utc)
        self.asset = PromptAsset(
            id="prompt-001",
            tenant_id="tenant-001",
            prompt_key="legal_chat_default",
            prompt_type="legal_chat",
            name="Legal chat",
            status="published",
            current_version_label="v1.0",
            current_version_id="promptver-001",
            published_version_id="promptver-001",
            published_at=now,
            published_by="admin-001",
            version_count=1,
            created_at=now,
            updated_at=now,
        )
        self.version = PromptAssetVersion(
            id="promptver-001",
            tenant_id="tenant-001",
            prompt_asset_id="prompt-001",
            prompt_key="legal_chat_default",
            prompt_type="legal_chat",
            version_label="v1.0",
            prompt_content="Answer using contract context.",
            system_prompt="You are a contract assistant.",
            variables_json=[{"name": "message"}],
            output_schema_json={"type": "text"},
            model_config_json={"provider": "openai", "model": "gpt-5.4"},
            evaluation_status="passed",
            evaluation_score=0.91,
            is_current=True,
            is_published=True,
            created_at=now,
            published_at=now,
            published_by="admin-001",
        )
        self.service = PromptRuntimeService(None)
        self.service.repository = StubPromptAssetRepository([self.asset])
        self.service.version_repository = StubPromptAssetVersionRepository([self.version])
        self.service.run_link_repository = StubPromptAssetRunLinkRepository()
        self.current_user = {"id": "user-001", "tenant_id": "tenant-001", "role": "member"}

    def test_resolve_published_prompt_prefers_key_and_returns_runtime_trace(self):
        runtime = self.service.resolve_published_prompt(
            self.current_user,
            "legal_chat",
            preferred_keys=["legal_chat_review", "legal_chat_default"],
            fallback_prompt_version="legal_chat.shared_ai_stream.v1",
        )

        self.assertTrue(runtime["available"])
        self.assertEqual(runtime["prompt_key"], "legal_chat_default")
        self.assertEqual(runtime["prompt_version"], "legal_chat_default.v1.0")
        self.assertEqual(runtime["template_key"], "legal_chat_default")
        self.assertEqual(runtime["provider"], "openai")
        self.assertEqual(runtime["model"], "gpt-5.4")
        self.assertEqual(runtime["match_strategy"], "preferred_key")

    def test_resolve_published_prompt_falls_back_when_missing(self):
        runtime = self.service.resolve_published_prompt(
            {"id": "user-002", "tenant_id": "tenant-other"},
            "contract_review",
            preferred_keys=["contract_review_default"],
            fallback_prompt_version="legal-review.v1",
        )

        self.assertFalse(runtime["available"])
        self.assertEqual(runtime["prompt_version"], "legal-review.v1")
        self.assertEqual(runtime["source"], "local_default")

    def test_record_run_link_does_not_require_admin_user(self):
        runtime = self.service.resolve_published_prompt(
            self.current_user,
            "legal_chat",
            preferred_keys=["legal_chat_default"],
        )
        run_link = self.service.record_run_link(
            self.current_user,
            runtime,
            run_type="legal_chat",
            run_id="conv-001",
            resource_type="conversation",
            resource_id="conv-001",
            trace={"message_length": 12},
        )

        self.assertIsNotNone(run_link)
        self.assertEqual(self.service.run_link_repository.items[0].created_by, "user-001")
        self.assertEqual(self.service.run_link_repository.items[0].trace_json["prompt_runtime"]["prompt_key"], "legal_chat_default")


if __name__ == "__main__":
    unittest.main()