readysite / pkg / assistant / providers / anthropic / anthropic_test.go
8.6 KB
anthropic_test.go
package anthropic_test

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

	"github.com/readysite/readysite/pkg/assistant"
	"github.com/readysite/readysite/pkg/assistant/providers/anthropic"
)

func TestNew_EmptyKey(t *testing.T) {
	_, err := anthropic.New("")
	if !errors.Is(err, assistant.ErrNoAPIKey) {
		t.Fatalf("expected ErrNoAPIKey, got %v", err)
	}
}

func TestNew_ValidKey(t *testing.T) {
	ai, err := anthropic.New("sk-ant-test-key")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if ai == nil {
		t.Fatal("expected non-nil assistant")
	}
}

func TestChat(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Verify request
		if r.Method != "POST" {
			t.Errorf("expected POST, got %s", r.Method)
		}
		if r.URL.Path != "/messages" {
			t.Errorf("expected /messages, got %s", r.URL.Path)
		}
		if r.Header.Get("x-api-key") != "sk-ant-test-key" {
			t.Errorf("expected x-api-key 'sk-ant-test-key', got %q", r.Header.Get("x-api-key"))
		}
		if r.Header.Get("anthropic-version") != "2023-06-01" {
			t.Errorf("expected anthropic-version '2023-06-01', got %q", r.Header.Get("anthropic-version"))
		}
		if r.Header.Get("Content-Type") != "application/json" {
			t.Errorf("expected Content-Type 'application/json', got %q", r.Header.Get("Content-Type"))
		}

		// Verify request body
		body, _ := io.ReadAll(r.Body)
		var req map[string]any
		json.Unmarshal(body, &req)

		if req["model"] != "claude-3-opus-20240229" {
			t.Errorf("expected model claude-3-opus-20240229, got %v", req["model"])
		}

		// Anthropic requires max_tokens
		if req["max_tokens"] == nil {
			t.Error("expected max_tokens to be set")
		}

		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{
			"content": [
				{
					"type": "text",
					"text": "Hello from Anthropic!"
				}
			],
			"stop_reason": "end_turn",
			"usage": {
				"input_tokens": 12,
				"output_tokens": 18
			}
		}`)
	}))
	defer server.Close()

	ai, err := anthropic.New("sk-ant-test-key", anthropic.WithBaseURL(server.URL))
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	resp, err := ai.Chat(context.Background(), assistant.ChatRequest{
		Model: "claude-3-opus-20240229",
		Messages: []assistant.Message{
			assistant.NewUserMessage("Hello"),
		},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if resp.Content != "Hello from Anthropic!" {
		t.Errorf("expected 'Hello from Anthropic!', got %q", resp.Content)
	}
	if resp.FinishReason != "stop" {
		t.Errorf("expected finish reason 'stop', got %q", resp.FinishReason)
	}
	if resp.Usage.PromptTokens != 12 {
		t.Errorf("expected 12 prompt tokens, got %d", resp.Usage.PromptTokens)
	}
	if resp.Usage.CompletionTokens != 18 {
		t.Errorf("expected 18 completion tokens, got %d", resp.Usage.CompletionTokens)
	}
	if resp.Usage.TotalTokens != 30 {
		t.Errorf("expected 30 total tokens, got %d", resp.Usage.TotalTokens)
	}
}

func TestChat_Error(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusUnauthorized)
		fmt.Fprint(w, `{
			"error": {
				"type": "authentication_error",
				"message": "Invalid API key"
			}
		}`)
	}))
	defer server.Close()

	ai, err := anthropic.New("sk-ant-bad-key", anthropic.WithBaseURL(server.URL))
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	_, err = ai.Chat(context.Background(), assistant.ChatRequest{
		Model: "claude-3-opus-20240229",
		Messages: []assistant.Message{
			assistant.NewUserMessage("Hello"),
		},
	})
	if err == nil {
		t.Fatal("expected error, got nil")
	}

	var apiErr *assistant.APIError
	if !errors.As(err, &apiErr) {
		t.Fatalf("expected APIError, got %T: %v", err, err)
	}
	if apiErr.StatusCode != http.StatusUnauthorized {
		t.Errorf("expected status 401, got %d", apiErr.StatusCode)
	}
	if apiErr.Type != "authentication_error" {
		t.Errorf("expected type 'authentication_error', got %q", apiErr.Type)
	}
	if apiErr.Message != "Invalid API key" {
		t.Errorf("expected message 'Invalid API key', got %q", apiErr.Message)
	}

	// Verify IsAuthError works
	if !assistant.IsAuthError(err) {
		t.Error("expected IsAuthError to return true for 401")
	}
}

func TestChat_ToolCalls(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Verify tools were sent in request
		body, _ := io.ReadAll(r.Body)
		var req map[string]any
		json.Unmarshal(body, &req)

		tools, ok := req["tools"].([]any)
		if !ok || len(tools) != 1 {
			t.Errorf("expected 1 tool in request, got %v", req["tools"])
		}

		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{
			"content": [
				{
					"type": "text",
					"text": "I'll check the weather for you."
				},
				{
					"type": "tool_use",
					"id": "toolu_abc123",
					"name": "get_weather",
					"input": {"location": "San Francisco"}
				}
			],
			"stop_reason": "tool_use",
			"usage": {
				"input_tokens": 20,
				"output_tokens": 35
			}
		}`)
	}))
	defer server.Close()

	ai, err := anthropic.New("sk-ant-test-key", anthropic.WithBaseURL(server.URL))
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	resp, err := ai.Chat(context.Background(), assistant.ChatRequest{
		Model: "claude-3-opus-20240229",
		Messages: []assistant.Message{
			assistant.NewUserMessage("What is the weather in San Francisco?"),
		},
		Tools: []assistant.Tool{
			assistant.NewTool("get_weather", "Get the current weather").
				String("location", "City name", true).
				Build(),
		},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if resp.Content != "I'll check the weather for you." {
		t.Errorf("expected content 'I'll check the weather for you.', got %q", resp.Content)
	}

	if len(resp.ToolCalls) != 1 {
		t.Fatalf("expected 1 tool call, got %d", len(resp.ToolCalls))
	}
	tc := resp.ToolCalls[0]
	if tc.ID != "toolu_abc123" {
		t.Errorf("expected tool call ID 'toolu_abc123', got %q", tc.ID)
	}
	if tc.Name != "get_weather" {
		t.Errorf("expected tool name 'get_weather', got %q", tc.Name)
	}

	var args struct {
		Location string `json:"location"`
	}
	if err := tc.ParseArguments(&args); err != nil {
		t.Fatalf("failed to parse arguments: %v", err)
	}
	if args.Location != "San Francisco" {
		t.Errorf("expected location 'San Francisco', got %q", args.Location)
	}

	// Anthropic converts stop_reason "tool_use" to finish_reason "tool_calls"
	if resp.FinishReason != "tool_calls" {
		t.Errorf("expected finish reason 'tool_calls', got %q", resp.FinishReason)
	}
}

func TestStream(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Verify stream=true in request
		body, _ := io.ReadAll(r.Body)
		var req map[string]any
		json.Unmarshal(body, &req)

		if req["stream"] != true {
			t.Errorf("expected stream=true, got %v", req["stream"])
		}

		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")

		flusher, ok := w.(http.Flusher)
		if !ok {
			t.Fatal("expected ResponseWriter to implement Flusher")
		}

		// Send Anthropic SSE events
		events := []string{
			`data: {"type":"message_start","message":{"id":"msg_abc","type":"message","role":"assistant","content":[],"model":"claude-3-opus-20240229","usage":{"input_tokens":12,"output_tokens":0}}}`,
			`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
			`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}`,
			`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" "}}`,
			`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"World"}}`,
			`data: {"type":"content_block_stop","index":0}`,
			`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":12,"output_tokens":8}}`,
			`data: {"type":"message_stop"}`,
		}
		for _, event := range events {
			fmt.Fprintf(w, "%s\n\n", event)
			flusher.Flush()
		}
	}))
	defer server.Close()

	ai, err := anthropic.New("sk-ant-test-key", anthropic.WithBaseURL(server.URL))
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	stream, err := ai.Stream(context.Background(), assistant.ChatRequest{
		Model: "claude-3-opus-20240229",
		Messages: []assistant.Message{
			assistant.NewUserMessage("Hello"),
		},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	defer stream.Close()

	resp, err := stream.Collect()
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if resp.Content != "Hello World" {
		t.Errorf("expected 'Hello World', got %q", resp.Content)
	}
	if resp.FinishReason != "stop" {
		t.Errorf("expected finish reason 'stop', got %q", resp.FinishReason)
	}
}
← Back