llmgw/internal/router/responses_test.go

package router

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
)

func TestAResponsesRequestBecomesAChatOne(t *testing.T) {
	// What Codex sends in the middle of a task: its instructions, the
	// person's turn, a call it made and what the call returned, and its tools
	// — one a function, one free-form.
	var in respRequest
	if err := json.Unmarshal([]byte(`{
	  "model": "coder", "instructions": "You are Codex.", "stream": true,
	  "input": [
	    {"type":"message","role":"developer","content":[{"type":"input_text","text":"Be careful."}]},
	    {"type":"message","role":"user","content":[{"type":"input_text","text":"List the files."}]},
	    {"type":"reasoning","summary":[]},
	    {"type":"function_call","call_id":"c1","name":"shell","arguments":"{\"command\":[\"ls\"]}"},
	    {"type":"function_call_output","call_id":"c1","output":"a.go\nb.go"},
	    {"type":"custom_tool_call","call_id":"c2","name":"apply_patch","input":"*** Begin Patch"},
	    {"type":"custom_tool_call_output","call_id":"c2","output":"done"}
	  ],
	  "tools": [
	    {"type":"function","name":"shell","description":"Run a command","parameters":{"type":"object","properties":{"command":{"type":"array"}}}},
	    {"type":"custom","name":"apply_patch","description":"Apply a patch"},
	    {"type":"web_search"}
	  ],
	  "tool_choice":"auto", "store": false, "reasoning": {"effort":"medium"}
	}`), &in); err != nil {
		t.Fatal(err)
	}
	chat, custom, err := toChat(in)
	if err != nil {
		t.Fatal(err)
	}
	b, _ := json.Marshal(chat)
	var got struct {
		Messages []struct {
			Role       string `json:"role"`
			Content    string `json:"content"`
			ToolCallID string `json:"tool_call_id"`
			ToolCalls  []struct {
				ID       string `json:"id"`
				Function struct{ Name, Arguments string }
			} `json:"tool_calls"`
		} `json:"messages"`
		Tools []struct {
			Function struct{ Name string } `json:"function"`
		} `json:"tools"`
		StreamOptions struct {
			IncludeUsage bool `json:"include_usage"`
		} `json:"stream_options"`
	}
	if err := json.Unmarshal(b, &got); err != nil {
		t.Fatal(err)
	}
	roles := []string{}
	for _, m := range got.Messages {
		roles = append(roles, m.Role)
	}
	if strings.Join(roles, ",") != "system,system,user,assistant,tool,assistant,tool" {
		t.Errorf("roles = %v", roles)
	}
	if got.Messages[3].ToolCalls[0].Function.Name != "shell" || got.Messages[4].ToolCallID != "c1" || got.Messages[4].Content != "a.go\nb.go" {
		t.Errorf("the shell call and its output: %+v %+v", got.Messages[3], got.Messages[4])
	}
	// A custom tool is a function of one string, and its call carries it.
	if !strings.Contains(got.Messages[5].ToolCalls[0].Function.Arguments, `"input":"*** Begin Patch"`) || !custom["apply_patch"] {
		t.Errorf("custom call = %+v", got.Messages[5])
	}
	// web_search is the server's own tool; nothing here can run it.
	if len(got.Tools) != 2 || !got.StreamOptions.IncludeUsage {
		t.Errorf("tools %+v, usage asked %v", got.Tools, got.StreamOptions.IncludeUsage)
	}
}

// streamingEngine answers chat completions with a stream: some text, a tool
// call in pieces, and usage.
func streamingEngine(t *testing.T, model string) *httptest.Server {
	t.Helper()
	mux := http.NewServeMux()
	mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) {
		fmt.Fprintf(w, "# TYPE vllm:num_requests_running gauge\nvllm:num_requests_running 0\n")
	})
	mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		var req struct {
			Stream bool `json:"stream"`
		}
		_ = json.NewDecoder(r.Body).Decode(&req)
		if !req.Stream {
			w.Header().Set("Content-Type", "application/json")
			fmt.Fprintf(w, `{"model":%q,"choices":[{"message":{"content":"Hello.","tool_calls":[{"id":"call_9","type":"function","function":{"name":"shell","arguments":"{\"command\":[\"ls\"]}"}}]}}],"usage":{"prompt_tokens":11,"completion_tokens":7}}`, model)
			return
		}
		w.Header().Set("Content-Type", "text/event-stream")
		for _, c := range []string{
			`{"choices":[{"delta":{"content":"Look"}}]}`,
			`{"choices":[{"delta":{"content":"ing."}}]}`,
			`{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"shell","arguments":"{\"comm"}}]}}]}`,
			`{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"and\":[\"ls\"]}"}}]}}]}`,
			`{"choices":[{"delta":{},"finish_reason":"tool_calls"}]}`,
			`{"choices":[],"usage":{"prompt_tokens":20,"completion_tokens":9}}`,
		} {
			fmt.Fprintf(w, "data: %s\n\n", c)
		}
		fmt.Fprint(w, "data: [DONE]\n\n")
	})
	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)
	return srv
}

func TestAResponsesStreamIsTranslatedFromChat(t *testing.T) {
	srv := streamingEngine(t, "coder")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: srv.URL, Models: []string{"coder"}}))
	waitForEndpoint(t, rt)

	req := httptest.NewRequest("POST", "/v1/responses", strings.NewReader(
		`{"model":"coder","stream":true,"input":"list the files","tools":[{"type":"function","name":"shell","parameters":{"type":"object"}}]}`))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	rt.Handler().ServeHTTP(w, req)
	if w.Code != 200 {
		t.Fatalf("%d %s", w.Code, w.Body.String())
	}
	var types []string
	var completed map[string]any
	var args string
	for _, block := range strings.Split(w.Body.String(), "\n\n") {
		for _, line := range strings.Split(block, "\n") {
			if !strings.HasPrefix(line, "data: ") {
				continue
			}
			var ev map[string]any
			if err := json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &ev); err != nil {
				t.Fatalf("not JSON: %s", line)
			}
			types = append(types, ev["type"].(string))
			if ev["type"] == "response.completed" {
				completed = ev["response"].(map[string]any)
			}
			if ev["type"] == "response.function_call_arguments.done" {
				args = ev["arguments"].(string)
			}
		}
	}
	want := []string{"response.created", "response.in_progress", "response.output_item.added", "response.content_part.added",
		"response.output_text.delta", "response.output_text.delta", "response.output_item.added",
		"response.function_call_arguments.delta", "response.function_call_arguments.delta",
		"response.output_text.done", "response.content_part.done", "response.output_item.done",
		"response.function_call_arguments.done", "response.output_item.done", "response.completed"}
	if strings.Join(types, " ") != strings.Join(want, " ") {
		t.Errorf("events:\n got %v\nwant %v", types, want)
	}
	if args != `{"command":["ls"]}` {
		t.Errorf("arguments = %q", args)
	}
	out, _ := completed["output"].([]any)
	if len(out) != 2 {
		t.Fatalf("output = %v", completed["output"])
	}
	msg, call := out[0].(map[string]any), out[1].(map[string]any)
	if msg["type"] != "message" || call["type"] != "function_call" || call["name"] != "shell" || call["call_id"] != "call_1" {
		t.Errorf("output = %v", out)
	}
	if u := completed["usage"].(map[string]any); u["input_tokens"].(float64) != 20 || u["output_tokens"].(float64) != 9 {
		t.Errorf("usage = %v", u)
	}
}

func TestAResponsesRequestWithoutAStreamGetsAResponseObject(t *testing.T) {
	srv := streamingEngine(t, "coder")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: srv.URL, Models: []string{"coder"}}))
	waitForEndpoint(t, rt)
	req := httptest.NewRequest("POST", "/v1/responses", strings.NewReader(`{"model":"coder","input":"hi"}`))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	rt.Handler().ServeHTTP(w, req)
	var got struct {
		Object string `json:"object"`
		Status string `json:"status"`
		Output []struct {
			Type    string `json:"type"`
			Name    string `json:"name"`
			Content []struct {
				Text string `json:"text"`
			} `json:"content"`
		} `json:"output"`
		Usage struct {
			Input  int `json:"input_tokens"`
			Output int `json:"output_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil {
		t.Fatalf("%v: %s", err, w.Body.String())
	}
	if got.Object != "response" || got.Status != "completed" || len(got.Output) != 2 ||
		got.Output[0].Content[0].Text != "Hello." || got.Output[1].Name != "shell" || got.Usage.Input != 11 {
		t.Errorf("%s", w.Body.String())
	}
}