rony-llm-agent/pkg/llm/providers/llamacpp/client_test.go

package llamacpp

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

	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
)

func TestClient_Name(t *testing.T) {
	client := &Client{}
	if client.Name() != "llama.cpp" {
		t.Errorf("expected 'llama.cpp', got %q", client.Name())
	}
}

func TestClient_Capabilities(t *testing.T) {
	client := &Client{}
	caps := client.Capabilities()
	if !caps.SupportsTools {
		t.Error("expected SupportsTools to be true")
	}
	if caps.SupportsVision {
		t.Error("expected SupportsVision to be false")
	}
}

func TestClient_Generate(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(llamaChatResponse{
			ID:      "llama-123",
			Model:   "llama3",
			Choices: []llamaChoice{
				{
					Index:        0,
					FinishReason: "stop",
					Message: llamaMessageResult{
						Role:    "assistant",
						Content: "Hello from llama.cpp!",
					},
				},
			},
			Usage: llamaUsage{
				PromptTokens:     10,
				CompletionTokens: 5,
				TotalTokens:      15,
			},
		})
	}))
	defer server.Close()

	client, err := New(Config{
		BaseURL: server.URL + "/v1",
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	resp, err := client.Generate(context.Background(), llm.CompletionRequest{
		Model: "llama3",
		Messages: []llm.Message{
			{Role: llm.RoleUser, Content: "Hi"},
		},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if resp.Content != "Hello from llama.cpp!" {
		t.Errorf("expected 'Hello from llama.cpp!', got %q", resp.Content)
	}
	if resp.ID != "llama-123" {
		t.Errorf("expected 'llama-123', got %q", resp.ID)
	}
	if resp.Usage.InputTokens != 10 {
		t.Errorf("expected 10 input tokens, got %d", resp.Usage.InputTokens)
	}
}

func TestClient_Generate_Error(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("server error"))
	}))
	defer server.Close()

	client, err := New(Config{
		BaseURL: server.URL + "/v1",
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	_, err = client.Generate(context.Background(), llm.CompletionRequest{})
	if err == nil {
		t.Fatal("expected error, got nil")
	}
}

func TestClient_Generate_ToolCall(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(llamaChatResponse{
			ID:      "llama-tool-1",
			Model:   "llama3",
			Choices: []llamaChoice{
				{
					Index:        0,
					FinishReason: "tool_calls",
					Message: llamaMessageResult{
						Role:    "assistant",
						Content: "",
						ToolCalls: []llamaToolCall{
							{
								ID:   "call-1",
								Type: "function",
								Function: llamaFunction{
									Name:      "calculate",
									Arguments: `{"a":1,"b":2}`,
								},
							},
						},
					},
				},
			},
			Usage: llamaUsage{
				PromptTokens:     20,
				CompletionTokens: 10,
				TotalTokens:      30,
			},
		})
	}))
	defer server.Close()

	client, err := New(Config{
		BaseURL: server.URL + "/v1",
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	resp, err := client.Generate(context.Background(), llm.CompletionRequest{
		Model: "llama3",
		Messages: []llm.Message{
			{Role: llm.RoleUser, Content: "What is 1 + 2?"},
		},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(resp.ToolCalls) != 1 {
		t.Fatalf("expected 1 tool call, got %d", len(resp.ToolCalls))
	}
	if resp.ToolCalls[0].Name != "calculate" {
		t.Errorf("expected 'calculate', got %q", resp.ToolCalls[0].Name)
	}
}

func TestClient_Stream(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n"))
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\" world\"},\"finish_reason\":null}]}\n"))
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"!\"},\"finish_reason\":\"stop\"}]}\n"))
		w.Write([]byte("data: [DONE]\n"))
	}))
	defer server.Close()

	client, err := New(Config{
		BaseURL: server.URL + "/v1",
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	var chunks []llm.StreamChunk
	stream := client.Stream(context.Background(), llm.CompletionRequest{})
	for chunk, err := range stream {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		chunks = append(chunks, chunk)
	}

	if len(chunks) != 3 {
		t.Errorf("expected 3 chunks, got %d", len(chunks))
	}
	if chunks[0].Delta != "Hello" {
		t.Errorf("expected 'Hello', got %q", chunks[0].Delta)
	}
	if chunks[1].Delta != " world" {
		t.Errorf("expected ' world', got %q", chunks[1].Delta)
	}
	if chunks[2].Delta != "!" {
		t.Errorf("expected '!', got %q", chunks[2].Delta)
	}
}

func TestClient_Stream_FinishReason(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"done\"},\"finish_reason\":\"stop\"}]}\n"))
		w.Write([]byte("data: [DONE]\n"))
	}))
	defer server.Close()

	client, err := New(Config{
		BaseURL: server.URL + "/v1",
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	var chunks []llm.StreamChunk
	stream := client.Stream(context.Background(), llm.CompletionRequest{})
	for chunk, err := range stream {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		chunks = append(chunks, chunk)
	}

	if len(chunks) != 1 {
		t.Fatalf("expected 1 chunk, got %d", len(chunks))
	}
	if chunks[0].FinishReason != "stop" {
		t.Errorf("expected 'stop' finish reason, got %q", chunks[0].FinishReason)
	}
}

func TestClient_Stream_RequestsAndParsesUsage(t *testing.T) {
	var gotBody string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		gotBody = string(body)
		w.Header().Set("Content-Type", "text/event-stream")
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"},\"finish_reason\":\"stop\"}]}\n"))
		w.Write([]byte("data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":2,\"total_tokens\":12}}\n"))
		w.Write([]byte("data: [DONE]\n"))
	}))
	defer server.Close()

	client, err := New(Config{BaseURL: server.URL + "/v1"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	var chunks []llm.StreamChunk
	stream := client.Stream(context.Background(), llm.CompletionRequest{})
	for chunk, err := range stream {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		chunks = append(chunks, chunk)
	}

	if !strings.Contains(gotBody, `"stream_options":{"include_usage":true}`) {
		t.Errorf("expected the request to ask for usage via stream_options, got body: %s", gotBody)
	}

	if len(chunks) != 2 {
		t.Fatalf("expected 2 chunks (content + usage-only), got %d", len(chunks))
	}
	usage := chunks[len(chunks)-1].Usage
	if usage.InputTokens != 10 || usage.OutputTokens != 2 || usage.TotalTokens != 12 {
		t.Errorf("expected usage to be parsed from the final event, got %+v", usage)
	}
}