mock.go 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328
  1. package tester
  2. import (
  3. "time"
  4. "github.com/langgenius/dify-plugin-daemon/internal/core/dify_invocation"
  5. "github.com/langgenius/dify-plugin-daemon/internal/types/entities/model_entities"
  6. "github.com/langgenius/dify-plugin-daemon/internal/types/entities/tool_entities"
  7. "github.com/langgenius/dify-plugin-daemon/internal/utils/routine"
  8. "github.com/langgenius/dify-plugin-daemon/internal/utils/stream"
  9. )
  10. type MockedDifyInvocation struct{}
  11. func NewMockedDifyInvocation() dify_invocation.BackwardsInvocation {
  12. return &MockedDifyInvocation{}
  13. }
  14. func (m *MockedDifyInvocation) InvokeLLM(payload *dify_invocation.InvokeLLMRequest) (*stream.Stream[model_entities.LLMResultChunk], error) {
  15. stream := stream.NewStream[model_entities.LLMResultChunk](5)
  16. routine.Submit(nil, func() {
  17. stream.Write(model_entities.LLMResultChunk{
  18. Model: model_entities.LLMModel(payload.Model),
  19. PromptMessages: payload.PromptMessages,
  20. SystemFingerprint: "test",
  21. Delta: model_entities.LLMResultChunkDelta{
  22. Index: &[]int{1}[0],
  23. Message: model_entities.PromptMessage{
  24. Role: model_entities.PROMPT_MESSAGE_ROLE_ASSISTANT,
  25. Content: "hello",
  26. Name: "test",
  27. },
  28. },
  29. })
  30. time.Sleep(100 * time.Millisecond)
  31. stream.Write(model_entities.LLMResultChunk{
  32. Model: model_entities.LLMModel(payload.Model),
  33. PromptMessages: payload.PromptMessages,
  34. SystemFingerprint: "test",
  35. Delta: model_entities.LLMResultChunkDelta{
  36. Index: &[]int{1}[0],
  37. Message: model_entities.PromptMessage{
  38. Role: model_entities.PROMPT_MESSAGE_ROLE_ASSISTANT,
  39. Content: " world",
  40. Name: "test",
  41. },
  42. },
  43. })
  44. time.Sleep(100 * time.Millisecond)
  45. stream.Write(model_entities.LLMResultChunk{
  46. Model: model_entities.LLMModel(payload.Model),
  47. PromptMessages: payload.PromptMessages,
  48. SystemFingerprint: "test",
  49. Delta: model_entities.LLMResultChunkDelta{
  50. Index: &[]int{2}[0],
  51. Message: model_entities.PromptMessage{
  52. Role: model_entities.PROMPT_MESSAGE_ROLE_ASSISTANT,
  53. Content: " world",
  54. Name: "test",
  55. },
  56. },
  57. })
  58. time.Sleep(100 * time.Millisecond)
  59. stream.Write(model_entities.LLMResultChunk{
  60. Model: model_entities.LLMModel(payload.Model),
  61. PromptMessages: payload.PromptMessages,
  62. SystemFingerprint: "test",
  63. Delta: model_entities.LLMResultChunkDelta{
  64. Index: &[]int{3}[0],
  65. Message: model_entities.PromptMessage{
  66. Role: model_entities.PROMPT_MESSAGE_ROLE_ASSISTANT,
  67. Content: " !",
  68. Name: "test",
  69. },
  70. },
  71. })
  72. time.Sleep(100 * time.Millisecond)
  73. stream.Write(model_entities.LLMResultChunk{
  74. Model: model_entities.LLMModel(payload.Model),
  75. PromptMessages: payload.PromptMessages,
  76. SystemFingerprint: "test",
  77. Delta: model_entities.LLMResultChunkDelta{
  78. Index: &[]int{3}[0],
  79. Usage: &model_entities.LLMUsage{
  80. PromptTokens: &[]int{100}[0],
  81. CompletionTokens: &[]int{100}[0],
  82. TotalTokens: &[]int{200}[0],
  83. Latency: &[]float64{0.4}[0],
  84. Currency: &[]string{"USD"}[0],
  85. },
  86. },
  87. })
  88. stream.Close()
  89. })
  90. return stream, nil
  91. }
  92. func (m *MockedDifyInvocation) InvokeTextEmbedding(payload *dify_invocation.InvokeTextEmbeddingRequest) (*model_entities.TextEmbeddingResult, error) {
  93. result := model_entities.TextEmbeddingResult{
  94. Model: payload.Model,
  95. Usage: model_entities.EmbeddingUsage{
  96. Tokens: &[]int{100}[0],
  97. TotalTokens: &[]int{100}[0],
  98. Latency: &[]float64{0.1}[0],
  99. Currency: &[]string{"USD"}[0],
  100. },
  101. }
  102. for range payload.Texts {
  103. result.Embeddings = append(result.Embeddings, []float64{0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0})
  104. }
  105. return &result, nil
  106. }
  107. func (m *MockedDifyInvocation) InvokeRerank(payload *dify_invocation.InvokeRerankRequest) (*model_entities.RerankResult, error) {
  108. result := model_entities.RerankResult{
  109. Model: payload.Model,
  110. }
  111. for i, doc := range payload.Docs {
  112. result.Docs = append(result.Docs, model_entities.RerankDocument{
  113. Index: &[]int{i}[0],
  114. Score: &[]float64{0.1}[0],
  115. Text: &doc,
  116. })
  117. }
  118. return &result, nil
  119. }
  120. func (m *MockedDifyInvocation) InvokeTTS(payload *dify_invocation.InvokeTTSRequest) (*stream.Stream[model_entities.TTSResult], error) {
  121. stream := stream.NewStream[model_entities.TTSResult](5)
  122. routine.Submit(nil, func() {
  123. for i := 0; i < 10; i++ {
  124. stream.Write(model_entities.TTSResult{
  125. Result: "a1b2c3d4",
  126. })
  127. time.Sleep(100 * time.Millisecond)
  128. }
  129. stream.Close()
  130. })
  131. return stream, nil
  132. }
  133. func (m *MockedDifyInvocation) InvokeSpeech2Text(payload *dify_invocation.InvokeSpeech2TextRequest) (*model_entities.Speech2TextResult, error) {
  134. result := model_entities.Speech2TextResult{
  135. Result: "hello world",
  136. }
  137. return &result, nil
  138. }
  139. func (m *MockedDifyInvocation) InvokeModeration(payload *dify_invocation.InvokeModerationRequest) (*model_entities.ModerationResult, error) {
  140. result := model_entities.ModerationResult{
  141. Result: true,
  142. }
  143. return &result, nil
  144. }
  145. func (m *MockedDifyInvocation) InvokeTool(payload *dify_invocation.InvokeToolRequest) (*stream.Stream[tool_entities.ToolResponseChunk], error) {
  146. stream := stream.NewStream[tool_entities.ToolResponseChunk](5)
  147. routine.Submit(nil, func() {
  148. for i := 0; i < 10; i++ {
  149. stream.Write(tool_entities.ToolResponseChunk{
  150. Type: tool_entities.ToolResponseChunkTypeText,
  151. Message: map[string]any{
  152. "text": "hello world",
  153. },
  154. })
  155. time.Sleep(100 * time.Millisecond)
  156. }
  157. stream.Close()
  158. })
  159. return stream, nil
  160. }
  161. func (m *MockedDifyInvocation) InvokeApp(payload *dify_invocation.InvokeAppRequest) (*stream.Stream[map[string]any], error) {
  162. stream := stream.NewStream[map[string]any](5)
  163. routine.Submit(nil, func() {
  164. stream.Write(map[string]any{
  165. "event": "agent_message",
  166. "message_id": "5ad4cb98-f0c7-4085-b384-88c403be6290",
  167. "conversation_id": "45701982-8118-4bc5-8e9b-64562b4555f2",
  168. "answer": "なんで",
  169. "created_at": time.Now().Unix(),
  170. })
  171. time.Sleep(100 * time.Millisecond)
  172. stream.Write(map[string]any{
  173. "event": "agent_message",
  174. "message_id": "5ad4cb98-f0c7-4085-b384-88c403be6290",
  175. "conversation_id": "45701982-8118-4bc5-8e9b-64562b4555f2",
  176. "answer": "春日影",
  177. "created_at": time.Now().Unix(),
  178. })
  179. time.Sleep(100 * time.Millisecond)
  180. stream.Write(map[string]any{
  181. "event": "agent_message",
  182. "message_id": "5ad4cb98-f0c7-4085-b384-88c403be6290",
  183. "conversation_id": "45701982-8118-4bc5-8e9b-64562b4555f2",
  184. "answer": "やったの",
  185. "created_at": time.Now().Unix(),
  186. })
  187. time.Sleep(100 * time.Millisecond)
  188. stream.Write(map[string]any{
  189. "event": "message_end",
  190. "id": "5e52ce04-874b-4d27-9045-b3bc80def685",
  191. "conversation_id": "45701982-8118-4bc5-8e9b-64562b4555f2",
  192. "created_at": time.Now().Unix(),
  193. "metadata": map[string]any{
  194. "retriever_resources": []map[string]any{
  195. {
  196. "position": 1,
  197. "dataset_id": "101b4c97-fc2e-463c-90b1-5261a4cdcafb",
  198. "dataset_name": "あなた",
  199. "document_id": "8dd1ad74-0b5f-4175-b735-7d98bbbb4e00",
  200. "document_name": "ご自分のことばかりですのね",
  201. "score": 0.98457545,
  202. "content": "CRYCHICは壊れてしまいましたわ",
  203. },
  204. },
  205. "usage": map[string]any{
  206. "prompt_tokens": 1033,
  207. "prompt_unit_price": "0.001",
  208. "prompt_price_unit": "0.001",
  209. "prompt_price": "0.0010330",
  210. "completion_tokens": 135,
  211. "completion_unit_price": "0.002",
  212. "completion_price_unit": "0.001",
  213. "completion_price": "0.0002700",
  214. "total_tokens": 1168,
  215. "total_price": "0.0013030",
  216. "currency": "USD",
  217. "latency": 1.381760165997548,
  218. },
  219. },
  220. })
  221. time.Sleep(100 * time.Millisecond)
  222. stream.Write(map[string]any{
  223. "event": "message_file",
  224. "id": "5e52ce04-874b-4d27-9045-b3bc80def685",
  225. "conversation_id": "45701982-8118-4bc5-8e9b-64562b4555f2",
  226. "belongs_to": "assistant",
  227. "url": "https://www.youtube.com/watch?v=dQw4w9WgXcQ",
  228. "created_at": time.Now().Unix(),
  229. })
  230. time.Sleep(100 * time.Millisecond)
  231. stream.Close()
  232. })
  233. return stream, nil
  234. }
  235. func (m *MockedDifyInvocation) InvokeEncrypt(payload *dify_invocation.InvokeEncryptRequest) (map[string]any, error) {
  236. return payload.Data, nil
  237. }
  238. func (m *MockedDifyInvocation) InvokeParameterExtractor(payload *dify_invocation.InvokeParameterExtractorRequest) (*dify_invocation.InvokeNodeResponse, error) {
  239. resp := &dify_invocation.InvokeNodeResponse{
  240. ProcessData: map[string]any{},
  241. Outputs: map[string]any{},
  242. Inputs: map[string]any{
  243. "query": payload.Query,
  244. },
  245. }
  246. for _, parameter := range payload.Parameters {
  247. typ := parameter.Type
  248. if typ == "string" {
  249. resp.Outputs[parameter.Name] = "Never gonna give you up ~"
  250. } else if typ == "number" {
  251. resp.Outputs[parameter.Name] = 1234567890
  252. } else if typ == "bool" {
  253. resp.Outputs[parameter.Name] = true
  254. } else if typ == "select" {
  255. options := parameter.Options
  256. if len(options) == 0 {
  257. resp.Outputs[parameter.Name] = "Never gonna let you down ~"
  258. } else {
  259. resp.Outputs[parameter.Name] = options[0]
  260. }
  261. } else if typ == "array[string]" {
  262. resp.Outputs[parameter.Name] = []string{
  263. "Never gonna run around and desert you ~",
  264. "Never gonna make you cry ~",
  265. "Never gonna say goodbye ~",
  266. "Never gonna tell a lie and hurt you ~",
  267. }
  268. } else if typ == "array[number]" {
  269. resp.Outputs[parameter.Name] = []int{114, 514, 1919, 810}
  270. } else if typ == "array[bool]" {
  271. resp.Outputs[parameter.Name] = []bool{true, false, true, false, true, false, true, false, true, false}
  272. } else if typ == "array[object]" {
  273. resp.Outputs[parameter.Name] = []map[string]any{
  274. {
  275. "name": "お願い",
  276. "age": 55555,
  277. },
  278. {
  279. "name": "何でもするがら",
  280. "age": 99999,
  281. },
  282. }
  283. }
  284. }
  285. return resp, nil
  286. }
  287. func (m *MockedDifyInvocation) InvokeQuestionClassifier(payload *dify_invocation.InvokeQuestionClassifierRequest) (*dify_invocation.InvokeNodeResponse, error) {
  288. return &dify_invocation.InvokeNodeResponse{
  289. ProcessData: map[string]any{},
  290. Outputs: map[string]any{
  291. "class_name": payload.Classes[0].Name,
  292. },
  293. Inputs: map[string]any{},
  294. }, nil
  295. }
  296. func (m *MockedDifyInvocation) InvokeSummary(payload *dify_invocation.InvokeSummaryRequest) (*dify_invocation.InvokeSummaryResponse, error) {
  297. return &dify_invocation.InvokeSummaryResponse{
  298. Summary: payload.Text,
  299. }, nil
  300. }
  301. func (m *MockedDifyInvocation) UploadFile(payload *dify_invocation.UploadFileRequest) (*dify_invocation.UploadFileResponse, error) {
  302. return &dify_invocation.UploadFileResponse{
  303. URL: "https://www.youtube.com/watch?v=dQw4w9WgXcQ",
  304. }, nil
  305. }