diff --git a/tests/test_detect_created_schema.py b/tests/test_detect_created_schema.py new file mode 100644 index 000000000..f3df6b8ca --- /dev/null +++ b/tests/test_detect_created_schema.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright (c) 2026 sol pbc + +import importlib +import json +from pathlib import Path + +from jsonschema import Draft202012Validator + +import think.models as models + +detect_created_mod = importlib.import_module("think.detect_created") + +DETECT_CREATED_SCHEMA_PATH = ( + Path(__file__).resolve().parents[1] / "think" / "detect_created.schema.json" +) + + +def _load_detect_created_schema() -> dict: + return json.loads(DETECT_CREATED_SCHEMA_PATH.read_text(encoding="utf-8")) + + +def test_detect_created_schema_file_is_valid_draft_2020_12(): + Draft202012Validator.check_schema(_load_detect_created_schema()) + + +def test_detect_created_schema_accepts_and_rejects_expected_values(): + schema = _load_detect_created_schema() + validator = Draft202012Validator(schema) + valid = { + "day": "20240315", + "time": "143052", + "confidence": "high", + "source": "QuickTime:CreateDate", + "utc": True, + } + + assert validator.is_valid(valid) + assert not validator.is_valid( + { + "day": "20240315", + "time": "143052", + "confidence": "high", + "source": "QuickTime:CreateDate", + } + ) + assert not validator.is_valid({**valid, "day": "2024-03-15"}) + assert not validator.is_valid({**valid, "time": "14:30:52"}) + assert not validator.is_valid({**valid, "confidence": "certain"}) + assert not validator.is_valid({**valid, "extra": "x"}) + assert not validator.is_valid({**valid, "source": ""}) + + +def test_detect_created_passes_schema_to_generate(monkeypatch): + captured = {} + + def fake_generate(**kwargs): + captured.update(kwargs) + return ( + '{"day": "20240315", "time": "143052", "confidence": "high", ' + '"source": "QuickTime:CreateDate", "utc": false}' + ) + + monkeypatch.setattr(models, "generate", fake_generate) + monkeypatch.setattr( + detect_created_mod, + "_extract_metadata", + lambda path: "QuickTime Create Date : 2024:03:15 14:30:52", + ) + + result = detect_created_mod.detect_created("/dev/null") + + assert captured["json_schema"] is detect_created_mod._SCHEMA + assert result == { + "day": "20240315", + "time": "143052", + "confidence": "high", + "source": "QuickTime:CreateDate", + "utc": False, + } diff --git a/think/detect_created.py b/think/detect_created.py index 273e39965..dc8ebe554 100644 --- a/think/detect_created.py +++ b/think/detect_created.py @@ -15,6 +15,10 @@ from typing import Optional from .prompts import load_prompt +_SCHEMA = json.loads( + (Path(__file__).parent / "detect_created.schema.json").read_text(encoding="utf-8") +) + def _load_system_prompt() -> str: """Load the system prompt from detect_created.txt file.""" @@ -99,6 +103,7 @@ def detect_created( thinking_budget=4096, system_instruction=_load_system_prompt(), json_output=True, + json_schema=_SCHEMA, ) try: diff --git a/think/detect_created.schema.json b/think/detect_created.schema.json new file mode 100644 index 000000000..42fa3259e --- /dev/null +++ b/think/detect_created.schema.json @@ -0,0 +1,27 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "additionalProperties": false, + "required": ["day", "time", "confidence", "source", "utc"], + "properties": { + "day": { + "type": "string", + "pattern": "^\\d{8}$" + }, + "time": { + "type": "string", + "pattern": "^\\d{6}$" + }, + "confidence": { + "type": "string", + "enum": ["high", "medium", "low"] + }, + "source": { + "type": "string", + "minLength": 1 + }, + "utc": { + "type": "boolean" + } + } +}