Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 27 additions & 1 deletion tasks/ai/prompts.py
Original file line number Diff line number Diff line change
Expand Up @@ -284,13 +284,39 @@ def build_ollama_tool_calling_prompt(
Fill only fields the user asked for. Return ONLY the JSON object."""


def _inject_unique_items(schema: Dict) -> Dict:
"""Add ``uniqueItems: True`` to every array-typed property recursively.

Small Ollama models can loop the same value forever in structured-output
mode; ``uniqueItems`` prevents that. Injected here rather than in the
shared ``tools.py`` schema so that only the Ollama structured-output path
is affected -- Gemini, Mistral, and native OpenAI tool-calling use the
schema as-is.
"""
if not isinstance(schema, dict):
return schema
# Snapshot keys so we can safely mutate the dict while walking it.
for key, value in list(schema.items()):
if key == 'type' and value == 'array':
schema.setdefault('uniqueItems', True)
elif isinstance(value, dict):
_inject_unique_items(value)
elif isinstance(value, list):
for item in value:
if isinstance(item, dict):
_inject_unique_items(item)
return schema
Comment thread
NeptuneHub marked this conversation as resolved.


def build_tool_calls_schema(tools: List[Dict]) -> Dict:
branches: List[Dict] = []
for t in tools:
name = t.get('name')
if not name:
continue
arg_schema = copy.deepcopy(t.get('inputSchema') or {"type": "object"})
arg_schema = _inject_unique_items(
copy.deepcopy(t.get('inputSchema') or {"type": "object"})
)
arg_schema['additionalProperties'] = False
branches.append(
{
Expand Down
7 changes: 1 addition & 6 deletions tasks/ai/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
Main Features:
* get_mcp_tools builds the schema dynamically, exposing text_match modes only when CLAP/LYRICS are enabled; tool descriptions carry the routing rules (when to use each tool and when to use a sibling instead) so they work as the primary routing signal for small models, with genre/voice/mood enums from the canonical vocab.
* execute_mcp_tool converts normalized energy 0..1 to raw score units before search_database, expands female/male voice spelling variants deterministically, passes exclude_artists/exclude_genres through as hard SQL cuts, and rejects year-only text_match queries (routing them to search_database); all failures return a generic error, never a traceback.
* Array args carry maxItems/uniqueItems caps so small-model structured output cannot loop a value forever; exclusion fields document that excluded names never go in seeds or positive filters.
* Array args carry maxItems caps so small-model structured output cannot loop a value forever; the Ollama structured-output path in prompts.py adds uniqueItems dynamically. Exclusion fields document that excluded names never go in seeds or positive filters.
"""

import logging
Expand Down Expand Up @@ -493,14 +493,12 @@ def get_mcp_tools() -> List[Dict]:
"type": "array",
"items": {"type": "string", "enum": list(GENRE_VOCAB)},
"maxItems": 5,
"uniqueItems": True,
"description": "Music genres the user WANTS, e.g. ['jazz'] or ['rock', 'blues'].",
},
"voices": {
"type": "array",
"items": {"type": "string", "enum": list(VOICE_ENUM)},
"maxItems": 2,
"uniqueItems": True,
"description": (
"Vocal type: 'female vocalists' for any female-voice request, "
"'male vocalists' for any male-voice request."
Expand All @@ -510,7 +508,6 @@ def get_mcp_tools() -> List[Dict]:
"type": "array",
"items": {"type": "string", "enum": list(config.OTHER_FEATURE_LABELS)},
"maxItems": 3,
"uniqueItems": True,
"description": "How it feels, e.g. ['sad'] or ['danceable', 'party'].",
},
"tempo_min": {"type": "number", "description": "Min BPM 40-200, e.g. 120"},
Expand Down Expand Up @@ -552,7 +549,6 @@ def get_mcp_tools() -> List[Dict]:
"type": "array",
"items": {"type": "string"},
"maxItems": 10,
"uniqueItems": True,
"description": (
"Artists the user does NOT want ('no 50 Cent' -> ['50 Cent']). "
"Hard-removed from the results; never put these in artist or seeds."
Expand All @@ -562,7 +558,6 @@ def get_mcp_tools() -> List[Dict]:
"type": "array",
"items": {"type": "string", "enum": list(GENRE_VOCAB)},
"maxItems": 5,
"uniqueItems": True,
"description": (
"Genres the user does NOT want ('no rap' -> ['Hip-Hop']). "
"Hard-removed from the results; never put these in genres."
Expand Down
15 changes: 15 additions & 0 deletions test/unit/test_mcp_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -607,6 +607,21 @@ def test_unknown_tool_returns_error(self):


class TestToolSurface:
def test_shared_schemas_omit_ollama_only_unique_items(self):
ai_mod = _import_ai_mcp_client()

def contains_unique_items(value):
if isinstance(value, dict):
return 'uniqueItems' in value or any(
contains_unique_items(child) for child in value.values()
)
if isinstance(value, list):
return any(contains_unique_items(child) for child in value)
return False

for tool in ai_mod.get_mcp_tools():
assert not contains_unique_items(tool['inputSchema'])

def test_no_llm_facing_get_songs_or_dead_params(self):
ai_mod = _import_ai_mcp_client()
for tool in ai_mod.get_mcp_tools():
Expand Down
33 changes: 33 additions & 0 deletions test/unit/test_tool_plan.py
Original file line number Diff line number Diff line change
Expand Up @@ -593,6 +593,39 @@ def test_source_schema_not_mutated(self):
pr.build_tool_calls_schema(tools)
assert 'additionalProperties' not in tools[0]['inputSchema']

def test_unique_items_are_added_only_to_derived_argument_arrays(self):
pr = _prompts()
tools = _tools_fixture()

schema = pr.build_tool_calls_schema(tools)
branches = schema['properties']['tool_calls']['items']['oneOf']
argument_schemas = [branch['properties']['arguments'] for branch in branches]

def array_schemas(value):
if isinstance(value, dict):
if value.get('type') == 'array':
yield value
for child in value.values():
yield from array_schemas(child)
elif isinstance(value, list):
for child in value:
yield from array_schemas(child)

derived_arrays = [
array_schema
for argument_schema in argument_schemas
for array_schema in array_schemas(argument_schema)
]
source_arrays = [
array_schema
for tool in tools
for array_schema in array_schemas(tool['inputSchema'])
]

assert derived_arrays
assert all(array_schema['uniqueItems'] is True for array_schema in derived_arrays)
assert all('uniqueItems' not in array_schema for array_schema in source_arrays)


class TestPromptRendering:
def test_system_prompt_derives_tools_and_rules(self):
Expand Down
Loading