fix(tools): harden fenced and embedded tool selection JSON (#115)

LightAgent/tools.py

@@ -23,6 +23,69 @@
 _TOOL_NAME_PATTERN = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
 _SUPPORTED_TOOL_PARAM_TYPES = {"string", "integer", "number", "boolean", "array", "object"}
 
+_FENCED_JSON_BLOCK = re.compile(
+    r"```(?:json)?\s*(.*?)```", re.DOTALL | re.IGNORECASE
+)
+
+
+def _selected_tool_names(payload: Any) -> frozenset[str]:
+    if not isinstance(payload, dict):
+        raise ValueError("Tool selection must be a JSON object")
+    tools = payload.get("tools", [])
+    if not isinstance(tools, list):
+        raise ValueError("Tool selection 'tools' must be a list")
+    if any(not isinstance(tool, dict) or not isinstance(tool.get("name"), str) for tool in tools):
+        raise ValueError("Each selected tool must have a string name")
+    return frozenset(tool["name"].strip().lower() for tool in tools)
+
+
+def _extract_json_payload(text: str) -> str:
+    """Preserve complete JSON, otherwise extract an unambiguous tool selection."""
+    content = text.strip()
+    if not content:
+        return content
+
+    try:
+        json.loads(content)
+    except json.JSONDecodeError:
+        pass
+    else:
+        return content
+
+    candidates = []
+
+    def add_candidate(candidate, payload):
+        if isinstance(payload, dict) and "tools" in payload:
+            candidates.append((candidate, _selected_tool_names(payload)))
+
+    for fenced in _FENCED_JSON_BLOCK.finditer(content):
+        candidate = fenced.group(1).strip()
+        try:
+            parsed = json.loads(candidate)
+        except json.JSONDecodeError:
+            continue
+        add_candidate(candidate, parsed)
+
+    # raw_decode respects nested objects, quoted braces, and escapes. Skip the
+    # complete decoded object rather than interpreting its children as choices.
+    decoder = json.JSONDecoder()
+    position = 0
+    while (position := content.find("{", position)) >= 0:
+        try:
+            parsed, end = decoder.raw_decode(content, position)
+        except json.JSONDecodeError:
+            position += 1
+            continue
+        add_candidate(content[position:end], parsed)
+        position = end
+
+    if not candidates:
+        return content
+    if any(names != candidates[0][1] for _, names in candidates[1:]):
+        raise ValueError("Conflicting tool selections in model output")
+    return candidates[0][0]
+
+
 
 class ToolRegistry:
     """集中管理工具注册表，避免全局变量"""
@@ -181,19 +244,17 @@ def filter_tools(self, tool_reflection_result: str) -> List[Dict]:
         """根据内容过滤工具"""
         try:
             # 安全解析可能包含 Markdown 代码块的 JSON
-            refined_content = tool_reflection_result.strip()
-            if refined_content.startswith('```json') and refined_content.endswith('```'):
-                refined_content = refined_content[7:-3].strip()
+            refined_content = _extract_json_payload(tool_reflection_result)
 
             parsed_data = json.loads(refined_content)
-            valid_tools = {tool["name"].strip().lower() for tool in parsed_data.get("tools", [])}
+            valid_tools = _selected_tool_names(parsed_data)
 
             return [
                 schema for schema in self.openai_function_schemas
                 if isinstance(schema, dict) and
                    schema.get("function", {}).get("name", "").strip().lower() in valid_tools
             ]
-        except (json.JSONDecodeError, KeyError, AttributeError) as e:
+        except (ValueError, TypeError, AttributeError) as e:
             raise ValueError(f"工具过滤失败: {str(e)}") from e
 
 

docs/tools.md

@@ -169,6 +169,13 @@ agent can reduce the set of tools visible to the model after an initial
 planning step. The `filter_tools()` method accepts a JSON string listing the
 tools the planner selected and returns only the matching schemas.
 
+Complete JSON is parsed before any Markdown extraction, so code fences inside
+JSON strings are preserved. Otherwise, JSON or unlabelled fenced blocks and
+objects embedded in prose are accepted when they contain a valid `tools` list.
+Unrelated JSON objects are skipped. Conflicting tool selections and malformed
+selection structures raise `ValueError`; an empty selection remains valid.
+Filtering never adds unregistered tools.
+
 ### Creating Custom Tools
 
 #### Basic Custom Tool

tests/test_tool_schema_validation.py

@@ -1,3 +1,7 @@
+import json
+
+import pytest
+
 from LightAgent import LightAgent, ToolRegistry
 
 
@@ -101,6 +105,124 @@ def test_registry_does_not_report_same_function_registered_again():
     assert registry.validate_tools() == []
 
 
+def _make_filterable_registry():
+    registry = ToolRegistry()
+    registry.register_tool(make_tool(name="search_database"))
+    registry.register_tool(make_tool(name="get_weather"))
+    return registry
+
+
+def test_filter_tools_accepts_plain_json():
+    registry = _make_filterable_registry()
+
+    result = registry.filter_tools("{\"tools\": [{\"name\": \"get_weather\"}]}")
+
+    names = [s["function"]["name"] for s in result]
+    assert names == ["get_weather"]
+
+
+def test_filter_tools_accepts_json_fenced_markdown_block():
+    registry = _make_filterable_registry()
+
+    result = registry.filter_tools('```json\n{"tools": [{"name": "search_database"}]}\n```')
+
+    names = [s["function"]["name"] for s in result]
+    assert names == ["search_database"]
+
+
+def test_filter_tools_accepts_unlabeled_markdown_code_block():
+    """LLMs frequently wrap JSON in a plain ``` fence without the `json` tag."""
+    registry = _make_filterable_registry()
+
+    result = registry.filter_tools('```\n{"tools": [{"name": "get_weather"}]}\n```')
+
+    names = [s["function"]["name"] for s in result]
+    assert names == ["get_weather"]
+
+
+def test_filter_tools_accepts_fence_with_surrounding_prose():
+    registry = _make_filterable_registry()
+
+    result = registry.filter_tools(
+        'Here are the tools I need:\n```json\n{"tools": [{"name": "search_database"}]}\n```'
+    )
+
+    names = [s["function"]["name"] for s in result]
+    assert names == ["search_database"]
+
+
+def test_filter_tools_accepts_json_embedded_in_prose_without_fence():
+    registry = _make_filterable_registry()
+
+    result = registry.filter_tools(
+        "I need the weather. {\"tools\": [{\"name\": \"get_weather\"}]} Let me fetch it."
+    )
+
+    names = [s["function"]["name"] for s in result]
+    assert names == ["get_weather"]
+
+
+def test_filter_tools_still_reports_truly_invalid_content():
+    registry = _make_filterable_registry()
+
+    with pytest.raises(ValueError):
+        registry.filter_tools("this is not json at all")
+
+
+def test_filter_tools_reports_invalid_content_inside_a_fence():
+    registry = _make_filterable_registry()
+
+    with pytest.raises(ValueError):
+        registry.filter_tools("```\nthis is not json either\n```")
+
+
+def test_filter_tools_preserves_valid_json_with_fences_inside_strings():
+    payload = {"tools": [{"name": "get_weather"}], "note": "```json {} ```"}
+    result = _make_filterable_registry().filter_tools(json.dumps(payload))
+    assert [item["function"]["name"] for item in result] == ["get_weather"]
+
+
+@pytest.mark.parametrize("prefix", ["Context {}. Selection: ", "Context {not JSON}. Selection: "])
+def test_filter_tools_skips_unrelated_braces_before_selection(prefix):
+    text = prefix + json.dumps({"tools": [{"name": "get_weather"}], "note": 'quoted } brace and "text"'})
+    result = _make_filterable_registry().filter_tools(text + " Done.")
+    assert [item["function"]["name"] for item in result] == ["get_weather"]
+
+
+def test_filter_tools_skips_invalid_first_fence():
+    text = '```\nnot JSON\n```\n```json\n{"tools":[{"name":"get_weather"}]}\n```'
+    result = _make_filterable_registry().filter_tools(text)
+    assert [item["function"]["name"] for item in result] == ["get_weather"]
+
+
+def test_filter_tools_rejects_conflicting_selections():
+    text = '{"tools":[{"name":"get_weather"}]} or {"tools":[{"name":"search_database"}]}'
+    with pytest.raises(ValueError, match="Conflicting tool selections"):
+        _make_filterable_registry().filter_tools(text)
+
+
+def test_filter_tools_accepts_repeated_equivalent_selections():
+    text = '{"tools":[{"name":" GET_WEATHER "}]} again {"tools":[{"name":"get_weather"}]}'
+    result = _make_filterable_registry().filter_tools(text)
+    assert [item["function"]["name"] for item in result] == ["get_weather"]
+
+
+@pytest.mark.parametrize("payload", [{}, {"tools": []}])
+def test_filter_tools_preserves_empty_selection_compatibility(payload):
+    assert _make_filterable_registry().filter_tools(json.dumps(payload)) == []
+
+
+@pytest.mark.parametrize("payload", [[], {"tools": None}, {"tools": {}}, {"tools": [None]},
+                                      {"tools": [{}]}, {"tools": [{"name": 42}]}])
+def test_filter_tools_reports_invalid_selection_shapes(payload):
+    with pytest.raises(ValueError):
+        _make_filterable_registry().filter_tools(json.dumps(payload))
+
+
+def test_filter_tools_cannot_add_unregistered_names():
+    assert _make_filterable_registry().filter_tools('{"tools":[{"name":"unregistered"}]}') == []
+
+
 def test_lightagent_exposes_registered_tool_diagnostics():
     agent = LightAgent(
         model="gpt-4o-mini",