use-workflow-template.ts 1.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  1. import { generateNewNode } from '../utils'
  2. import {
  3. NODE_WIDTH_X_OFFSET,
  4. START_INITIAL_POSITION,
  5. } from '../constants'
  6. import { useIsChatMode } from './use-workflow'
  7. import { useNodesInitialData } from './use-nodes-data'
  8. export const useWorkflowTemplate = () => {
  9. const isChatMode = useIsChatMode()
  10. const nodesInitialData = useNodesInitialData()
  11. const startNode = generateNewNode({
  12. data: nodesInitialData.start,
  13. position: START_INITIAL_POSITION,
  14. })
  15. if (isChatMode) {
  16. const llmNode = generateNewNode({
  17. id: 'llm',
  18. data: {
  19. ...nodesInitialData.llm,
  20. memory: {
  21. window: { enabled: false, size: 10 },
  22. },
  23. selected: true,
  24. },
  25. position: {
  26. x: START_INITIAL_POSITION.x + NODE_WIDTH_X_OFFSET,
  27. y: START_INITIAL_POSITION.y,
  28. },
  29. } as any)
  30. const answerNode = generateNewNode({
  31. id: 'answer',
  32. data: {
  33. ...nodesInitialData.answer,
  34. answer: `{{#${llmNode.id}.text#}}`,
  35. },
  36. position: {
  37. x: START_INITIAL_POSITION.x + NODE_WIDTH_X_OFFSET * 2,
  38. y: START_INITIAL_POSITION.y,
  39. },
  40. } as any)
  41. const startToLlmEdge = {
  42. id: `${startNode.id}-${llmNode.id}`,
  43. source: startNode.id,
  44. sourceHandle: 'source',
  45. target: llmNode.id,
  46. targetHandle: 'target',
  47. }
  48. const llmToAnswerEdge = {
  49. id: `${llmNode.id}-${answerNode.id}`,
  50. source: llmNode.id,
  51. sourceHandle: 'source',
  52. target: answerNode.id,
  53. targetHandle: 'target',
  54. }
  55. return {
  56. nodes: [startNode, llmNode, answerNode],
  57. edges: [startToLlmEdge, llmToAnswerEdge],
  58. }
  59. }
  60. else {
  61. return {
  62. nodes: [startNode],
  63. edges: [],
  64. }
  65. }
  66. }