use-workflow-template.ts 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  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 { newNode: startNode } = generateNewNode({
  12. data: nodesInitialData.start,
  13. position: START_INITIAL_POSITION,
  14. })
  15. if (isChatMode) {
  16. const { newNode: llmNode } = generateNewNode({
  17. id: 'llm',
  18. data: {
  19. ...nodesInitialData.llm,
  20. memory: {
  21. window: { enabled: false, size: 10 },
  22. query_prompt_template: '{{#sys.query#}}',
  23. },
  24. selected: true,
  25. },
  26. position: {
  27. x: START_INITIAL_POSITION.x + NODE_WIDTH_X_OFFSET,
  28. y: START_INITIAL_POSITION.y,
  29. },
  30. } as any)
  31. const { newNode: answerNode } = generateNewNode({
  32. id: 'answer',
  33. data: {
  34. ...nodesInitialData.answer,
  35. answer: `{{#${llmNode.id}.text#}}`,
  36. },
  37. position: {
  38. x: START_INITIAL_POSITION.x + NODE_WIDTH_X_OFFSET * 2,
  39. y: START_INITIAL_POSITION.y,
  40. },
  41. } as any)
  42. const startToLlmEdge = {
  43. id: `${startNode.id}-${llmNode.id}`,
  44. source: startNode.id,
  45. sourceHandle: 'source',
  46. target: llmNode.id,
  47. targetHandle: 'target',
  48. }
  49. const llmToAnswerEdge = {
  50. id: `${llmNode.id}-${answerNode.id}`,
  51. source: llmNode.id,
  52. sourceHandle: 'source',
  53. target: answerNode.id,
  54. targetHandle: 'target',
  55. }
  56. return {
  57. nodes: [startNode, llmNode, answerNode],
  58. edges: [startToLlmEdge, llmToAnswerEdge],
  59. }
  60. }
  61. else {
  62. return {
  63. nodes: [startNode],
  64. edges: [],
  65. }
  66. }
  67. }