9611202221

Committed 21 Jun 2024 09:10AM UTC coverage: 89.953% (-0.1%) from 90.059%

Build # 9611202221

Build Type

push

github

Committed by

web-flow

Commit Message

fix(JsonSchemaValidator): fix recursive loop and general LLM (claude, mistral...) compatibility (#7556)

* Feat: Fix recursive conversion in JsonSchemaValidator (autofix generated by ClaudeOpus). Modify the behaviour to build the error template in a single user_message instead of two separate. Modify the behaviour to only include latest message instead of full history (very costly if long looping pipeline)

* Feat: Fix recursive conversion in JsonSchemaValidator (autofix generated by ClaudeOpus). Modify the behaviour to build the error template in a single user_message instead of two separate. Modify the behaviour to only include latest message instead of full history (very costly if long looping pipeline)

* reno

* fix test

* Verify provided message contains JSON object to begin with

* Minor detail

---------

Co-authored-by: Vladimir Blagojevic <dovlex@gmail.com>

Run Details

6912 of 7684 relevant lines covered (89.95%)

0.9 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

haystack/components/validators/json_schema.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import json
from typing import Any, Dict, List, Optional

from haystack import component
from haystack.dataclasses import ChatMessage
from haystack.lazy_imports import LazyImport

with LazyImport(message="Run 'pip install jsonschema'") as jsonschema_import:
    from jsonschema import ValidationError, validate


def is_valid_json(s: str) -> bool:
    """
    Check if the provided string is a valid JSON.

    :param s: The string to be checked.
    :returns: `True` if the string is a valid JSON; otherwise, `False`.
    """
    try:
        json.loads(s)
    except ValueError:
        return False
    return True


@component
class JsonSchemaValidator:
    """
    Validates JSON content of `ChatMessage` against a specified [JSON Schema](https://json-schema.org/).

    If JSON content of a message conforms to the provided schema, the message is passed along the "validated" output.
    If the JSON content does not conform to the schema, the message is passed along the "validation_error" output.
    In the latter case, the error message is constructed using the provided `error_template` or a default template.
    These error ChatMessages can be used by LLMs in Haystack 2.x recovery loops.

    Usage example:

    ```python
    from typing import List

    from haystack import Pipeline
    from haystack.components.generators.chat import OpenAIChatGenerator
    from haystack.components.others import Multiplexer
    from haystack.components.validators import JsonSchemaValidator
    from haystack import component
    from haystack.dataclasses import ChatMessage


    @component
    class MessageProducer:

        @component.output_types(messages=List[ChatMessage])
        def run(self, messages: List[ChatMessage]) -> dict:
            return {"messages": messages}


    p = Pipeline()
    p.add_component("llm", OpenAIChatGenerator(model="gpt-4-1106-preview",
                                               generation_kwargs={"response_format": {"type": "json_object"}}))
    p.add_component("schema_validator", JsonSchemaValidator())
    p.add_component("mx_for_llm", Multiplexer(List[ChatMessage]))
    p.add_component("message_producer", MessageProducer())

    p.connect("message_producer.messages", "mx_for_llm")
    p.connect("mx_for_llm", "llm")
    p.connect("llm.replies", "schema_validator.messages")
    p.connect("schema_validator.validation_error", "mx_for_llm")

    result = p.run(data={
        "message_producer": {
            "messages":[ChatMessage.from_user("Generate JSON for person with name 'John' and age 30")]},
            "schema_validator": {
                "json_schema": {
                    "type": "object",
                    "properties": {"name": {"type": "string"},
                    "age": {"type": "integer"}
                }
            }
        }
    })
    print(result)
    >> {'schema_validator': {'validated': [ChatMessage(content='\\n{\\n  "name": "John",\\n  "age": 30\\n}',
    role=<ChatRole.ASSISTANT: 'assistant'>, name=None, meta={'model': 'gpt-4-1106-preview', 'index': 0,
    'finish_reason': 'stop', 'usage': {'completion_tokens': 17, 'prompt_tokens': 20, 'total_tokens': 37}})]}}
    ```
    """

    # Default error description template
    default_error_template = (
        "The following generated JSON does not conform to the provided schema.\n"
        "Generated JSON: {failing_json}\n"
        "Error details:\n- Message: {error_message}\n"
        "- Error Path in JSON: {error_path}\n"
        "- Schema Path: {error_schema_path}\n"
        "Please match the following schema:\n"
        "{json_schema}\n"
        "and provide the corrected JSON content ONLY. Please do not output anything else than the raw corrected "
        "JSON string, this is the most important part of the task. Don't use any markdown and don't add any comment."
    )

    def __init__(self, json_schema: Optional[Dict[str, Any]] = None, error_template: Optional[str] = None):
        """
        Initialize the JsonSchemaValidator component.

        :param json_schema: A dictionary representing the [JSON schema](https://json-schema.org/) against which
            the messages' content is validated.
        :param error_template: A custom template string for formatting the error message in case of validation failure.
        """
        jsonschema_import.check()
        self.json_schema = json_schema
        self.error_template = error_template

    @component.output_types(validated=List[ChatMessage], validation_error=List[ChatMessage])
    def run(
        self,
        messages: List[ChatMessage],
        json_schema: Optional[Dict[str, Any]] = None,
        error_template: Optional[str] = None,
    ) -> Dict[str, List[ChatMessage]]:
        """
        Validates the last of the provided messages against the specified json schema.

        If it does, the message is passed along the "validated" output. If it does not, the message is passed along
        the "validation_error" output.

        :param messages: A list of ChatMessage instances to be validated. The last message in this list is the one
            that is validated.
        :param json_schema: A dictionary representing the [JSON schema](https://json-schema.org/)
            against which the messages' content is validated. If not provided, the schema from the component init
            is used.
        :param error_template: A custom template string for formatting the error message in case of validation. If not
            provided, the `error_template` from the component init is used.
        :return:  A dictionary with the following keys:
            - "validated": A list of messages if the last message is valid.
            - "validation_error": A list of messages if the last message is invalid.
        :raises ValueError: If no JSON schema is provided or if the message content is not a dictionary or a list of
            dictionaries.
        """
        last_message = messages[-1]
        if not is_valid_json(last_message.content):
            return {
                "validation_error": [
                    ChatMessage.from_user(
                        f"The message '{last_message.content}' is not a valid JSON object. "
                        f"Please provide only a valid JSON object in string format."
                        f"Don't use any markdown and don't add any comment."
                    )
                ]
            }

        last_message_content = json.loads(last_message.content)
        json_schema = json_schema or self.json_schema
        error_template = error_template or self.error_template or self.default_error_template

        if not json_schema:
            raise ValueError("Provide a JSON schema for validation either in the run method or in the component init.")
        # fc payload is json object but subtree `parameters` is string - we need to convert to json object
        # we need complete json to validate it against schema
        last_message_json = self._recursive_json_to_object(last_message_content)
        using_openai_schema: bool = self._is_openai_function_calling_schema(json_schema)
        if using_openai_schema:
            validation_schema = json_schema["parameters"]
        else:
            validation_schema = json_schema
        try:
            last_message_json = [last_message_json] if not isinstance(last_message_json, list) else last_message_json
            for content in last_message_json:
                if using_openai_schema:
                    validate(instance=content["function"]["arguments"], schema=validation_schema)
                else:
                    validate(instance=content, schema=validation_schema)

            return {"validated": [last_message]}
        except ValidationError as e:
            error_path = " -> ".join(map(str, e.absolute_path)) if e.absolute_path else "N/A"
            error_schema_path = " -> ".join(map(str, e.absolute_schema_path)) if e.absolute_schema_path else "N/A"

            error_template = error_template or self.default_error_template

            recovery_prompt = self._construct_error_recovery_message(
                error_template,
                str(e),
                error_path,
                error_schema_path,
                validation_schema,
                failing_json=last_message.content,
            )
            return {"validation_error": [ChatMessage.from_user(recovery_prompt)]}

    def _construct_error_recovery_message(
        self,
        error_template: str,
        error_message: str,
        error_path: str,
        error_schema_path: str,
        json_schema: Dict[str, Any],
        failing_json: str,
    ) -> str:
        """
        Constructs an error recovery message using a specified template or the default one if none is provided.

        :param error_template: A custom template string for formatting the error message in case of validation failure.
        :param error_message: The error message returned by the JSON schema validator.
        :param error_path: The path in the JSON content where the error occurred.
        :param error_schema_path: The path in the JSON schema where the error occurred.
        :param json_schema: The JSON schema against which the content is validated.
        :param failing_json: The generated invalid JSON string.
        """
        error_template = error_template or self.default_error_template

        return error_template.format(
            error_message=error_message,
            error_path=error_path,
            error_schema_path=error_schema_path,
            json_schema=json_schema,
            failing_json=failing_json,
        )

    def _is_openai_function_calling_schema(self, json_schema: Dict[str, Any]) -> bool:
        """
        Checks if the provided schema is a valid OpenAI function calling schema.

        :param json_schema: The JSON schema to check
        :return: `True` if the schema is a valid OpenAI function calling schema; otherwise, `False`.
        """
        return all(key in json_schema for key in ["name", "description", "parameters"])

    def _recursive_json_to_object(self, data: Any) -> Any:
        """
        Convert any string values that are valid JSON objects into dictionary objects.

        Returns a new data structure.

        :param data: The data structure to be traversed.
        :return: A new data structure with JSON strings converted to dictionary objects.
        """
        if isinstance(data, list):
            return [self._recursive_json_to_object(item) for item in data]

        if isinstance(data, dict):
            new_dict = {}
            for key, value in data.items():
                if isinstance(value, str):
                    try:
                        json_value = json.loads(value)
                        if isinstance(json_value, (dict, list)):
                            new_dict[key] = self._recursive_json_to_object(json_value)
                        else:
                            new_dict[key] = value  # Preserve the original string value
                    except json.JSONDecodeError:
                        new_dict[key] = value
                elif isinstance(value, dict):
                    new_dict[key] = self._recursive_json_to_object(value)
                else:
                    new_dict[key] = value
            return new_dict

        # If it's neither a list nor a dictionary, return the value directly
        raise ValueError("Input must be a dictionary or a list of dictionaries.")

1	# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	import json	×
6	from typing import Any, Dict, List, Optional	×
7
8	from haystack import component	×
9	from haystack.dataclasses import ChatMessage	×
10	from haystack.lazy_imports import LazyImport	×
11
12	with LazyImport(message="Run 'pip install jsonschema'") as jsonschema_import:	×
13	from jsonschema import ValidationError, validate	×
14
15
16	def is_valid_json(s: str) -> bool:	×
17	"""
18	Check if the provided string is a valid JSON.
19
20	:param s: The string to be checked.
21	:returns: `True` if the string is a valid JSON; otherwise, `False`.
22	"""
23	try:	×
24	json.loads(s)	×
25	except ValueError:	×
26	return False	×
27	return True	×
28
29
30	@component	×
31	class JsonSchemaValidator:	×
32	"""
33	Validates JSON content of `ChatMessage` against a specified [JSON Schema](https://json-schema.org/).
34
35	If JSON content of a message conforms to the provided schema, the message is passed along the "validated" output.
36	If the JSON content does not conform to the schema, the message is passed along the "validation_error" output.
37	In the latter case, the error message is constructed using the provided `error_template` or a default template.
38	These error ChatMessages can be used by LLMs in Haystack 2.x recovery loops.
39
40	Usage example:
41
42	```python
43	from typing import List
44
45	from haystack import Pipeline
46	from haystack.components.generators.chat import OpenAIChatGenerator
47	from haystack.components.others import Multiplexer
48	from haystack.components.validators import JsonSchemaValidator
49	from haystack import component
50	from haystack.dataclasses import ChatMessage
51
52
53	@component
54	class MessageProducer:
55
56	@component.output_types(messages=List[ChatMessage])
57	def run(self, messages: List[ChatMessage]) -> dict:
58	return {"messages": messages}
59
60
61	p = Pipeline()
62	p.add_component("llm", OpenAIChatGenerator(model="gpt-4-1106-preview",
63	generation_kwargs={"response_format": {"type": "json_object"}}))
64	p.add_component("schema_validator", JsonSchemaValidator())
65	p.add_component("mx_for_llm", Multiplexer(List[ChatMessage]))
66	p.add_component("message_producer", MessageProducer())
67
68	p.connect("message_producer.messages", "mx_for_llm")
69	p.connect("mx_for_llm", "llm")
70	p.connect("llm.replies", "schema_validator.messages")
71	p.connect("schema_validator.validation_error", "mx_for_llm")
72
73	result = p.run(data={
74	"message_producer": {
75	"messages":[ChatMessage.from_user("Generate JSON for person with name 'John' and age 30")]},
76	"schema_validator": {
77	"json_schema": {
78	"type": "object",
79	"properties": {"name": {"type": "string"},
80	"age": {"type": "integer"}
81	}
82	}
83	}
84	})
85	print(result)
86	>> {'schema_validator': {'validated': [ChatMessage(content='\\n{\\n "name": "John",\\n "age": 30\\n}',
87	role=<ChatRole.ASSISTANT: 'assistant'>, name=None, meta={'model': 'gpt-4-1106-preview', 'index': 0,
88	'finish_reason': 'stop', 'usage': {'completion_tokens': 17, 'prompt_tokens': 20, 'total_tokens': 37}})]}}
89	```
90	"""
91
92	# Default error description template
93	default_error_template = (	×
94	"The following generated JSON does not conform to the provided schema.\n"
95	"Generated JSON: {failing_json}\n"
96	"Error details:\n- Message: {error_message}\n"
97	"- Error Path in JSON: {error_path}\n"
98	"- Schema Path: {error_schema_path}\n"
99	"Please match the following schema:\n"
100	"{json_schema}\n"
101	"and provide the corrected JSON content ONLY. Please do not output anything else than the raw corrected "
102	"JSON string, this is the most important part of the task. Don't use any markdown and don't add any comment."
103	)
104
105	def __init__(self, json_schema: Optional[Dict[str, Any]] = None, error_template: Optional[str] = None):	×
106	"""
107	Initialize the JsonSchemaValidator component.
108
109	:param json_schema: A dictionary representing the [JSON schema](https://json-schema.org/) against which
110	the messages' content is validated.
111	:param error_template: A custom template string for formatting the error message in case of validation failure.
112	"""
113	jsonschema_import.check()	×
114	self.json_schema = json_schema	×
115	self.error_template = error_template	×
116
117	@component.output_types(validated=List[ChatMessage], validation_error=List[ChatMessage])	×
118	def run(	×
119	self,
120	messages: List[ChatMessage],
121	json_schema: Optional[Dict[str, Any]] = None,
122	error_template: Optional[str] = None,
123	) -> Dict[str, List[ChatMessage]]:
124	"""
125	Validates the last of the provided messages against the specified json schema.
126
127	If it does, the message is passed along the "validated" output. If it does not, the message is passed along
128	the "validation_error" output.
129
130	:param messages: A list of ChatMessage instances to be validated. The last message in this list is the one
131	that is validated.
132	:param json_schema: A dictionary representing the [JSON schema](https://json-schema.org/)
133	against which the messages' content is validated. If not provided, the schema from the component init
134	is used.
135	:param error_template: A custom template string for formatting the error message in case of validation. If not
136	provided, the `error_template` from the component init is used.
137	:return: A dictionary with the following keys:
138	- "validated": A list of messages if the last message is valid.
139	- "validation_error": A list of messages if the last message is invalid.
140	:raises ValueError: If no JSON schema is provided or if the message content is not a dictionary or a list of
141	dictionaries.
142	"""
143	last_message = messages[-1]	×
144	if not is_valid_json(last_message.content):	×
145	return {	×
146	"validation_error": [
147	ChatMessage.from_user(
148	f"The message '{last_message.content}' is not a valid JSON object. "
149	f"Please provide only a valid JSON object in string format."
150	f"Don't use any markdown and don't add any comment."
151	)
152	]
153	}
154
155	last_message_content = json.loads(last_message.content)	×
156	json_schema = json_schema or self.json_schema	×
157	error_template = error_template or self.error_template or self.default_error_template	×
158
159	if not json_schema:	×
160	raise ValueError("Provide a JSON schema for validation either in the run method or in the component init.")	×
161	# fc payload is json object but subtree `parameters` is string - we need to convert to json object
162	# we need complete json to validate it against schema
163	last_message_json = self._recursive_json_to_object(last_message_content)	×
164	using_openai_schema: bool = self._is_openai_function_calling_schema(json_schema)	×
165	if using_openai_schema:	×
166	validation_schema = json_schema["parameters"]	×
167	else:
168	validation_schema = json_schema	×
169	try:	×
170	last_message_json = [last_message_json] if not isinstance(last_message_json, list) else last_message_json	×
171	for content in last_message_json:	×
172	if using_openai_schema:	×
173	validate(instance=content["function"]["arguments"], schema=validation_schema)	×
174	else:
175	validate(instance=content, schema=validation_schema)	×
176
177	return {"validated": [last_message]}	×
178	except ValidationError as e:	×
179	error_path = " -> ".join(map(str, e.absolute_path)) if e.absolute_path else "N/A"	×
180	error_schema_path = " -> ".join(map(str, e.absolute_schema_path)) if e.absolute_schema_path else "N/A"	×
181
182	error_template = error_template or self.default_error_template	×
183
184	recovery_prompt = self._construct_error_recovery_message(	×
185	error_template,
186	str(e),
187	error_path,
188	error_schema_path,
189	validation_schema,
190	failing_json=last_message.content,
191	)
192	return {"validation_error": [ChatMessage.from_user(recovery_prompt)]}	×
193
194	def _construct_error_recovery_message(	×
195	self,
196	error_template: str,
197	error_message: str,
198	error_path: str,
199	error_schema_path: str,
200	json_schema: Dict[str, Any],
201	failing_json: str,
202	) -> str:
203	"""
204	Constructs an error recovery message using a specified template or the default one if none is provided.
205
206	:param error_template: A custom template string for formatting the error message in case of validation failure.
207	:param error_message: The error message returned by the JSON schema validator.
208	:param error_path: The path in the JSON content where the error occurred.
209	:param error_schema_path: The path in the JSON schema where the error occurred.
210	:param json_schema: The JSON schema against which the content is validated.
211	:param failing_json: The generated invalid JSON string.
212	"""
213	error_template = error_template or self.default_error_template	×
214
215	return error_template.format(	×
216	error_message=error_message,
217	error_path=error_path,
218	error_schema_path=error_schema_path,
219	json_schema=json_schema,
220	failing_json=failing_json,
221	)
222
223	def _is_openai_function_calling_schema(self, json_schema: Dict[str, Any]) -> bool:	×
224	"""
225	Checks if the provided schema is a valid OpenAI function calling schema.
226
227	:param json_schema: The JSON schema to check
228	:return: `True` if the schema is a valid OpenAI function calling schema; otherwise, `False`.
229	"""
230	return all(key in json_schema for key in ["name", "description", "parameters"])	×
231
232	def _recursive_json_to_object(self, data: Any) -> Any:	×
233	"""
234	Convert any string values that are valid JSON objects into dictionary objects.
235
236	Returns a new data structure.
237
238	:param data: The data structure to be traversed.
239	:return: A new data structure with JSON strings converted to dictionary objects.
240	"""
241	if isinstance(data, list):	×
242	return [self._recursive_json_to_object(item) for item in data]	×
243
244	if isinstance(data, dict):	×
245	new_dict = {}	×
246	for key, value in data.items():	×
247	if isinstance(value, str):	×
248	try:	×
249	json_value = json.loads(value)	×
250	if isinstance(json_value, (dict, list)):	×
251	new_dict[key] = self._recursive_json_to_object(json_value)	×
252	else:
253	new_dict[key] = value # Preserve the original string value	×
254	except json.JSONDecodeError:	×
255	new_dict[key] = value	×
256	elif isinstance(value, dict):	×
257	new_dict[key] = self._recursive_json_to_object(value)	×
258	else:
259	new_dict[key] = value	×
260	return new_dict	×
261
262	# If it's neither a list nor a dictionary, return the value directly
263	raise ValueError("Input must be a dictionary or a list of dictionaries.")	×

deepset-ai / haystack / 9611202221

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous