mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 23:58:34 +00:00
bfe2c44f55
* feat: add docstring * feat: add new docstring * fix: linting --------- Co-authored-by: João Moura <joaomdmoura@gmail.com>
48 lines
1.7 KiB
Python
48 lines
1.7 KiB
Python
import json
|
|
from typing import Any, Type
|
|
|
|
import regex
|
|
from pydantic import BaseModel, ValidationError
|
|
|
|
from crewai.agents.parser import OutputParserException
|
|
|
|
"""Parser for converting text outputs into Pydantic models."""
|
|
|
|
class CrewPydanticOutputParser:
|
|
"""Parses text outputs into specified 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
|