Merge pull request #314 from lightspeed-core/TamiTakamiya/replace-json-loads-with-ast-literal_eval

Replace json.loads with ast.literal_eval

Author: TamiTakamiya

src/app/endpoints/streaming_query.py

@@ -1,9 +1,9 @@
 """Handler for REST API call to provide answer to streaming query."""
 
+import ast
 import json
-import logging
 import re
-from json import JSONDecodeError
+import logging
 from typing import Any, AsyncIterator, Iterator
 
 from cachetools import TTLCache  # type: ignore
@@ -362,12 +362,12 @@ def _handle_tool_execution_event(
                                 summary = summary[:newline_pos]
                         for match in METADATA_PATTERN.findall(text_content_item.text):
                             try:
-                                meta = json.loads(match.replace("'", '"'))
+                                meta = ast.literal_eval(match)
                                 if "document_id" in meta:
                                     metadata_map[meta["document_id"]] = meta
-                            except JSONDecodeError:
+                            except Exception:  # pylint: disable=broad-except
                                 logger.debug(
-                                    "JSONDecodeError was thrown in processing %s",
+                                    "An exception was thrown in processing %s",
                                     match,
                                 )
 

tests/unit/app/endpoints/test_streaming_query.py

@@ -55,19 +55,22 @@
 """,
     """Result 1
 Content: ABC
-Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Doc1', 'document_id': 'doc-1'}
+Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Doc1', 'document_id': 'doc-1', \
+'source': None}
 """,
     """Result 2
 Content: ABC
-Metadata: {'docs_url': 'https://example.com/doc2', 'title': 'Doc2', 'document_id': 'doc-2'}
+Metadata: {'docs_url': 'https://example.com/doc2', 'title': 'Doc2', 'document_id': 'doc-2', \
+'source': None}
 """,
     """END of knowledge_search tool results.
 """,
     # Following metadata contains an intentionally incorrect keyword "Title" (instead of "title")
     # and it is not picked as a referenced document.
     """Result 3
 Content: ABC
-Metadata: {'docs_url': 'https://example.com/doc3', 'Title': 'Doc3', 'document_id': 'doc-3'}
+Metadata: {'docs_url': 'https://example.com/doc3', 'Title': 'Doc3', 'document_id': 'doc-3', \
+'source': None}
 """,
     """The above results were retrieved to help answer the user\'s query: "Sample Query".
 Use them as supporting information only in answering this query.