entities.py 1.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. from enum import Enum
  2. from pydantic import BaseModel, Field
  3. from core.workflow.entities.base_node_data_entities import BaseNodeData
  4. class AnswerNodeData(BaseNodeData):
  5. """
  6. Answer Node Data.
  7. """
  8. answer: str = Field(..., description="answer template string")
  9. class GenerateRouteChunk(BaseModel):
  10. """
  11. Generate Route Chunk.
  12. """
  13. class ChunkType(Enum):
  14. VAR = "var"
  15. TEXT = "text"
  16. type: ChunkType = Field(..., description="generate route chunk type")
  17. class VarGenerateRouteChunk(GenerateRouteChunk):
  18. """
  19. Var Generate Route Chunk.
  20. """
  21. type: GenerateRouteChunk.ChunkType = GenerateRouteChunk.ChunkType.VAR
  22. """generate route chunk type"""
  23. value_selector: list[str] = Field(..., description="value selector")
  24. class TextGenerateRouteChunk(GenerateRouteChunk):
  25. """
  26. Text Generate Route Chunk.
  27. """
  28. type: GenerateRouteChunk.ChunkType = GenerateRouteChunk.ChunkType.TEXT
  29. """generate route chunk type"""
  30. text: str = Field(..., description="text")
  31. class AnswerNodeDoubleLink(BaseModel):
  32. node_id: str = Field(..., description="node id")
  33. source_node_ids: list[str] = Field(..., description="source node ids")
  34. target_node_ids: list[str] = Field(..., description="target node ids")
  35. class AnswerStreamGenerateRoute(BaseModel):
  36. """
  37. AnswerStreamGenerateRoute entity
  38. """
  39. answer_dependencies: dict[str, list[str]] = Field(
  40. ..., description="answer dependencies (answer node id -> dependent answer node ids)"
  41. )
  42. answer_generate_route: dict[str, list[GenerateRouteChunk]] = Field(
  43. ..., description="answer generate route (answer node id -> generate route chunks)"
  44. )