fix: json parse in router chain output (#243)

This commit is contained in:
John Wang 2023-05-29 18:25:01 +08:00 committed by GitHub
parent 091beffae7
commit af2138e8b8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 9 additions and 6 deletions

View File

@ -84,13 +84,16 @@ class RouterOutputParser(BaseOutputParser[Dict[str, str]]):
def parse_json_markdown(self, json_string: str) -> dict: def parse_json_markdown(self, json_string: str) -> dict:
# Remove the triple backticks if present # Remove the triple backticks if present
json_string = json_string.replace("```json", "").replace("```", "") start_index = json_string.find("```json")
end_index = json_string.find("```", start_index + len("```json"))
# Strip whitespace and newlines from the start and end if start_index != -1 and end_index != -1:
json_string = json_string.strip() extracted_content = json_string[start_index + len("```json"):end_index].strip()
# Parse the JSON string into a Python dictionary # Parse the JSON string into a Python dictionary
parsed = json.loads(json_string) parsed = json.loads(extracted_content)
else:
raise Exception("Could not find JSON block in the output.")
return parsed return parsed

View File

@ -90,7 +90,7 @@ class MultiDatasetRouterChain(Chain):
callback_manager=llm_callback_manager callback_manager=llm_callback_manager
) )
destinations = [f"{d.id}: {d.description}" for d in datasets] destinations = ["{}: {}".format(d.id, d.description.replace('\n', ' ')) for d in datasets]
destinations_str = "\n".join(destinations) destinations_str = "\n".join(destinations)
router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format( router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(
destinations=destinations_str destinations=destinations_str