annoscript 0.1.0__py3-none-any.whl

annoscript/test_examples.py

@@ -0,0 +1,149 @@
+"""Client-side example snapshots and compilation; no backend dependencies."""
+
+from copy import deepcopy
+import json
+import unittest
+from unittest.mock import patch
+
+
+from annoscript_examples.fixtures import REVIEW_INPUT as REVIEW, LAUNCH_INPUT as LAUNCH
+
+
+from annoscript.compiler import compile_program
+from annoscript.ir import Program
+from annoscript.anno import Example, ExtractResponse, GenerateResponse, GenerateListResponse
+from annoscript._test_ir_runtime import evaluate
+
+
+LOCATION = Example(LAUNCH, ["广州"])
+
+
+def captured(anno):
+    return anno.extract("地点", source_field="text", max_items=2,
+                       max_length_per_item=80, examples=[LOCATION]).items
+
+
+def constructed(anno):
+    state = {"text": "示例在广州发布"}
+    answer = ["广州"]
+    example = Example(state=state, answer=answer)
+    state["text"] = "示例条目"
+    answer[0] = "条目"
+    return anno.extract("地点", source_field="text", max_items=2,
+                       max_length_per_item=80, examples=[example]).items
+
+
+def different_examples(anno):
+    # Deliberately vary only the answer to check identity, not business accuracy.
+    first = anno.choice("情感", {"正面": "积极", "中性": "中立"},
+                       examples=[Example(REVIEW, "正面")]).choice
+    second = anno.choice("情感", {"正面": "积极", "中性": "中立"},
+                        examples=[Example(REVIEW, "中性")]).choice
+    third = anno.choice("情感", {"正面": "积极", "中性": "中立"},
+                       examples=[Example(REVIEW, "正面")]).choice
+    return [first, second, third]
+
+
+def named_generations(anno):
+    title = anno.generate("标题", max_length=60, field_name="event_title").text
+    dimensions = anno.generate_list("维度", max_items=8, max_length_per_item=24,
+                                   field_name="tags").items
+    locations = anno.extract("地点", source_field="text", max_items=8,
+                            max_length_per_item=64, field_name="locations").items
+    return [title, dimensions, locations]
+
+
+def nullable_example(anno):
+    value = anno.generate("缺依据为 null", nullable=True, min_length=1, max_length=24,
+                         examples=[Example(state="示例正文", answer=None)]).text
+    if value is not None:
+        return value.strip()
+    return value
+
+
+def requirements_for_all_operations(anno):
+    requirements = ["仅依据正文。", "保留原意。"]
+    a = anno.noul("判断", requirements=requirements).noul
+    b = anno.choice("选择", {"yes": "是"}, requirements=requirements).choice
+    c = anno.score("评分", ["低", "高"], requirements=requirements).score
+    d = anno.generate("标签", max_length=24, requirements=requirements).text
+    e = anno.generate_list("维度", max_items=2, max_length_per_item=24,
+                          requirements=requirements).items
+    f = anno.extract("地点", source_field="text", max_items=2, max_length_per_item=24,
+                    requirements=requirements).items
+    return [a, b, c, d, e, f]
+
+
+class ExampleClientTests(unittest.TestCase):
+    def test_requirements_survive_client_ir_for_all_operations(self):
+        from annoscript.anno import NoulResponse, ChoiceResponse, ScoreResponse
+
+        class Probe:
+            def __init__(self):
+                self.calls = []
+
+            def record(self, operation, options, response):
+                self.calls.append(operation)
+                assert options["requirements"] == ["仅依据正文。", "保留原意。"]
+                return response
+
+            def noul(self, instructions, **options):
+                return self.record("noul", options, NoulResponse(.8))
+
+            def choice(self, instructions, criteria, **options):
+                return self.record("choice", options, ChoiceResponse("yes", .9, {"yes": 1.}))
+
+            def score(self, instructions, criteria, **options):
+                return self.record("score", options, ScoreResponse({0: 0., 1: 1.}, 1., dict(enumerate(criteria)), .9))
+
+            def generate(self, instructions, **options):
+                return self.record("generate", options, GenerateResponse("标签"))
+
+            def generate_list(self, instructions, **options):
+                return self.record("generate_list", options, GenerateListResponse(["维度"]))
+
+            def extract(self, instructions, **options):
+                return self.record("extract", options, ExtractResponse(["广州"]))
+
+        for optimize in (False, True):
+            probe = Probe()
+            program = Program.from_json(compile_program(requirements_for_all_operations, optimize=optimize).to_json())
+            self.assertEqual(evaluate(program.to_dict(), probe), [.8, "yes", 1., "标签", ["维度"], ["广州"]])
+            self.assertEqual(probe.calls, ["noul", "choice", "score", "generate", "generate_list", "extract"])
+
+    def test_nullable_example_and_none_survive_client_ir(self):
+        class Probe:
+            state = "正文"
+
+            def generate(self, instructions, **options):
+                assert options["nullable"] is True
+                example = options["examples"][0]
+                example = example.to_dict() if type(example) is Example else example
+                assert example["answer"] is None
+                return GenerateResponse(None)
+
+        for optimize in (False, True):
+            program = Program.from_json(compile_program(nullable_example, optimize=optimize).to_json())
+            self.assertIsNone(evaluate(program.to_dict(), Probe()))
+
+    def test_field_names_survive_client_ir_roundtrip(self):
+        class Probe:
+            state = {"text": "广州"}
+
+            def generate(self, instructions, **options):
+                return GenerateResponse(options["field_name"])
+
+            def generate_list(self, instructions, **options):
+                return GenerateListResponse([options["field_name"]])
+
+            def extract(self, instructions, **options):
+                return ExtractResponse([options["field_name"]])
+
+        for optimize in (False, True):
+            program = Program.from_json(compile_program(named_generations, optimize=optimize).to_json())
+            self.assertEqual(evaluate(program.to_dict(), Probe()),
+                             ["event_title", ["tags"], ["locations"]])
+
+
+if __name__ == "__main__":
+    unittest.main()

annoscript/test_field_names.py

@@ -0,0 +1,221 @@
+"""Public task names through the real wire, planner, runtime and durable jobs."""
+
+import asyncio
+from dataclasses import replace
+import json
+from pathlib import Path
+import tempfile
+import unittest
+
+import httpx
+
+from annoscript._data import thaw
+from annoscript.backends import InferenceContext, LLMAPIBackend, InferenceError
+from annoscript.backends.base import normalize_call
+from annoscript.compiler import compile_program
+from annoscript.ir import Program
+from annoscript.plan import ExecutionPlan
+from annoscript.plan_lowering import lower_to_plan
+from annoscript.runtime import AnnoRuntime, RunInput
+from annoscript.server.releases import Releases
+from annoscript.server.store import Store
+from annoscript.server.worker import Worker
+from annoscript.test_jobs_server import configurations, payload
+from annoscript.test_runtime import FakeBackend
+from annoscript._test_evaluation_http import FixtureLLMBackend, llm_calls
+
+
+def named_pair(anno):
+    first = anno.generate("same", max_length=10, max_num_tokens=32, field_name="first").text
+    second = anno.generate("same", max_length=10, max_num_tokens=32, field_name="second").text
+    again = anno.generate("same", max_length=10, max_num_tokens=32, field_name="first").text
+    return [first, second, again]
+
+
+def completion(content):
+    return httpx.Response(200, json={"choices": [{"finish_reason": "stop", "message": {"content": content}}],
+                                     "usage": {"prompt_tokens": 7, "completion_tokens": 2}})
+
+
+def named_handler(sent, *, repair=False):
+    def handle(request):
+        body = json.loads(request.content)
+        sent.append(body)
+        batched, calls = llm_calls(body)
+        if batched:
+            answer = {call.arguments["field_name"]: call.arguments["field_name"] for _, call in calls}
+            if repair:
+                answer.pop("second")
+            return completion(json.dumps({"answers": answer}))
+        return completion(calls[0][1].arguments["field_name"])
+    return handle
+
+
+class FieldNameTests(unittest.IsolatedAsyncioTestCase):
+    def make_call(self, operation="generate", name=None, identity="run1-q1", context=None, **options):
+        context = context or InferenceContext({"text": 'q1 q10 甲"\\\n😀'})
+        defaults = {"max_length": 10} if operation == "generate" else {"max_items": 2, "max_length_per_item": 10}
+        if operation == "extract":
+            defaults["source_field"] = "text"
+        if name is not None:
+            defaults["field_name"] = name
+        return normalize_call(identity, operation, ("保留正文中的 q1 与 q10",), {**defaults, **options}, context)
+
+    async def test_none_and_omitted_preserve_exact_payload_and_binding(self):
+        ctx = InferenceContext({"text": "甲"})
+        backend = LLMAPIBackend("http://offline.invalid", "fixture", "fixture", profile="qwen_vllm")
+        try:
+            for op in ("generate", "generate_list", "extract"):
+                absent = self.make_call(op, context=ctx)
+                explicit = self.make_call(op, context=ctx, field_name=None)
+                self.assertEqual(absent, explicit)
+                self.assertNotIn("field_name", absent.arguments)
+                a, b = backend._prepare(ctx, [absent], 128), backend._prepare(ctx, [explicit], 128)
+                self.assertEqual((a.binding, a.body, a.payload), (b.binding, b.body, b.payload))
+                other = self.make_call(op, identity="run1-q2", context=ctx)
+                self.assertEqual(backend._batch_payload(ctx, [absent, other], 256),
+                                 backend._batch_payload(ctx, [explicit, other], 256))
+        finally:
+            await backend.aclose()
+
+    async def test_validation_and_collisions_never_send_http(self):
+        sent = []
+        ctx = InferenceContext({"text": "甲"})
+        backend = LLMAPIBackend("http://offline.invalid", "fixture", "fixture",
+                                transport=httpx.MockTransport(lambda r: sent.append(r)))
+        try:
+            for op in ("generate", "generate_list", "extract"):
+                for value in ("", " \t\n", False, 1, [], {}):
+                    with self.subTest(op=op, value=value), self.assertRaisesRegex(ValueError, "field_name"):
+                        self.make_call(op, context=ctx, field_name=value)
+            for names in (("same", "same"), ("q2", None)):
+                calls = [self.make_call(name=n, identity=str(i), context=ctx) for i, n in enumerate(names)]
+                with self.assertRaisesRegex(ValueError, "duplicate model-visible"):
+                    await backend.infer_batch(ctx, calls, max_tokens=128)
+                with self.assertRaisesRegex(ValueError, "duplicate model-visible"):
+                    backend.plan_batch(ctx, calls, max_calls=8, max_input_bytes=100000,
+                                       max_output_tokens=4096, default_tokens=1024)
+            self.assertEqual(sent, [])
+        finally:
+            await backend.aclose()
+
+    async def test_mixed_names_examples_grammar_and_reordered_answers(self):
+        ctx = InferenceContext({"text": 'q1 q10 甲"\\\n😀'})
+        strange = ' 地点"\\\n😀 '
+        calls = [self.make_call(name="q10", context=ctx),
+                 self.make_call("extract", name=strange, identity="run1-q2", context=ctx,
+                                examples=[{"state": {"text": "广州"}, "answer": ["广州"]}]),
+                 self.make_call("generate_list", identity="run1-q3", context=ctx)]
+        body = {"answers": {"q3": [], strange: ["甲"], "q10": "标题"}}
+        for profile in ("legacy", "qwen_vllm"):
+            sent, events = [], []
+            def handle(request):
+                sent.append(json.loads(request.content))
+                return completion(json.dumps(body, ensure_ascii=False))
+            backend = LLMAPIBackend("http://offline.invalid", "fixture", "fixture", profile=profile,
+                                    transport=httpx.MockTransport(handle), diagnostic_observer=events.append)
+            try:
+                result = await backend.infer_batch(ctx, calls, max_tokens=256)
+                self.assertEqual(result.errors, {})
+                self.assertEqual(result.responses[calls[0].id].text, "标题")
+                self.assertEqual(result.responses[calls[1].id].items, ["甲"])
+                self.assertEqual(result.responses[calls[2].id].items, [])
+                prompt = sent[0]["messages"][0]["content"]
+                self.assertIn('保留正文中的 q1 与 q10', prompt)
+                self.assertIn('任务 ' + json.dumps(strange, ensure_ascii=False), prompt)
+                self.assertIn('示例输出：' + json.dumps({"answers": {strange: ["广州"]}}, ensure_ascii=False, separators=(",", ":")), prompt)
+                self.assertEqual(json.loads(sent[0]["messages"][1]["content"]), thaw(ctx.state))
+                self.assertTrue(any(e.get("wire_call_ids") == {"q10": calls[0].id, strange: calls[1].id, "q3": calls[2].id}
+                                    for e in events))
+                if profile == "qwen_vllm":
+                    root = sent[0]["structured_outputs"]["grammar"].splitlines()[0]
+                    self.assertIn("q10", root)
+                    self.assertNotIn("run1", root)
+            finally:
+                await backend.aclose()
+
+    async def test_single_and_repair_only_rename_task_identity(self):
+        ctx = InferenceContext({"text": "甲"})
+        for op, text in (("generate", "标题"), ("extract", '{"items":["甲"]}'),
+                         ("generate_list", '{"items":[]}')):
+            sent, events = [], []
+            def handle(request):
+                sent.append(json.loads(request.content))
+                return completion(text)
+            backend = LLMAPIBackend("http://offline.invalid", "fixture", "fixture", profile="qwen_vllm",
+                                    transport=httpx.MockTransport(handle), diagnostic_observer=events.append)
+            try:
+                plain = self.make_call(op, context=ctx, max_num_tokens=64)
+                named = self.make_call(op, name="summary", context=ctx, max_num_tokens=64)
+                for feedback in (None, "invalid answer"):
+                    expected = backend._single_payload(ctx, plain, 64, feedback)
+                    await backend.infer(ctx, named, repair_feedback=feedback)
+                    actual = sent[-1]
+                    actual["messages"][0]["content"] = actual["messages"][0]["content"].replace('任务 "summary"', '任务 "run1-q1"')
+                    self.assertEqual(actual, expected)
+                self.assertTrue(any(e.get("wire_call_ids") == {"summary": named.id} for e in events))
+            finally:
+                await backend.aclose()
+
+    async def test_names_bind_preparation_and_remain_stable_in_split_subsets(self):
+        ctx = InferenceContext({"text": "甲"})
+        backend = LLMAPIBackend("http://offline.invalid", "fixture", "fixture", profile="qwen_vllm")
+        try:
+            calls = [self.make_call(name=n, identity=str(i), context=ctx) for i, n in enumerate(("first", "second", "third"))]
+            first = backend._prepare(ctx, calls[:2], 128, batch=True)
+            second = backend._prepare(ctx, calls[1:], 128, batch=True)
+            self.assertEqual(dict(first.evidence["wire_call_ids"]), {"first": "0", "second": "1"})
+            self.assertEqual(dict(second.evidence["wire_call_ids"]), {"second": "1", "third": "2"})
+            changed = replace(calls[0], arguments={**calls[0].arguments, "field_name": "different"})
+            self.assertNotEqual(first.binding, backend._prepare(ctx, [changed, calls[1]], 128, batch=True).binding)
+        finally:
+            await backend.aclose()
+
+    async def test_ir_runtime_reuse_and_repair_in_all_modes(self):
+        for optimize in (False, True):
+            program = Program.from_json(compile_program(named_pair, optimize=optimize).to_json())
+            plan = ExecutionPlan.from_json(lower_to_plan(program).to_json())
+            for mode in ("strict", "eager", "speculative"):
+                for repair in (False, True):
+                    sent = []
+                    backend = FixtureLLMBackend("http://offline.invalid", "fixture", "fixture", profile="qwen_vllm",
+                                                transport=httpx.MockTransport(named_handler(sent, repair=repair)))
+                    async with AnnoRuntime(FakeBackend(), backend, trace=True, max_llm_batch_calls=16,
+                                          max_llm_batch_output_tokens=4096) as runtime:
+                        results = [r async for r in await runtime.submit(plan, [RunInput(0, "甲")], mode=mode)]
+                    self.assertEqual(results[0].status, "success", results[0])
+                    self.assertEqual(results[0].value, ["first", "second", "first"])
+                    self.assertEqual(sum(len(llm_calls(b)[1]) for b in sent), 3 if repair and mode != "strict" else 2)
+                    self.assertTrue(all(b["max_tokens"] == (192 if llm_calls(b)[0] else 32) for b in sent))
+                    if repair and mode != "strict":
+                        self.assertEqual(len(sent), 2)
+                        self.assertIn('任务 "second"', sent[-1]["messages"][0]["content"])
+
+
+class DurableFieldNameTests(unittest.TestCase):
+    def test_names_roundtrip_through_worker_and_collect_does_not_infer(self):
+        with tempfile.TemporaryDirectory() as directory:
+            store = Store(Path(directory) / "jobs.sqlite")
+            releases = Releases(configurations(), "test")
+            sent = []
+            def factory(_):
+                return FakeBackend(), FixtureLLMBackend("http://offline.invalid", "fixture", "fixture", profile="qwen_vllm",
+                                                        transport=httpx.MockTransport(named_handler(sent)))
+            request = payload(("甲", "乙"), mode="speculative", program=named_pair)
+            job = store.submit("alice", "named", request, releases)
+            worker = Worker(store, releases, backend_factory=factory)
+            async def run():
+                while (claim := store.claim("worker", releases)) is not None:
+                    await worker.execute(claim)
+            asyncio.run(run())
+            before = len(sent)
+            for _ in range(2):
+                page = store.results("alice", job["job_id"])
+                self.assertEqual([r["value"] for r in page["results"]], [["first", "second", "first"]] * 2)
+                self.assertEqual(store.submit("alice", "named", request, releases)["job_id"], job["job_id"])
+                self.assertIsNone(store.claim("worker", releases))
+            self.assertEqual(len(sent), before)
+
+
+if __name__ == "__main__":
+    unittest.main()

annoscript/test_formatter.py

@@ -0,0 +1,153 @@
+from dataclasses import replace
+import io
+import os
+from pathlib import Path
+import subprocess
+import sys
+import tempfile
+import unittest
+from unittest.mock import patch
+
+from annoscript import ir
+from annoscript.compiler import compile_program
+from annoscript.formatter import format_ir
+from annoscript.inspect import main
+
+
+_ROOT = Path(__file__).resolve().parents[1]
+
+
+def example_program():
+    span = ir.SourceSpan("missing.source.module", "annotate", 1, 1, 12, 1)
+    source = ir.Source(span, (ir.SourceSpan("caller", "entry", 5, 1, 5, 20),))
+    result = ir.Var("decision")
+    body = (
+        ir.Assign("text", ir.GetItem(ir.State(), ir.Literal("text")), id="s1", source=source),
+        ir.Infer("decision", "choice", (ir.Literal("判断事件"), ir.DictExpr((
+            (ir.Literal("yes"), ir.Literal("事件")), (ir.Literal("no"), ir.Literal("非事件")),
+        ))), (), id="s2", source=source),
+        ir.If(ir.Compare(ir.Field(result, "choice"), ("eq",), (ir.Literal("yes"),)), (
+            ir.Constrain(ir.Compare(ir.Intrinsic("len", (ir.Var("text"),)), ("gt",), (ir.Literal(0),)),
+                         ir.Literal("必须有正文"), id="s3", source=source),
+            ir.Return(ir.Var("text"), id="s4", source=source),
+        ), (
+            ir.Return(ir.Literal(None), id="s5", source=source),
+        ), id="s6", source=source),
+    )
+    return ir.Program("demo", ir.ProgramConfig("公共规则"), body, span)
+
+
+class FormatterTests(unittest.TestCase):
+    def test_readable_model_call_branch_and_constraint(self):
+        program = example_program()
+        self.assertEqual(format_ir(program), "\n".join((
+            "# AnnoScript IR v1 (display notation)",
+            "system_prompt '公共规则'",
+            "program 'demo':",
+            "    text = anno.state['text']",
+            "    decision = anno.choice('判断事件', {'yes': '事件', 'no': '非事件'})",
+            "    if (decision.choice == 'yes'):",
+            "        constrain((len(text) > 0), message='必须有正文')",
+            "        return text",
+            "    else:",
+            "        return None",
+            "",
+        )))
+        self.assertEqual(format_ir(ir.Program.from_json(program.to_json())), format_ir(program))
+
+    def test_source_comments_are_optional(self):
+        program = example_program()
+        self.assertNotIn("missing.source.module", format_ir(program))
+        text = format_ir(program, show_source=True)
+        self.assertIn("# [s2] missing.source.module:1:1 (annotate)", text)
+        self.assertIn("inlined from caller:5:1 (entry)", text)
+        self.assertIn("decision = anno.choice", text)
+
+    def test_all_compiled_constructs_render_without_executing_models(self):
+        from annoscript import test_compiler as examples
+        for function in (
+            examples.loop_helpers, examples.expression_order, examples.short_circuit,
+            examples.constrained, examples.all_operations, examples.formatting,
+        ):
+            with self.subTest(function=function.__name__):
+                program = compile_program(function)
+                before = program.to_dict()
+                text = format_ir(program)
+                self.assertTrue(text.endswith("\n"))
+                self.assertEqual(before, program.to_dict())
+                self.assertEqual(text, format_ir(ir.Program.from_json(program.to_json())))
+        loops = format_ir(compile_program(examples.loop_helpers))
+        self.assertIn("for ", loops)
+        formatting = format_ir(compile_program(examples.formatting))
+        self.assertIn("format_string(", formatting)
+        self.assertIn("format_value(", formatting)
+        self.assertIn("conversion='r'", formatting)
+        self.assertIn("json.dumps(", formatting)
+        calls = format_ir(compile_program(examples.all_operations))
+        for operation in ("noul", "choice", "score", "extract", "generate", "generate_list"):
+            self.assertIn(f"anno.{operation}(", calls)
+
+    def test_empty_blocks_literals_and_unusual_names(self):
+        program = example_program()
+        statement = ir.Return(ir.TupleExpr((ir.Literal("中文\n'{}\\"),)), id="s1", source=ir.Source(program.source))
+        program = replace(program, config=ir.ProgramConfig(), body=(statement,))
+        text = format_ir(program)
+        self.assertIn(repr("中文\n'{}\\"), text)
+        self.assertEqual(len(text.splitlines()), 3)
+        self.assertTrue(text.rstrip().endswith(",)"))
+        self.assertNotIn("system_prompt", text)
+        self.assertEqual(format_ir(replace(program, body=())).splitlines()[-1], "    pass")
+        statement = replace(statement, value=ir.Var("has space"))
+        self.assertIn("var['has space']", format_ir(replace(program, body=(statement,))))
+
+
+class InspectCommandTests(unittest.TestCase):
+    def run_command(self, *args, stdin=None):
+        return subprocess.run(
+            [sys.executable, "-m", "annoscript.inspect", *args], input=stdin,
+            text=True, encoding="utf-8", capture_output=True, cwd=_ROOT,
+            env={**os.environ, "PYTHONDONTWRITEBYTECODE": "1"}, timeout=10,
+        )
+
+    def test_arbitrary_path_and_source_flag(self):
+        program = example_program()
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "任意 program.json"
+            path.write_text(program.to_json(), encoding="utf-8")
+            result = self.run_command(str(path), "--show-source")
+        self.assertEqual(result.returncode, 0, result.stderr)
+        self.assertEqual(result.stderr, "")
+        self.assertEqual(result.stdout, format_ir(program, show_source=True))
+
+    def test_stdin_and_help(self):
+        program = example_program()
+        result = self.run_command("-", stdin=program.to_json())
+        self.assertEqual(result.returncode, 0, result.stderr)
+        self.assertEqual(result.stdout, format_ir(program))
+        result = self.run_command("--help")
+        self.assertEqual(result.returncode, 0)
+        self.assertIn("--show-source", result.stdout)
+
+    def test_invalid_files_report_errors_without_tracebacks(self):
+        for text in ("not JSON", '{"format":"lev-ir","version":99}'):
+            with self.subTest(text=text):
+                result = self.run_command("-", stdin=text)
+                self.assertEqual(result.returncode, 2)
+                self.assertEqual(result.stdout, "")
+                self.assertIn("annoscript.inspect:", result.stderr)
+                self.assertNotIn("Traceback", result.stderr)
+        with tempfile.TemporaryDirectory() as directory:
+            result = self.run_command(str(Path(directory) / "missing.json"))
+        self.assertEqual(result.returncode, 2)
+        self.assertIn("annoscript.inspect:", result.stderr)
+
+    def test_main_can_be_embedded_without_exiting_process(self):
+        output = io.StringIO()
+        with patch("sys.stdin", io.StringIO(example_program().to_json())), patch("sys.stdout", output):
+            status = main(["-"])
+        self.assertEqual(status, 0)
+        self.assertEqual(output.getvalue(), format_ir(example_program()))
+
+
+if __name__ == "__main__":
+    unittest.main()

annoscript/test_fusion_scheduling.py

@@ -0,0 +1,109 @@
+"""Shared admission and nullable semantics; synthetic/local HTTP only."""
+import asyncio
+import json
+import tempfile
+import unittest
+from pathlib import Path
+
+import httpx
+from annoscript._dispatch_pool import RequestSlots
+from annoscript.backends import BatchResult, InferenceOutput, DecisionAPIBackend, LLMAPIBackend
+from annoscript.compiler import compile_program
+from annoscript.anno import ChoiceResponse, GenerateResponse
+from annoscript.plan_lowering import lower_to_plan
+from annoscript.runtime import AnnoRuntime, RunInput
+from annoscript.server.worker import Worker
+from annoscript.server.store import Store
+from annoscript.server.releases import Releases
+from annoscript.test_jobs_server import configurations
+from annoscript.remote import ir_sha256
+
+
+def join_ready(anno):
+    first = anno.choice("ready immediately", {"yes": "yes", "no": "no"}).choice
+    text = anno.generate("synthetic", max_length=10).text
+    second = anno.choice(text, {"yes": "yes", "no": "no"}).choice
+    return [first, second]
+
+
+class SlotTests(unittest.IsolatedAsyncioTestCase):
+    async def test_fifo_cancel_duplicate_wait_and_idempotent_release(self):
+        p = RequestSlots(generation=1, decision=1)
+        woke = []
+        r = p.acquire("llm", "a", lambda: woke.append("a"))
+        self.assertIsNone(p.acquire("llm", "b", lambda: woke.append("b")))
+        self.assertIsNone(p.acquire("llm", "c", lambda: woke.append("c")))
+        self.assertIsNone(p.acquire("llm", "b", lambda: woke.append("duplicate")))
+        r(); r()
+        self.assertEqual(woke, ["b"])
+        self.assertIsNone(p.acquire("llm", "c", lambda: None))
+        p.cancel("b")
+        c = p.acquire("llm", "c", lambda: None); c()
+        self.assertEqual(p.active, {"llm": 0, "decision": 0})
+        self.assertFalse(p.waiting("b"))
+
+    async def test_waiting_batches_are_selected_at_admission(self):
+        pool = RequestSlots(generation=8, decision=1)
+        release = pool.acquire("decision", "external", lambda: None)
+        batches = []
+        generated = asyncio.Event()
+        class Backend:
+            async def infer(self, ctx, call):
+                generated.set()
+                return InferenceOutput(GenerateResponse("newly ready"))
+            async def infer_batch(self, ctx, calls):
+                batches.append([c.arguments["instructions"] for c in calls])
+                return BatchResult({c.id: ChoiceResponse("yes", 1., {"yes": 1., "no": 0.}) for c in calls})
+            async def aclose(self): pass
+        b = Backend()
+        async with AnnoRuntime(b, b, request_slots=pool, trace=True, max_decision_inflight=1) as runtime:
+            stream = await runtime.submit(lower_to_plan(compile_program(join_ready)), [RunInput(0, "fixture")])
+            await generated.wait()
+            for _ in range(8): await asyncio.sleep(0)
+            self.assertFalse(batches)
+            release()
+            result = [r async for r in stream]
+            self.assertEqual(result[0].value, ["yes", "yes"])
+            self.assertEqual(batches, [["ready immediately", "newly ready"]])
+            self.assertTrue(any(e['event'] == 'slot_wait' for e in runtime.trace))
+        self.assertEqual(sum(pool.active.values()), 0)
+
+    async def test_cancel_waiter_does_not_send_or_leak(self):
+        pool = RequestSlots(generation=1, decision=1)
+        a = pool.acquire('decision', 'held', lambda: None)
+        b = pool.acquire('llm', 'held', lambda: None)
+        class Never:
+            async def infer(self, *a, **k): raise AssertionError('sent')
+            async def infer_batch(self, *a, **k): raise AssertionError('sent')
+            async def aclose(self): pass
+        async with AnnoRuntime(Never(), Never(), request_slots=pool) as rt:
+            stream = await rt.submit(lower_to_plan(compile_program(join_ready)), [RunInput(0, 'fixture')])
+            for _ in range(8): await asyncio.sleep(0)
+            await stream.aclose()
+        a(); b()
+        self.assertFalse(any(pool.waiters.values()))
+        self.assertEqual(sum(pool.active.values()), 0)
+
+
+class SharedWorkerTests(unittest.IsolatedAsyncioTestCase):
+    pass
+
+
+from annoscript.test_generation_policy import P0ProtocolWorkerTests, P0DiagnosticWorkerTests
+
+class SharedProtocolTests(P0ProtocolWorkerTests):
+    def run_http(self, handler, *, observer=None):
+        w = super().run_http(handler, observer=observer)
+        w.concurrency = 8
+        w.request_slots = RequestSlots(generation=8, decision=4)
+        return w
+
+class SharedDiagnosticTests(P0DiagnosticWorkerTests):
+    def worker(self, handler, **kw):
+        w = super().worker(handler, **kw)
+        w.concurrency = 8
+        w.request_slots = RequestSlots(generation=8, decision=4)
+        return w
+
+
+if __name__ == '__main__': unittest.main()