621b0099d6
- Updated .env.example to include new configuration options for the UltiAI branding and API endpoints. - Enhanced Nginx configuration to support new API routes for the MCP and WebSocket connections. - Introduced sub-filters for branding adjustments in Nginx responses. - Added new JavaScript patch for API endpoint adjustments. - Implemented tests for new API functionalities and improved error handling in the AI gateway.
113 lines
3.1 KiB
Go
113 lines
3.1 KiB
Go
package ai
|
|
|
|
import (
|
|
"encoding/json"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestRepairChatCompletionBody_preservesToolCalls(t *testing.T) {
|
|
body := []byte(`{
|
|
"model":"gpt-4o",
|
|
"messages":[
|
|
{"role":"user","content":"list mail"},
|
|
{"role":"assistant","content":"","tool_calls":[{"id":"call_1","type":"function","function":{"name":"mail_list","arguments":"{\"limit\":3}"}}]},
|
|
{"role":"tool","tool_call_id":"call_1","content":"{\"items\":[]}"}
|
|
]
|
|
}`)
|
|
|
|
repaired, err := repairChatCompletionBody(body)
|
|
if err != nil {
|
|
t.Fatalf("repairChatCompletionBody() error = %v", err)
|
|
}
|
|
|
|
var parsed map[string]any
|
|
if err := json.Unmarshal(repaired, &parsed); err != nil {
|
|
t.Fatalf("json.Unmarshal() error = %v", err)
|
|
}
|
|
messages := parsed["messages"].([]any)
|
|
assistant := messages[1].(map[string]any)
|
|
toolCalls, ok := assistant["tool_calls"].([]any)
|
|
if !ok || len(toolCalls) != 1 {
|
|
t.Fatalf("expected assistant tool_calls preserved, got %#v", assistant)
|
|
}
|
|
}
|
|
|
|
func TestRepairChatCompletionBody_insertsMissingAssistantToolCalls(t *testing.T) {
|
|
body := []byte(`{
|
|
"model":"gpt-4o",
|
|
"messages":[
|
|
{"role":"system","content":"prompt"},
|
|
{"role":"user","content":"list mail"},
|
|
{"role":"tool","tool_call_id":"call_abc","content":"{\"items\":[]}"}
|
|
]
|
|
}`)
|
|
|
|
repaired, err := repairChatCompletionBody(body)
|
|
if err != nil {
|
|
t.Fatalf("repairChatCompletionBody() error = %v", err)
|
|
}
|
|
|
|
var parsed map[string]any
|
|
if err := json.Unmarshal(repaired, &parsed); err != nil {
|
|
t.Fatalf("json.Unmarshal() error = %v", err)
|
|
}
|
|
messages := parsed["messages"].([]any)
|
|
if len(messages) != 4 {
|
|
t.Fatalf("messages len = %d, want 4", len(messages))
|
|
}
|
|
assistant := messages[2].(map[string]any)
|
|
if messageRole(assistant) != "assistant" {
|
|
t.Fatalf("messages[2].role = %q, want assistant", messageRole(assistant))
|
|
}
|
|
toolCalls, ok := assistant["tool_calls"].([]any)
|
|
if !ok || len(toolCalls) != 1 {
|
|
t.Fatalf("expected inserted tool_calls, got %#v", assistant)
|
|
}
|
|
call := toolCalls[0].(map[string]any)
|
|
if call["id"] != "call_abc" {
|
|
t.Fatalf("tool_call id = %v, want call_abc", call["id"])
|
|
}
|
|
}
|
|
|
|
func TestRepairChatCompletionBody_stripsNothingWhenAlreadyValid(t *testing.T) {
|
|
body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
|
|
repaired, err := repairChatCompletionBody(body)
|
|
if err != nil {
|
|
t.Fatalf("repairChatCompletionBody() error = %v", err)
|
|
}
|
|
if string(repaired) != string(body) {
|
|
t.Fatalf("expected unchanged body, got %s", repaired)
|
|
}
|
|
}
|
|
|
|
func TestRepairToolMessages_doesNotDuplicateAssistant(t *testing.T) {
|
|
messages := []any{
|
|
map[string]any{"role": "user", "content": "list"},
|
|
map[string]any{
|
|
"role": "assistant",
|
|
"content": "",
|
|
"tool_calls": []any{
|
|
map[string]any{"id": "call_1", "type": "function"},
|
|
},
|
|
},
|
|
map[string]any{"role": "tool", "tool_call_id": "call_1", "content": "{}"},
|
|
}
|
|
|
|
repaired := repairToolMessages(messages)
|
|
if len(repaired) != 3 {
|
|
t.Fatalf("len = %d, want 3", len(repaired))
|
|
}
|
|
if !strings.Contains(string(mustJSON(repaired)), "call_1") {
|
|
t.Fatal("expected original tool_call_id preserved")
|
|
}
|
|
}
|
|
|
|
func mustJSON(v any) []byte {
|
|
b, err := json.Marshal(v)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return b
|
|
}
|