entities.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. from typing import Any, Literal, Union
  2. from pydantic import BaseModel, field_validator
  3. from pydantic_core.core_schema import ValidationInfo
  4. from core.tools.entities.tool_entities import ToolProviderType
  5. from core.workflow.nodes.base.entities import BaseNodeData
  6. class ToolEntity(BaseModel):
  7. provider_id: str
  8. provider_type: ToolProviderType
  9. provider_name: str # redundancy
  10. tool_name: str
  11. tool_label: str # redundancy
  12. tool_configurations: dict[str, Any]
  13. plugin_unique_identifier: str | None = None # redundancy
  14. @field_validator("tool_configurations", mode="before")
  15. @classmethod
  16. def validate_tool_configurations(cls, value, values: ValidationInfo):
  17. if not isinstance(value, dict):
  18. raise ValueError("tool_configurations must be a dictionary")
  19. for key in values.data.get("tool_configurations", {}):
  20. value = values.data.get("tool_configurations", {}).get(key)
  21. if not isinstance(value, str | int | float | bool):
  22. raise ValueError(f"{key} must be a string")
  23. return value
  24. class ToolNodeData(BaseNodeData, ToolEntity):
  25. class ToolInput(BaseModel):
  26. # TODO: check this type
  27. value: Union[Any, list[str]]
  28. type: Literal["mixed", "variable", "constant"]
  29. @field_validator("type", mode="before")
  30. @classmethod
  31. def check_type(cls, value, validation_info: ValidationInfo):
  32. typ = value
  33. value = validation_info.data.get("value")
  34. if typ == "mixed" and not isinstance(value, str):
  35. raise ValueError("value must be a string")
  36. elif typ == "variable":
  37. if not isinstance(value, list):
  38. raise ValueError("value must be a list")
  39. for val in value:
  40. if not isinstance(val, str):
  41. raise ValueError("value must be a list of strings")
  42. elif typ == "constant" and not isinstance(value, str | int | float | bool):
  43. raise ValueError("value must be a string, int, float, or bool")
  44. return typ
  45. tool_parameters: dict[str, ToolInput]