From 13088ea8c5f9f2432116a5cce284fb18d3ac5b8e Mon Sep 17 00:00:00 2001 From: vedantsahai18 Date: Tue, 31 Dec 2024 16:56:54 -0500 Subject: [PATCH 1/4] feat(agents-api): Add an extract_json() custom function --- agents-api/agents_api/activities/utils.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/agents-api/agents_api/activities/utils.py b/agents-api/agents_api/activities/utils.py index 3094e2e78..4f18c3bb9 100644 --- a/agents-api/agents_api/activities/utils.py +++ b/agents-api/agents_api/activities/utils.py @@ -103,6 +103,15 @@ def chunk_doc(string: str) -> list[str]: return [" ".join([sent.text for sent in chunk]) for chunk in doc._.chunks] +def safe_extract_json(string: str) -> dict: + if len(string) > MAX_STRING_LENGTH: + msg = f"String exceeds maximum length of {MAX_STRING_LENGTH}" + raise ValueError(msg) + extarct_String = string[string.find("```json")+7:string.find("```", string.find("```json")+7)] + return json.loads(extarct_String) + + + # Restricted set of allowed functions ALLOWED_FUNCTIONS = { # Basic Python builtins @@ -131,6 +140,7 @@ def chunk_doc(string: str) -> list[str]: "load_yaml": safe_yaml_load, "dump_json": json.dumps, "dump_yaml": yaml.dump, + "extract_json": safe_extract_json, # Regex and NLP functions (using re2 which is safe against ReDoS) "search_regex": lambda pattern, string: re2.search(pattern, string), "match_regex": lambda pattern, string: bool(re2.fullmatch(pattern, string)), From e03b16ee813ef1bb1444503959c16c9f68f5f275 Mon Sep 17 00:00:00 2001 From: vedantsahai18 Date: Tue, 31 Dec 2024 17:01:16 -0500 Subject: [PATCH 2/4] chore: typo fix --- agents-api/agents_api/activities/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/agents-api/agents_api/activities/utils.py b/agents-api/agents_api/activities/utils.py index 4f18c3bb9..effc23600 100644 --- a/agents-api/agents_api/activities/utils.py +++ b/agents-api/agents_api/activities/utils.py @@ -103,12 +103,12 @@ def chunk_doc(string: str) -> list[str]: return [" ".join([sent.text for sent in chunk]) for chunk in doc._.chunks] -def safe_extract_json(string: str) -> dict: +def safe_extract_json(string: str): if len(string) > MAX_STRING_LENGTH: msg = f"String exceeds maximum length of {MAX_STRING_LENGTH}" raise ValueError(msg) - extarct_String = string[string.find("```json")+7:string.find("```", string.find("```json")+7)] - return json.loads(extarct_String) + extracted_string = string[string.find("```json")+7:string.find("```", string.find("```json")+7)] + return json.loads(extracted_string) From 2da2c3905907f8c7fa7312ebe4ef40716a7b8af1 Mon Sep 17 00:00:00 2001 From: Vedantsahai18 Date: Tue, 31 Dec 2024 22:02:27 +0000 Subject: [PATCH 3/4] refactor: Lint agents-api (CI) --- agents-api/agents_api/activities/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/agents-api/agents_api/activities/utils.py b/agents-api/agents_api/activities/utils.py index effc23600..8123490fe 100644 --- a/agents-api/agents_api/activities/utils.py +++ b/agents-api/agents_api/activities/utils.py @@ -107,11 +107,12 @@ def safe_extract_json(string: str): if len(string) > MAX_STRING_LENGTH: msg = f"String exceeds maximum length of {MAX_STRING_LENGTH}" raise ValueError(msg) - extracted_string = string[string.find("```json")+7:string.find("```", string.find("```json")+7)] + extracted_string = string[ + string.find("```json") + 7 : string.find("```", string.find("```json") + 7) + ] return json.loads(extracted_string) - # Restricted set of allowed functions ALLOWED_FUNCTIONS = { # Basic Python builtins From 9cffe68442fe466c19c8359d88cf1da4207e8511 Mon Sep 17 00:00:00 2001 From: vedantsahai18 Date: Thu, 2 Jan 2025 02:16:13 -0500 Subject: [PATCH 4/4] chore: check for JSON code markers --- agents-api/agents_api/activities/utils.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/agents-api/agents_api/activities/utils.py b/agents-api/agents_api/activities/utils.py index 8123490fe..a9d4a11f2 100644 --- a/agents-api/agents_api/activities/utils.py +++ b/agents-api/agents_api/activities/utils.py @@ -107,9 +107,14 @@ def safe_extract_json(string: str): if len(string) > MAX_STRING_LENGTH: msg = f"String exceeds maximum length of {MAX_STRING_LENGTH}" raise ValueError(msg) - extracted_string = string[ - string.find("```json") + 7 : string.find("```", string.find("```json") + 7) - ] + # Check if the string contains JSON code block markers + if "```json" in string: + extracted_string = string[ + string.find("```json") + 7 : string.find("```", string.find("```json") + 7) + ] + else: + # If no markers, try to parse the whole string as JSON + extracted_string = string return json.loads(extracted_string)