mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 15:48:29 +00:00
0b0f2d30ab
* sort imports * update --------- Co-authored-by: Brandon Hancock (bhancock_ai) <109994880+bhancockio@users.noreply.github.com> Co-authored-by: Eduardo Chiarotti <dudumelgaco@hotmail.com>
47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
import json
|
|
from typing import Any, Type
|
|
|
|
import regex
|
|
from pydantic import BaseModel, ValidationError
|
|
|
|
from crewai.agents.parser import OutputParserException
|
|
|
|
|
|
class CrewPydanticOutputParser:
|
|
"""Parses the text into pydantic models"""
|
|
|
|
pydantic_object: Type[BaseModel]
|
|
|
|
def parse_result(self, result: str) -> Any:
|
|
result = self._transform_in_valid_json(result)
|
|
|
|
# Treating edge case of function calling llm returning the name instead of tool_name
|
|
json_object = json.loads(result)
|
|
if "tool_name" not in json_object:
|
|
json_object["tool_name"] = json_object.get("name", "")
|
|
result = json.dumps(json_object)
|
|
|
|
try:
|
|
return self.pydantic_object.model_validate(json_object)
|
|
except ValidationError as e:
|
|
name = self.pydantic_object.__name__
|
|
msg = f"Failed to parse {name} from completion {json_object}. Got: {e}"
|
|
raise OutputParserException(error=msg)
|
|
|
|
def _transform_in_valid_json(self, text) -> str:
|
|
text = text.replace("```", "").replace("json", "")
|
|
json_pattern = r"\{(?:[^{}]|(?R))*\}"
|
|
matches = regex.finditer(json_pattern, text)
|
|
|
|
for match in matches:
|
|
try:
|
|
# Attempt to parse the matched string as JSON
|
|
json_obj = json.loads(match.group())
|
|
# Return the first successfully parsed JSON object
|
|
json_obj = json.dumps(json_obj)
|
|
return str(json_obj)
|
|
except json.JSONDecodeError:
|
|
# If parsing fails, skip to the next match
|
|
continue
|
|
return text
|