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:
# 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
json_string = json_string.strip()
if start_index != -1 and end_index != -1:
extracted_content = json_string[start_index + len("```json"):end_index].strip()
# Parse the JSON string into a Python dictionary
parsed = json.loads(json_string)
# Parse the JSON string into a Python dictionary
parsed = json.loads(extracted_content)
else:
raise Exception("Could not find JSON block in the output.")
return parsed

View File

@ -90,7 +90,7 @@ class MultiDatasetRouterChain(Chain):
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)
router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(
destinations=destinations_str