diff --git a/graphiti_core/llm_client/openai_generic_client.py b/graphiti_core/llm_client/openai_generic_client.py index d462c4e3dc..010eb5b6ad 100644 --- a/graphiti_core/llm_client/openai_generic_client.py +++ b/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) diff --git a/tests/llm_client/test_openai_generic_client.py b/tests/llm_client/test_openai_generic_client.py index 3355f25164..cbefa6efa6 100644 --- a/tests/llm_client/test_openai_generic_client.py +++ b/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