fix(llm): json_object guidance as example instance, not raw schema paste

graphiti_core/llm_client/openai_generic_client.py

@@ -37,6 +37,110 @@
 StructuredOutputMode = Literal['json_schema', 'json_object']
 
 
+def _resolve(schema: dict[str, Any], defs: dict[str, Any]) -> dict[str, Any]:
+    """Follow a top-level $ref to its $defs target (one hop is all pydantic emits)."""
+    if '$ref' in schema:
+        return defs.get(schema['$ref'].split('/')[-1], {})
+    return schema
+
+
+def _example_value(schema: dict[str, Any], defs: dict[str, Any], depth: int = 0) -> Any:
+    """Minimal plausible instance value for a JSON-Schema fragment."""
+    if depth > 6:
+        return None
+    if '$ref' in schema:
+        return _example_value(_resolve(schema, defs), defs, depth + 1)
+    if schema.get('enum'):
+        return schema['enum'][0]
+    for combinator in ('anyOf', 'oneOf', 'allOf'):
+        options = schema.get(combinator)
+        if options:
+            non_null = [o for o in options if o.get('type') != 'null']
+            return _example_value(non_null[0] if non_null else options[0], defs, depth + 1)
+    schema_type = schema.get('type')
+    if schema_type == 'array':
+        # Arrays of objects get one exemplar item so the model sees the item shape
+        # (the old schema paste carried it via $defs). Arrays of scalars stay empty —
+        # the shape is obvious, and a scalar exemplar like [0] reads as a suggested
+        # answer rather than a shape hint.
+        items = _resolve(schema.get('items', {}), defs)
+        if items.get('type') == 'object' or 'properties' in items:
+            return [_example_value(items, defs, depth + 1)]
+        return []
+    if schema_type == 'string':
+        return ''
+    if schema_type == 'integer':
+        return 0
+    if schema_type == 'number':
+        return 0
+    if schema_type == 'boolean':
+        return False
+    if schema_type == 'object' or 'properties' in schema:
+        return {
+            name: _example_value(sub, defs, depth + 1)
+            for name, sub in schema.get('properties', {}).items()
+        }
+    return None
+
+
+def example_instruction_for_model(response_model: type[BaseModel]) -> str:
+    """Prompt tail describing the expected reply as an example instance, not a schema.
+
+    Appending ``model_json_schema()`` verbatim leads weaker ``json_object``
+    providers (e.g. DeepSeek) to occasionally echo the schema wrapper itself
+    back — responses whose top-level keys are ``properties``/``required`` —
+    which then fail response-model validation and abort the whole episode.
+    An example-shaped skeleton plus a plain-text field list gives the model
+    nothing schema-shaped to parrot. Measured on deepseek-v4-flash with the
+    ``dedupe_edges.resolve_edge`` prompt: ~10% of calls failed with the schema
+    paste vs ~0-2% with the example form (N=20 per arm, temperature 1).
+    """
+    schema = response_model.model_json_schema()
+    defs = schema.get('$defs', {})
+    properties: dict[str, Any] = schema.get('properties', {})
+    required = set(schema.get('required', properties.keys()))
+    example = {name: _example_value(sub, defs) for name, sub in properties.items()}
+
+    def describe(name: str, sub: dict[str, Any], req: set[str], indent: str) -> list[str]:
+        resolved = _resolve(sub, defs)
+        field_type = sub.get('type', resolved.get('type', 'object'))
+        line = f'{indent}- {name} ({field_type}, {"required" if name in req else "optional"})'
+        description = sub.get('description', resolved.get('description', ''))
+        if description:
+            line += f': {description}'
+        lines = [line]
+        # For arrays of objects, describe the item's fields one level deep so the
+        # model knows the item keys (the old schema paste carried this via $defs).
+        if field_type == 'array':
+            items = _resolve(resolved.get('items', sub.get('items', {})), defs)
+            item_props = items.get('properties')
+            if item_props:
+                item_required = set(items.get('required', item_props.keys()))
+                lines.append(f'{indent}  each item has:')
+                for item_name, item_sub in item_props.items():
+                    item_resolved = _resolve(item_sub, defs)
+                    item_type = item_sub.get('type', item_resolved.get('type', 'object'))
+                    item_line = (
+                        f'{indent}  - {item_name} ({item_type}, '
+                        f'{"required" if item_name in item_required else "optional"})'
+                    )
+                    item_desc = item_sub.get('description', item_resolved.get('description', ''))
+                    if item_desc:
+                        item_line += f': {item_desc}'
+                    lines.append(item_line)
+        return lines
+
+    field_lines: list[str] = []
+    for name, sub in properties.items():
+        field_lines.extend(describe(name, sub, required, ''))
+    return (
+        '\n\nRespond ONLY with a JSON object of exactly this shape — fill in the values; '
+        'it is the answer instance, not a schema (do not output "properties" or '
+        '"required" keys). Arrays may contain zero or more items:\n\n'
+        f'{json.dumps(example)}\n\nFields:\n' + '\n'.join(field_lines)
+    )
+
+
 class OpenAIGenericClient(LLMClient):
     """
     OpenAIClient is a client class for interacting with OpenAI's language models.
@@ -188,16 +292,13 @@ async def generate_response(
         if max_tokens is None:
             max_tokens = self.max_tokens
 
-        # In json_object fallback mode the API does not enforce the schema, so embed it in
-        # the prompt to guide the model. In json_schema mode the schema is enforced via
-        # response_format, so no prompt injection is needed.
+        # In json_object fallback mode the API does not enforce the schema, so guide the
+        # model with an example-shaped instruction. Deliberately NOT the raw JSON Schema:
+        # weaker providers (e.g. DeepSeek) sometimes echo a pasted schema back verbatim
+        # ('properties'/'required' as top-level keys), failing validation. In json_schema
+        # mode the schema is enforced via response_format, so no prompt injection is needed.
         if response_model is not None and self.structured_output_mode == 'json_object':
-            serialized_model = json.dumps(response_model.model_json_schema())
-            messages[
-                -1
-            ].content += (
-                f'\n\nRespond with a JSON object in the following format:\n\n{serialized_model}'
-            )
+            messages[-1].content += example_instruction_for_model(response_model)
 
         # Add multilingual extraction instructions
         messages[0].content += get_extraction_language_instruction(group_id)

tests/llm_client/test_openai_generic_client.py

@@ -77,21 +77,24 @@ async def test_json_schema_mode_does_not_inject_schema_into_prompt():
     await client.generate_response(messages, response_model=ResponseModel)
 
     sent_user_content = completions.create_calls[0]['messages'][-1]['content']
-    assert 'Respond with a JSON object in the following format' not in sent_user_content
+    assert 'Respond ONLY with a JSON object of exactly this shape' not in sent_user_content
 
 
 @pytest.mark.asyncio
-async def test_json_object_mode_uses_json_object_and_injects_schema():
+async def test_json_object_mode_injects_example_instance_not_schema():
     client, completions = _make_client(structured_output_mode='json_object')
 
     await client.generate_response(_messages(), response_model=ResponseModel)
 
     call = completions.create_calls[0]
     assert call['response_format'] == {'type': 'json_object'}
-    # The schema must be injected into the prompt since the API will not enforce it.
+    # The API does not enforce the schema in json_object mode, so the prompt must carry
+    # guidance — but as an example instance, NOT the raw schema (weak providers echo a
+    # pasted schema back verbatim: 'properties'/'required' as top-level keys).
     sent_user_content = call['messages'][-1]['content']
-    assert 'Respond with a JSON object in the following format' in sent_user_content
-    assert json.dumps(ResponseModel.model_json_schema()) in sent_user_content
+    assert 'Respond ONLY with a JSON object of exactly this shape' in sent_user_content
+    assert json.dumps({'foo': ''}) in sent_user_content
+    assert json.dumps(ResponseModel.model_json_schema()) not in sent_user_content
 
 
 @pytest.mark.asyncio
@@ -103,7 +106,8 @@ async def test_no_response_model_uses_json_object_without_injection():
     call = completions.create_calls[0]
     assert call['response_format'] == {'type': 'json_object'}
     assert (
-        'Respond with a JSON object in the following format' not in call['messages'][-1]['content']
+        'Respond ONLY with a JSON object of exactly this shape'
+        not in call['messages'][-1]['content']
     )
     assert result == {'any': 'thing'}
 
@@ -156,6 +160,95 @@ async def test_strips_markdown_code_fence_before_parsing():
     assert result == {'foo': 'bar'}
 
 
+def _extract_example_json(instruction: str) -> dict:
+    """The example instance is the only line in the instruction starting with '{'."""
+    for line in instruction.splitlines():
+        if line.startswith('{'):
+            return json.loads(line)
+    raise AssertionError('no example JSON line found in instruction')
+
+
+def test_example_instruction_is_instance_shaped_for_edge_duplicate():
+    from graphiti_core.llm_client.openai_generic_client import example_instruction_for_model
+    from graphiti_core.prompts.dedupe_edges import EdgeDuplicate
+
+    instruction = example_instruction_for_model(EdgeDuplicate)
+    example = _extract_example_json(instruction)
+    assert set(example.keys()) == set(EdgeDuplicate.model_fields.keys())
+    # The example itself must validate against the model.
+    EdgeDuplicate.model_validate(example)
+    # Field descriptions carried over as plain text for semantic guidance.
+    assert 'duplicate_facts' in instruction
+    assert 'EXISTING FACTS' in instruction
+
+
+def test_example_instruction_recurses_nested_models_and_optionals():
+    from pydantic import Field
+
+    from graphiti_core.llm_client.openai_generic_client import example_instruction_for_model
+
+    class Nested(BaseModel):
+        label: str = Field(..., description='a label')
+        score: float = Field(..., description='a score')
+
+    class Outer(BaseModel):
+        items: list[int] = Field(..., description='list of idx values')
+        nested: Nested = Field(..., description='nested model')
+        maybe: str | None = Field(None, description='optional field')
+
+    example = _extract_example_json(example_instruction_for_model(Outer))
+    assert example['items'] == []
+    assert example['nested'] == {'label': '', 'score': 0}
+    Outer.model_validate(example)
+
+
+def test_example_value_handles_enum_anyof_and_scalars():
+    from graphiti_core.llm_client.openai_generic_client import _example_value
+
+    assert _example_value({'enum': ['a', 'b']}, {}) == 'a'
+    assert _example_value({'anyOf': [{'type': 'null'}, {'type': 'integer'}]}, {}) == 0
+    assert _example_value({'type': 'boolean'}, {}) is False
+
+
+def test_array_of_objects_gets_exemplar_item_and_field_docs():
+    # The main extraction models are arrays of objects; the model needs the item
+    # shape (the old schema paste carried it via $defs). Scalar arrays stay [].
+    from graphiti_core.llm_client.openai_generic_client import example_instruction_for_model
+    from graphiti_core.prompts.extract_nodes import ExtractedEntities
+
+    instruction = example_instruction_for_model(ExtractedEntities)
+    example = _extract_example_json(instruction)
+    assert len(example['extracted_entities']) == 1
+    item = example['extracted_entities'][0]
+    assert set(item.keys()) == {'name', 'entity_type_id', 'episode_indices'}
+    ExtractedEntities.model_validate(example)
+    # Item fields documented one level deep.
+    assert 'each item has:' in instruction
+    assert 'entity_type_id' in instruction
+    # Scalar-array fields keep an empty-list example (no suggested answer bias).
+    from graphiti_core.prompts.dedupe_edges import EdgeDuplicate
+
+    dedupe_example = _extract_example_json(example_instruction_for_model(EdgeDuplicate))
+    assert dedupe_example == {'duplicate_facts': [], 'contradicted_facts': []}
+
+
+def test_extracted_edges_exemplar_carries_item_shape():
+    # The other high-volume extraction model named in review: ExtractedEdges is
+    # edges: list[Edge] where Edge mixes required strs, optional nullable strs,
+    # and a scalar int array — the exemplar must validate and document all keys.
+    from graphiti_core.llm_client.openai_generic_client import example_instruction_for_model
+    from graphiti_core.prompts.extract_edges import Edge, ExtractedEdges
+
+    instruction = example_instruction_for_model(ExtractedEdges)
+    example = _extract_example_json(instruction)
+    assert len(example['edges']) == 1
+    item = example['edges'][0]
+    assert set(item.keys()) == set(Edge.model_fields.keys())
+    ExtractedEdges.model_validate(example)
+    for key in ('source_entity_name', 'target_entity_name', 'relation_type', 'fact'):
+        assert key in instruction
+
+
 @pytest.mark.asyncio
 async def test_non_retryable_error_is_not_retried():
     # The old hand-rolled re-prompt loop is gone. Retry is now delegated to the base