cogito/tools_native_test.go

package cogito

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/sashabaranov/go-openai"
)

// audioTM is a local TypedMultimedia test double carrying an audio payload.
// It is defined in THIS package on purpose: the clients-package double lives in
// a different package, and package cogito cannot import clients (clients imports
// cogito — an internal-test import cycle).
type audioTM struct{ data, format string }

func (a audioTM) URL() string          { return "" }
func (a audioTM) MediaKind() MediaKind { return MediaAudio }
func (a audioTM) Data() string         { return a.data }
func (a audioTM) Format() string       { return a.format }

// nativeStreamLLM is an in-package StreamingLLM + NativePartsAware double that
// mirrors what the real LocalAI client does: it stashes the pending native
// parts via SetPendingNativeParts and, on CreateChatCompletionStream, serializes
// any stashed audio into the request body it POSTs to the SSE server — so the
// audio genuinely reaches the wire only if the seam forwarded the parts to it.
type nativeStreamLLM struct {
	url     string
	pending []NativePart
}

func (l *nativeStreamLLM) Ask(ctx context.Context, f Fragment) (Fragment, error) {
	return f, nil
}

func (l *nativeStreamLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (LLMReply, LLMUsage, error) {
	return LLMReply{}, LLMUsage{}, nil
}

func (l *nativeStreamLLM) SetPendingNativeParts(parts []NativePart) {
	l.pending = parts
}

func (l *nativeStreamLLM) CreateChatCompletionStream(ctx context.Context, req openai.ChatCompletionRequest) (<-chan StreamEvent, error) {
	// Build a request body that embeds the stashed audio as an input_audio
	// content part, exactly the shape the real client emits on the wire.
	type inputAudio struct {
		Format string `json:"format"`
		Data   string `json:"data"`
	}
	type contentPart struct {
		Type       string      `json:"type"`
		Text       string      `json:"text,omitempty"`
		InputAudio *inputAudio `json:"input_audio,omitempty"`
	}
	var content []contentPart
	for _, p := range l.pending {
		if p.Kind == MediaAudio {
			content = append(content, contentPart{
				Type:       "input_audio",
				InputAudio: &inputAudio{Format: p.Format, Data: p.Data},
			})
		}
	}
	body, err := json.Marshal(map[string]any{"content": content})
	if err != nil {
		return nil, err
	}
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, l.url, strings.NewReader(string(body)))
	if err != nil {
		return nil, err
	}
	resp, err := http.DefaultClient.Do(httpReq)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, resp.Body)

	out := make(chan StreamEvent, 2)
	out <- StreamEvent{Type: StreamEventContent, Content: "ok"}
	out <- StreamEvent{Type: StreamEventDone}
	close(out)
	return out, nil
}

// TestExecuteToolsStreamingSerializesNativeAudio drives the REAL wrapped
// streaming seam: it wraps a StreamingLLM with newCountingLLM exactly as
// ExecuteTools does (tools.go:1227), then calls the unexported askWithStreaming
// seam and asserts the audio reaches the wire.
//
// Without the SetPendingNativeParts forwarder on the usage-counting wrapper
// (usage_counter.go), askWithStreaming's llm.(NativePartsAware) assertion fails
// on the *countingStreamingLLM wrapper, the stash is never set on the underlying
// client, and the audio is dropped — this test FAILS. With the forwarder the
// parts are forwarded to the client, serialized, and the test PASSES.
func TestExecuteToolsStreamingSerializesNativeAudio(t *testing.T) {
	var body string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		b, _ := io.ReadAll(r.Body)
		body = string(b)
		w.Header().Set("Content-Type", "text/event-stream")
		fl, _ := w.(http.Flusher)
		_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\n"))
		if fl != nil {
			fl.Flush()
		}
		_, _ = w.Write([]byte("data: [DONE]\n\n"))
		if fl != nil {
			fl.Flush()
		}
	}))
	defer srv.Close()

	client := &nativeStreamLLM{url: srv.URL}
	// Wrap exactly as the ExecuteTools loop does (tools.go:1227).
	wrapped := newCountingLLM(client, &usageCounter{})

	f := Fragment{}.AddMessage("user", "hear this", audioTM{data: "SND64", format: "wav"})

	// Drive the real wrapped streaming seam. A non-nil callback is required for
	// askWithStreaming to take the streaming path.
	if _, err := askWithStreaming(context.Background(), wrapped, f, func(StreamEvent) {}); err != nil {
		t.Fatalf("askWithStreaming: %v", err)
	}

	if !strings.Contains(body, `"type":"input_audio"`) || !strings.Contains(body, `"data":"SND64"`) {
		t.Fatalf("streaming request through the counting wrapper missing input_audio:\n%s", body)
	}
}