diff --git a/cragents/_utils.py b/cragents/_utils.py index d4630e0..8e1dfc6 100644 --- a/cragents/_utils.py +++ b/cragents/_utils.py @@ -154,7 +154,7 @@ def build_grammar(generation_sequence: Sequence[GenerationSequenceElement]) -> s tool_names = [f'"{tool_name}"' for tool_name in element.tool_names] custom_defs.append(f"FUNCTION_NAME: ({' | '.join(tool_names)})") - grammar = "\n".join([start_def] + custom_defs + default_defs) + grammar = "\n".join([start_def.strip()] + custom_defs + default_defs) return grammar diff --git a/tests/test_agent.py b/tests/test_agent.py index 828064f..deb7bc8 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -146,29 +146,29 @@ async def test_mixed_output_type(): async def test_set_guide_requires_openai_model(): agent = CRAgent(TestModel()) with pytest.raises(RuntimeError, match="OpenAIChatModel required"): - await agent.set_guide([Anchor("hi ")]) + await agent.set_guide([Anchor("hi")]) async def test_set_guide_creates_model_settings_when_none(): agent = CRAgent(model) assert agent.model_settings is None - await agent.set_guide([Anchor("hi ")]) + await agent.set_guide([Anchor("hi")]) assert agent.model_settings is not None assert "extra_body" in agent.model_settings async def test_set_guide_preserves_existing_model_settings(): agent = CRAgent(model, model_settings=OpenAIChatModelSettings(temperature=0.5)) - await agent.set_guide([Anchor("hi ")]) + await agent.set_guide([Anchor("hi")]) assert agent.model_settings["temperature"] == 0.5 assert "extra_body" in agent.model_settings async def test_set_guide_overwrites_on_second_call(): agent = CRAgent(model) - await agent.set_guide([Anchor("first ")]) + await agent.set_guide([Anchor("first")]) first_grammar = agent.model_settings["extra_body"]["structured_outputs"]["grammar"] - await agent.set_guide([Anchor("second ")]) + await agent.set_guide([Anchor("second")]) second_grammar = agent.model_settings["extra_body"]["structured_outputs"]["grammar"] assert first_grammar != second_grammar assert "second" in second_grammar diff --git a/tests/test_grammar.py b/tests/test_grammar.py index 4ae39c8..ddd95fd 100644 --- a/tests/test_grammar.py +++ b/tests/test_grammar.py @@ -9,125 +9,152 @@ def test_grammar_single_anchor(): grammar = build_grammar([Anchor("hello ")]) assert grammar == snapshot("""\ -start: "hello " \nFREE: /[\\S\\s]*/ +start: "hello " +FREE: /[\\S\\s]*/ NL: /\\n/""") def test_grammar_multiple_anchors(): grammar = build_grammar([Anchor("foo "), Anchor("bar ")]) assert grammar == snapshot("""\ -start: "foo " "bar " \nFREE: /[\\S\\s]*/ -NL: /\\n/""") +start: "foo " "bar " +FREE: /[\\S\\s]*/ +NL: /\\n/\ +""") def test_grammar_single_free(): grammar = build_grammar([Free()]) assert grammar == snapshot("""\ -start: FREE \nFREE: /[\\S\\s]*/ -NL: /\\n/""") +start: FREE +FREE: /[\\S\\s]*/ +NL: /\\n/\ +""") def test_grammar_anchor_then_free(): grammar = build_grammar([Anchor("Result: "), Free()]) assert grammar == snapshot("""\ -start: "Result: " FREE \nFREE: /[\\S\\s]*/ -NL: /\\n/""") +start: "Result: " FREE +FREE: /[\\S\\s]*/ +NL: /\\n/\ +""") def test_grammar_single_constrain(): grammar = build_grammar([Constrain(max_newlines=2, max_char_captures=3)]) assert grammar == snapshot("""\ -start: block_1 \nblock_1: p_1{1,2} +start: block_1 +block_1: p_1{1,2} p_1: s_1{1,3} NL NL s_1[lazy]: /[^\\.\\n]+/ ( "." ) FREE: /[\\S\\s]*/ -NL: /\\n/""") +NL: /\\n/\ +""") def test_grammar_constrain_custom_chars(): grammar = build_grammar([Constrain(max_newlines=1, max_char_captures=2, chars_to_capture="!?")]) assert grammar == snapshot("""\ -start: block_1 \nblock_1: p_1{1,1} +start: block_1 +block_1: p_1{1,1} p_1: s_1{1,2} NL NL s_1[lazy]: /[^!\\?\\n]+/ ( "!" | "?" ) FREE: /[\\S\\s]*/ -NL: /\\n/""") +NL: /\\n/\ +""") def test_grammar_constrain_regex_special_chars(): # "." is the default and must be escaped to "\\." in the character class grammar = build_grammar([Constrain(max_newlines=1, max_char_captures=1, chars_to_capture=".")]) assert grammar == snapshot("""\ -start: block_1 \nblock_1: p_1{1,1} +start: block_1 +block_1: p_1{1,1} p_1: s_1{1,1} NL NL s_1[lazy]: /[^\\.\\n]+/ ( "." ) FREE: /[\\S\\s]*/ -NL: /\\n/""") +NL: /\\n/\ +""") def test_grammar_multiple_constrains_uid_increments(): grammar = build_grammar([Constrain(1, 1), Constrain(2, 3)]) assert grammar == snapshot("""\ -start: block_1 block_2 \nblock_1: p_1{1,1} +start: block_1 block_2 +block_1: p_1{1,1} p_1: s_1{1,1} NL NL s_1[lazy]: /[^\\.\\n]+/ ( "." ) block_2: p_2{1,2} p_2: s_2{1,3} NL NL s_2[lazy]: /[^\\.\\n]+/ ( "." ) FREE: /[\\S\\s]*/ -NL: /\\n/""") +NL: /\\n/\ +""") def test_grammar_think_empty_sequence(): grammar = build_grammar([Think([])]) assert grammar == snapshot("""\ -start: NL \nFREE: /[\\S\\s]*/ -NL: /\\n/""") +start: NL +FREE: /[\\S\\s]*/ +NL: /\\n/\ +""") def test_grammar_think_with_anchor(): grammar = build_grammar([Think([Anchor("step one ")])]) assert grammar == snapshot("""\ -start: NL "step one " \nFREE: /[\\S\\s]*/ -NL: /\\n/""") +start: NL "step one " +FREE: /[\\S\\s]*/ +NL: /\\n/\ +""") def test_grammar_think_with_constrain(): grammar = build_grammar([Think([Constrain(1, 2)])]) assert grammar == snapshot("""\ -start: NL block_1 \nblock_1: p_1{1,1} +start: NL block_1 +block_1: p_1{1,1} p_1: s_1{1,2} NL NL s_1[lazy]: /[^\\.\\n]+/ ( "." ) FREE: /[\\S\\s]*/ -NL: /\\n/""") +NL: /\\n/\ +""") def test_grammar_think_with_free(): grammar = build_grammar([Think([Free()])]) assert grammar == snapshot("""\ -start: NL FREE \nFREE: /[\\S\\s]*/ -NL: /\\n/""") +start: NL FREE +FREE: /[\\S\\s]*/ +NL: /\\n/\ +""") def test_grammar_think_uid_shared_with_outer(): # Outer Constrain gets uid=1, inner Think's Constrain gets uid=2 grammar = build_grammar([Constrain(1, 1), Think([Constrain(2, 3)])]) assert grammar == snapshot("""\ -start: block_1 NL block_2 \nblock_1: p_1{1,1} +start: block_1 NL block_2 +block_1: p_1{1,1} p_1: s_1{1,1} NL NL s_1[lazy]: /[^\\.\\n]+/ ( "." ) block_2: p_2{1,2} p_2: s_2{1,3} NL NL s_2[lazy]: /[^\\.\\n]+/ ( "." ) FREE: /[\\S\\s]*/ -NL: /\\n/""") +NL: /\\n/\ +""") def test_grammar_think_custom_tokens(): grammar = build_grammar([Think([], start_token="<|think|>", stop_token="<|/think|>")]) assert grammar == snapshot("""\ -start: <|think|> NL <|/think|> \nFREE: /[\\S\\s]*/ -NL: /\\n/""") +start: <|think|> NL <|/think|> +FREE: /[\\S\\s]*/ +NL: /\\n/\ +""") def test_grammar_use_tools_explicit_schema():