Repository: run-llama/llama_index
Pull request: Updated SelectionOutputParser to handle JSON objects and arrays
This commit updates the `SelectionOutputParser` class in the llama-index library to correctly handle both JSON objects and arrays. Previously, the `_marshal_llm_to_json` method was designed to extract a JSON array from a string, but it failed when the output was a JSON object, resulting in a `json.JSONDecodeError`.

The `_marshal_llm_to_json` method has been modified to check the first character of the string to determine whether it's dealing with a JSON object or array, and adjust its behavior accordingly.

Additionally, the `parse` method has been updated to handle the case where the output is a single JSON object (not a list). It now checks if the output is a dictionary, and if so, wraps it in a list.

These changes ensure that the `SelectionOutputParser` class can correctly parse both JSON objects and arrays, regardless of the number of tools used in the `LLMSingleSelector`. This should prevent the `json.JSONDecodeError` from occurring in similar scenarios in the future.

# Description

This change addresses an issue in the `SelectionOutputParser` class of the llama-index library where it was unable to correctly parse JSON objects. The `_marshal_llm_to_json` method was designed to extract a JSON array from a string, but it failed when the output was a JSON object. This resulted in a `json.JSONDecodeError: Expecting value: line 1 column 1 (char 0)` error. The motivation for this change was to ensure that the `SelectionOutputParser` class can correctly parse both JSON objects and arrays, regardless of the number of tools used in the `LLMSingleSelector`.

Fixes #6607

## Type of Change

- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] This change requires a documentation update

# How Has This Been Tested?

The changes were tested by running the `SelectionOutputParser` with both JSON objects and arrays. The tests confirmed that the class can now correctly parse both types of JSON data. Instructions for reproducing the tests are included in the bug report.

- [ ] Added new unit/integration tests
- [ ] Added new notebook (that tests end-to-end)
- [x] I stared at the code and made sure it makes sense

# Suggested Checklist:

- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes


Problem statement:
[Bug]: SelectionOutputParser incorrectly parsing Router Query Engine
### Bug Description

The `SelectionOutputParser` class in the llama-index library was unable to correctly parse JSON objects when used with a router query engine and a `LLMSingleSelector` with only one tool. The `_marshal_llm_to_json` method was designed to extract a JSON array from a string, but it failed when the output was a JSON object. This resulted in a `json.JSONDecodeError: Expecting value: line 1 column 1 (char 0)` error.

### Version

The issue was observed in llama-index version v0.6.34.post1.

### Steps to Reproduce

1. Set up a router query engine with a `LLMSingleSelector` with only one tool.
2. Use the GPT-4 model for the tool.
3. Create an instance of the `SelectionOutputParser` class.
4. Call the `parse` method on this instance, passing a string that represents a JSON object (not an array). For example: `{"choice": 1, "reason": "The provided choice does not directly address the question, but it is the only choice available. The question asks for a specific task of data structuring, which is not covered in the given choice."}`
4b. also if you attempt to ask a question that does have a tool that matches it still errors

output pre marshal {"choice": 1, "reason": "The provided choice does not directly address the question, but it is the only choice available. The question asks for a specific task of data structuring, which is not covered in the given choice."}
output post marshal 
Error parsing JSON response: Expecting value: line 1 column 1 (char 0)

### Relevant Logs/Tracbacks

```shell
{
  "type": "array",
  "items": [
    {
      "type": "object",
      "properties": {
        "choice": {
          "type": "integer",
          "value": 1
        },	
        "reason": {
          "type": "string",
          "value": "The choice is relevant as it refers to the structuring of XXX data, which is the task at hand. The task involves XXX into a specific JSON format."
        }
      },
      "required": [
        "choice",
        "reason"
      ],
      "additionalProperties": false
    }
  ]
}
[2023-06-27 13:50:39 -0400] [79292] [ERROR] Error handling request /send_message
Traceback (most recent call last):
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/gunicorn/workers/gthread.py", line 271, in handle
    keepalive = self.handle_request(req, conn)
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/gunicorn/workers/gthread.py", line 323, in handle_request
    respiter = self.wsgi(environ, resp.start_response)
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/flask/app.py", line 2213, in __call__
    return self.wsgi_app(environ, start_response)
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/flask/app.py", line 2193, in wsgi_app
    response = self.handle_exception(e)
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/flask/app.py", line 2190, in wsgi_app
    response = self.full_dispatch_request()
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/flask/app.py", line 1486, in full_dispatch_request
    rv = self.handle_user_exception(e)
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/flask/app.py", line 1484, in full_dispatch_request
    rv = self.dispatch_request()
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/flask/app.py", line 1469, in dispatch_request
    return self.ensure_sync(self.view_functions[rule.endpoint])(**view_args)
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/flask_login/utils.py", line 290, in decorated_view
    return current_app.ensure_sync(func)(*args, **kwargs)
  File "/Users/shere/Documents/GitHub/dev-ai-chat-template/app/app.py", line 348, in send_message
    response = get_bot_response(user_question, user_id, project_name, chat_history_file_path, project_attributes, direct_ai_chat)
  File "/Users/shere/Documents/GitHub/dev-ai-chat-template/app/llm_handler.py", line 284, in get_bot_response
    response = handle_index_type(
  File "/Users/shere/Documents/GitHub/dev-ai-chat-template/app/llm_handler.py", line 304, in handle_index_type
    response_obj = query_router_index(
  File "/Users/shere/Documents/GitHub/dev-ai-chat-template/app/llm_handler.py", line 1672, in query_router_index
    response = query_engine.query(user_question)
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/llama_index/indices/query/base.py", line 23, in query
    response = self._query(str_or_query_bundle)
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/llama_index/query_engine/router_query_engine.py", line 116, in _query
    result = self._selector.select(self._metadatas, query_bundle)
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/llama_index/selectors/types.py", line 76, in select
    return self._select(choices=metadatas, query=query_bundle)
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/llama_index/selectors/llm_selectors.py", line 103, in _select
    parse = self._prompt.output_parser.parse(prediction)
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/llama_index/output_parsers/selection.py", line 71, in parse
    answers = [Answer.from_dict(json_dict) for json_dict in json_list]
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/llama_index/output_parsers/selection.py", line 71, in <listcomp>
    answers = [Answer.from_dict(json_dict) for json_dict in json_list]
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/dataclasses_json/api.py", line 72, in from_dict
    return _decode_dataclass(cls, kvs, infer_missing)
  File "/Users/shere/virtualenvs/dev-ai-chat-template/lib/python3.10/site-packages/dataclasses_json/core.py", line 159, in _decode_dataclass
    field_value = kvs[field.name]
KeyError: 'choice'
```

Hints from issue discussion:
**Solution Outline**

1. **Problem Identification**: The issue was identified in the `SelectionOutputParser` class of the llama-index library. The `_marshal_llm_to_json` method was designed to extract a JSON array from a string, but it failed when the output was a JSON object. This resulted in a `json.JSONDecodeError: Expecting value: line 1 column 1 (char 0)` error.

2. **Root Cause Analysis**: The root cause of the issue was that the `_marshal_llm_to_json` method was not designed to handle JSON objects. It was only looking for JSON arrays in the input string, and returned an empty string when it couldn't find one, which caused the `json.loads` function to raise an error.

3. **Solution Design**: The solution was to modify the `_marshal_llm_to_json` method to handle both JSON objects and arrays. This involved checking the first character of the string to determine whether it's dealing with a JSON object or array, and adjusting its behavior accordingly.

4. **Implementation**: The `_marshal_llm_to_json` method was updated with the following code:

    ```python
    def _marshal_llm_to_json(self, output: str) -> str:
        output = output.strip()
        if output[0] == '{':
            left = output.find("{")
            right = output.rfind("}")
        elif output[0] == '[':
            left = output.find("[")
            right = output.rfind("]")
        else:
            raise ValueError('Invalid JSON string')
        output = output[left : right + 1]
        return output
    ```

5. **Additional Changes**: The `parse` method was also updated to handle the case where the output is a single JSON object (not a list). This was done by checking if the output is a dictionary, and if so, wrapping it in a list:

    ```python
    def parse(self, output: str) -> Any:
        output = self._marshal_llm_to_json(output)
        json_output = json.loads(output)
        if isinstance(json_output, dict):
            json_output = [json_output]
        answers = [Answer.from_dict(json_dict) for json_dict in json_output]
        return StructuredOutput(raw_output=output, parsed_output=answers)
    ```

6. **Result**: After these changes, the `SelectionOutputParser` class was able to correctly parse both JSON objects and arrays, regardless of the number of tools used in the `LLMSingleSelector`.
I think your solution makes sense, definitely feel free to raise a PR with this @sheresaidon !

Also just to note, we also have pydantic selectors that work with the new function calling API from OpenAI. These should be less prone to breaking, and I recommend trying them out :) https://gpt-index.readthedocs.io/en/latest/examples/query_engine/RouterQueryEngine.html#pydanticsingleselector

Commit message: Updated SelectionOutputParser to handle JSON objects and arrays

This commit updates the `SelectionOutputParser` class in the llama-index library to correctly handle both JSON objects and arrays. Previously, the `_marshal_llm_to_json` method was designed to extract a JSON array from a string, but it failed when the output was a JSON object, resulting in a `json.JSONDecodeError`.

The `_marshal_llm_to_json` method has been modified to check the first character of the string to determine whether it's dealing with a JSON object or array, and adjust its behavior accordingly.

Additionally, the `parse` method has been updated to handle the case where the output is a single JSON object (not a list). It now checks if the output is a dictionary, and if so, wraps it in a list.

These changes ensure that the `SelectionOutputParser` class can correctly parse both JSON objects and arrays, regardless of the number of tools used in the `LLMSingleSelector`. This should prevent the `json.JSONDecodeError` from occurring in similar scenarios in the future.

Patch to review:
diff --git a/llama_index/output_parsers/selection.py b/llama_index/output_parsers/selection.py
index 6441f8825a..e2c87f46de 100644
--- a/llama_index/output_parsers/selection.py
+++ b/llama_index/output_parsers/selection.py
@@ -48,27 +48,35 @@ class Answer(DataClassJsonMixin):
 
 class SelectionOutputParser(BaseOutputParser):
     def _marshal_llm_to_json(self, output: str) -> str:
-        """Extract a valid JSON array from a string.
-        Extracts a substring that represents a valid JSON array.
+        """Extract a valid JSON object or array from a string.
+        Extracts a substring that represents a valid JSON object or array.
 
         Args:
-            output: A string that may contain a valid JSON array surrounded by
+            output: A string that may contain a valid JSON object or array surrounded by
             extraneous characters or information.
 
         Returns:
-            A string representing a valid JSON array.
+            A string representing a valid JSON object or array.
 
         """
         output = output.strip()
-        left = output.find("[")
-        right = output.rfind("]")
+        if output[0] == '{':
+            left = output.find("{")
+            right = output.rfind("}")
+        elif output[0] == '[':
+            left = output.find("[")
+            right = output.rfind("]")
+        else:
+            raise ValueError('Invalid JSON string')
         output = output[left : right + 1]
         return output
 
     def parse(self, output: str) -> Any:
         output = self._marshal_llm_to_json(output)
-        json_list = json.loads(output)
-        answers = [Answer.from_dict(json_dict) for json_dict in json_list]
+        json_output = json.loads(output)
+        if isinstance(json_output, dict):
+            json_output = [json_output]
+        answers = [Answer.from_dict(json_dict) for json_dict in json_output]
         return StructuredOutput(raw_output=output, parsed_output=answers)
 
     def format(self, prompt_template: str) -> str:


Code review comments:
diff_hunk: @@ -48,27 +48,35 @@ class Answer(DataClassJsonMixin):
 
 class SelectionOutputParser(BaseOutputParser):
     def _marshal_llm_to_json(self, output: str) -> str:
-        """Extract a valid JSON array from a string.
-        Extracts a substring that represents a valid JSON array.
+        """Extract a valid JSON object or array from a string.
+        Extracts a substring that represents a valid JSON object or array.
 
         Args:
-            output: A string that may contain a valid JSON array surrounded by
+            output: A string that may contain a valid JSON object or array surrounded by
             extraneous characters or information.
 
         Returns:
-            A string representing a valid JSON array.
+            A string representing a valid JSON object or array.
 
         """
         output = output.strip()
-        left = output.find("[")
-        right = output.rfind("]")
+        if output[0] == '{':
path: llama_index/output_parsers/selection.py
text: @user1:
This feels a bit brittle personally, now that I look at this again, as sometimes the LLM may have extra output before the JSON starts

Maybe a better option is to find the `left` using both `[` and `{`, and whichever one is further left changes the symbol you use to find the `right`? Let me know if this makes sense

```python
left_sqaure = output.find("[")
left_brace = output.find("{")

if left_square < left_brace:
  <find right square>
else:
  <find right brace>
```