gemini_chat_completions_compat_service_test.go 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163
  1. package service
  2. import (
  3. "encoding/json"
  4. "testing"
  5. "github.com/stretchr/testify/require"
  6. )
  7. func TestGeminiResponseToChatCompletionsPreservesInlineData(t *testing.T) {
  8. tests := []struct {
  9. name string
  10. parts []any
  11. want string
  12. }{
  13. {
  14. name: "image only",
  15. parts: []any{
  16. map[string]any{"inlineData": map[string]any{"mimeType": "image/png", "data": "aW1hZ2U="}},
  17. },
  18. want: "![image](data:image/png;base64,aW1hZ2U=)",
  19. },
  20. {
  21. name: "text and image",
  22. parts: []any{
  23. map[string]any{"text": "rendered image:\n"},
  24. map[string]any{"inlineData": map[string]any{"mimeType": "image/webp", "data": "d2VicA=="}},
  25. },
  26. want: "rendered image:\n![image](data:image/webp;base64,d2VicA==)",
  27. },
  28. }
  29. for _, tt := range tests {
  30. t.Run(tt.name, func(t *testing.T) {
  31. geminiResp := map[string]any{
  32. "candidates": []any{map[string]any{
  33. "content": map[string]any{"parts": tt.parts},
  34. "finishReason": "STOP",
  35. }},
  36. }
  37. rawData, err := json.Marshal(geminiResp)
  38. require.NoError(t, err)
  39. got, _, err := geminiResponseToChatCompletions(geminiResp, "gemini-test", rawData, nil)
  40. require.NoError(t, err)
  41. require.Len(t, got.Choices, 1)
  42. var content string
  43. require.NoError(t, json.Unmarshal(got.Choices[0].Message.Content, &content))
  44. require.Equal(t, tt.want, content)
  45. require.Equal(t, "stop", got.Choices[0].FinishReason)
  46. })
  47. }
  48. }
  49. func TestGeminiResponseToChatCompletionsOmitsInvalidInlineData(t *testing.T) {
  50. tests := []struct {
  51. name string
  52. inlineData map[string]any
  53. }{
  54. {
  55. name: "unsupported MIME type",
  56. inlineData: map[string]any{"mimeType": "image/svg+xml", "data": "PHN2Zz48L3N2Zz4="},
  57. },
  58. {
  59. name: "malformed MIME type",
  60. inlineData: map[string]any{"mimeType": "image/png; charset=utf-8", "data": "aW1hZ2U="},
  61. },
  62. {
  63. name: "malformed base64",
  64. inlineData: map[string]any{"mimeType": "image/png", "data": "not-valid-base64!!!"},
  65. },
  66. }
  67. for _, tt := range tests {
  68. t.Run(tt.name, func(t *testing.T) {
  69. geminiResp := map[string]any{
  70. "candidates": []any{map[string]any{
  71. "content": map[string]any{"parts": []any{map[string]any{"text": "before"}, map[string]any{"inlineData": tt.inlineData}, map[string]any{"text": "after"}}},
  72. "finishReason": "STOP",
  73. }},
  74. }
  75. rawData, err := json.Marshal(geminiResp)
  76. require.NoError(t, err)
  77. got, _, err := geminiResponseToChatCompletions(geminiResp, "gemini-test", rawData, nil)
  78. require.NoError(t, err)
  79. var content string
  80. require.NoError(t, json.Unmarshal(got.Choices[0].Message.Content, &content))
  81. require.Equal(t, "beforeafter", content)
  82. })
  83. }
  84. }
  85. func TestConvertGeminiToClaudeMessageOmitsInlineDataForAnthropicMessages(t *testing.T) {
  86. geminiResp := map[string]any{
  87. "candidates": []any{map[string]any{
  88. "content": map[string]any{"parts": []any{
  89. map[string]any{"text": "before"},
  90. map[string]any{"inlineData": map[string]any{"mimeType": "image/png", "data": "aW1hZ2U="}},
  91. map[string]any{"functionCall": map[string]any{"name": "get_weather", "args": map[string]any{"city": "Paris"}}},
  92. map[string]any{"text": "after"},
  93. }},
  94. "finishReason": "STOP",
  95. }},
  96. }
  97. rawData, err := json.Marshal(geminiResp)
  98. require.NoError(t, err)
  99. withInlineData, _ := convertGeminiToClaudeMessage(geminiResp, "gemini-test", rawData, true)
  100. contentWithInlineData, ok := withInlineData["content"].([]any)
  101. require.True(t, ok)
  102. require.Len(t, contentWithInlineData, 4)
  103. require.Equal(t, map[string]any{"type": "text", "text": "before"}, contentWithInlineData[0])
  104. require.Equal(t, map[string]any{"type": "text", "text": "![image](data:image/png;base64,aW1hZ2U=)"}, contentWithInlineData[1])
  105. toolUse, ok := contentWithInlineData[2].(map[string]any)
  106. require.True(t, ok)
  107. require.Equal(t, "tool_use", toolUse["type"])
  108. require.Equal(t, "get_weather", toolUse["name"])
  109. require.Equal(t, map[string]any{"type": "text", "text": "after"}, contentWithInlineData[3])
  110. withoutInlineData, _ := convertGeminiToClaudeMessage(geminiResp, "gemini-test", rawData, false)
  111. contentWithoutInlineData, ok := withoutInlineData["content"].([]any)
  112. require.True(t, ok)
  113. require.Len(t, contentWithoutInlineData, 3)
  114. require.Equal(t, map[string]any{"type": "text", "text": "before"}, contentWithoutInlineData[0])
  115. toolUseWithoutInlineData, ok := contentWithoutInlineData[1].(map[string]any)
  116. require.True(t, ok)
  117. require.Equal(t, "tool_use", toolUseWithoutInlineData["type"])
  118. require.Equal(t, "get_weather", toolUseWithoutInlineData["name"])
  119. require.Equal(t, map[string]any{"type": "text", "text": "after"}, contentWithoutInlineData[2])
  120. }
  121. func TestGeminiResponseToChatCompletionsRetainsTextAndToolBehavior(t *testing.T) {
  122. geminiResp := map[string]any{
  123. "candidates": []any{map[string]any{
  124. "content": map[string]any{"parts": []any{
  125. map[string]any{"text": "checking"},
  126. map[string]any{"functionCall": map[string]any{
  127. "name": "get_weather",
  128. "args": map[string]any{"city": "Paris"},
  129. }},
  130. }},
  131. "finishReason": "STOP",
  132. }},
  133. }
  134. rawData, err := json.Marshal(geminiResp)
  135. require.NoError(t, err)
  136. got, _, err := geminiResponseToChatCompletions(geminiResp, "gemini-test", rawData, nil)
  137. require.NoError(t, err)
  138. require.Len(t, got.Choices, 1)
  139. choice := got.Choices[0]
  140. var content string
  141. require.NoError(t, json.Unmarshal(choice.Message.Content, &content))
  142. require.Equal(t, "checking", content)
  143. require.Equal(t, "tool_calls", choice.FinishReason)
  144. require.Len(t, choice.Message.ToolCalls, 1)
  145. require.Equal(t, "get_weather", choice.Message.ToolCalls[0].Function.Name)
  146. require.JSONEq(t, `{"city":"Paris"}`, choice.Message.ToolCalls[0].Function.Arguments)
  147. }