diff --git a/.claude/settings.local.json b/.claude/settings.local.json index 8e1bf294..0ba6cf51 100644 --- a/.claude/settings.local.json +++ b/.claude/settings.local.json @@ -4,7 +4,8 @@ "Bash(make test-all)", "Bash(make lint)", "Bash(go test:*)", - "Bash(make test:*)" + "Bash(make test:*)", + "WebFetch(domain:docs.x.ai)" ], "deny": [], "ask": [] diff --git a/.env.template b/.env.template index 2d730267..af69acd1 100644 --- a/.env.template +++ b/.env.template @@ -34,3 +34,6 @@ CACHE_TYPE=local # Google Gemini # GEMINI_API_KEY=... + +# xAI (Grok) +# XAI_API_KEY=... diff --git a/EXAMPLES.md b/EXAMPLES.md index 55f98355..82ceae49 100644 --- a/EXAMPLES.md +++ b/EXAMPLES.md @@ -11,6 +11,7 @@ PORT=8080 OPENAI_API_KEY=sk-your-openai-api-key-here ANTHROPIC_API_KEY=sk-ant-your-anthropic-api-key-here GEMINI_API_KEY=your-gemini-api-key-here +XAI_API_KEY=xai-your-xai-api-key-here ``` Note: You only need to provide the API key for the provider(s) you want to use. @@ -187,6 +188,73 @@ curl http://localhost:8080/v1/chat/completions \ }' ``` +## xAI Examples + +### Basic Responses API Request + +```bash +curl http://localhost:8080/v1/responses \ + -H "Content-Type: application/json" \ + -d '{ + "model": "grok-4-1-fast-non-reasoning", + "input": "What is the capital of France?" + }' +``` + +### Responses API Request with Instructions + +```bash +curl http://localhost:8080/v1/responses \ + -H "Content-Type: application/json" \ + -d '{ + "model": "grok-4-1-fast-non-reasoning", + "input": "Write a haiku about programming.", + "instructions": "You are a creative AI assistant who specializes in writing poetry.", + "temperature": 0.8, + "max_output_tokens": 200 + }' +``` + +### Responses API with Structured Input + +```bash +curl http://localhost:8080/v1/responses \ + -H "Content-Type: application/json" \ + -d '{ + "model": "grok-4-1-fast-non-reasoning", + "input": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Explain quantum computing in simple terms."} + ], + "max_output_tokens": 500 + }' +``` + +### Streaming Responses + +```bash +curl http://localhost:8080/v1/responses \ + -H "Content-Type: application/json" \ + -N \ + -d '{ + "model": "grok-4-1-fast-non-reasoning", + "input": "Tell me a short story about AI.", + "stream": true + }' +``` + +### Using Grok 2 Mini (Faster Model) + +```bash +curl http://localhost:8080/v1/responses \ + -H "Content-Type: application/json" \ + -d '{ + "model": "grok-4-1-fast-non-reasoning", + "input": "What are the benefits of renewable energy?", + "max_output_tokens": 1000 + }' +``` + ## List Available Models Get a list of all available models from configured providers: @@ -263,6 +331,15 @@ response = client.chat.completions.create( ) print(response.choices[0].message.content) +# Or use xAI Grok models! +response = client.chat.completions.create( + model="grok-4-1-fast-non-reasoning", + messages=[ + {"role": "user", "content": "Hello!"} + ] +) +print(response.choices[0].message.content) + # Streaming works too stream = client.chat.completions.create( model="claude-3-5-haiku-20241022", @@ -278,39 +355,46 @@ for chunk in stream: ## Node.js Example ```javascript -import OpenAI from 'openai'; +import OpenAI from "openai"; const client = new OpenAI({ - baseURL: 'http://localhost:8080/v1', - apiKey: 'not-needed' // API key is configured on the server side + baseURL: "http://localhost:8080/v1", + apiKey: "not-needed", // API key is configured on the server side }); // Use OpenAI models const response1 = await client.chat.completions.create({ - model: 'gpt-4o-mini', - messages: [{ role: 'user', content: 'Hello!' }] + model: "gpt-4o-mini", + messages: [{ role: "user", content: "Hello!" }], }); console.log(response1.choices[0].message.content); // Or use Anthropic models with the same interface! const response2 = await client.chat.completions.create({ - model: 'claude-3-5-sonnet-20241022', - messages: [{ role: 'user', content: 'Hello!' }] + model: "claude-3-5-sonnet-20241022", + messages: [{ role: "user", content: "Hello!" }], }); console.log(response2.choices[0].message.content); // Or use Google Gemini models! const response3 = await client.chat.completions.create({ - model: 'gemini-2.0-flash', - messages: [{ role: 'user', content: 'Hello!' }] + model: "gemini-2.0-flash", + messages: [{ role: "user", content: "Hello!" }], }); console.log(response3.choices[0].message.content); +// Or use xAI Grok models! +const response4 = await client.chat.completions.create({ + model: "grok-4-1-fast-non-reasoning", + messages: [{ role: "user", content: "Hello!" }], +}); +console.log(response4.choices[0].message.content); + // Streaming const stream = await client.chat.completions.create({ - model: 'claude-3-5-haiku-20241022', - messages: [{ role: 'user', content: 'Tell me a story' }], - stream: true + model: "claude-3-5-haiku-20241022", + messages: [{ role: "user", content: "Tell me a story" }], + stream: true, }); for await (const chunk of stream) { @@ -323,6 +407,7 @@ for await (const chunk of stream) { ## Available Models ### OpenAI + - `gpt-4o` - Most capable GPT-4 model - `gpt-4o-mini` - Fast and efficient GPT-4 model - `gpt-4-turbo` - Previous generation GPT-4 Turbo @@ -332,6 +417,7 @@ for await (const chunk of stream) { - `o1-mini` - Faster reasoning model ### Anthropic + - `claude-3-5-sonnet-20241022` - Latest Sonnet (best balance of speed and capability) - `claude-3-5-sonnet-20240620` - Previous Sonnet version - `claude-3-5-haiku-20241022` - Latest Haiku (fastest, most cost-effective) @@ -340,31 +426,44 @@ for await (const chunk of stream) { - `claude-3-haiku-20240307` - Previous Haiku generation ### Google Gemini + - `gemini-2.0-flash` - Latest Flash model (fast and efficient) - `gemini-1.5-pro` - Most capable Gemini model (large context window) - `gemini-1.5-flash` - Previous generation Flash model - `gemini-1.0-pro` - Original Gemini Pro model +### xAI + +- `grok-4-1-fast-non-reasoning` - Most capable Grok model (best for complex reasoning) + ## Tips 1. **Model Selection**: The gateway automatically routes requests to the correct provider based on the model name prefix: + - `gpt-` or `o1` → OpenAI - `claude-` → Anthropic - `gemini-` → Google Gemini + - `grok-` → xAI 2. **API Compatibility**: The gateway exposes an OpenAI-compatible API, so you can use existing OpenAI client libraries to access OpenAI, Anthropic, and Google Gemini models. 3. **Streaming**: All providers support streaming responses. The gateway automatically converts provider-specific streaming formats to match OpenAI's SSE format. 4. **System Messages**: + - Anthropic handles system messages differently. The gateway automatically extracts system messages and sends them in Anthropic's required format. - Gemini uses Google's OpenAI-compatible endpoint which handles system messages natively. 5. **Max Tokens**: + - Anthropic requires `max_tokens` to be set. If not provided, the gateway uses a default of 4096 tokens. - Gemini and OpenAI have optional `max_tokens` parameters. 6. **Context Windows**: + - Gemini 1.5 Pro offers an exceptionally large context window (up to 1M tokens), making it ideal for long-form content analysis. - OpenAI GPT-4 models typically support 8K-128K tokens depending on the variant. - Anthropic Claude models support up to 200K tokens. + - xAI Grok models support up to 128K tokens. + +7. **Responses API**: The `/v1/responses` endpoint is OpenAI-compatible and provides a unified interface across all providers. xAI internally converts Responses API requests to chat completion format since it doesn't natively support the Responses API. diff --git a/cmd/gomodel/main.go b/cmd/gomodel/main.go index b9d82f81..390f07ce 100644 --- a/cmd/gomodel/main.go +++ b/cmd/gomodel/main.go @@ -18,6 +18,7 @@ import ( _ "gomodel/internal/providers/anthropic" _ "gomodel/internal/providers/gemini" _ "gomodel/internal/providers/openai" + _ "gomodel/internal/providers/xai" "gomodel/internal/server" ) diff --git a/config/config.go b/config/config.go index aa2f5699..516f04d7 100644 --- a/config/config.go +++ b/config/config.go @@ -138,6 +138,12 @@ func Load() (*Config, error) { APIKey: apiKey, } } + if apiKey := viper.GetString("XAI_API_KEY"); apiKey != "" { + cfg.Providers["xai-primary"] = ProviderConfig{ + Type: "xai", + APIKey: apiKey, + } + } } return &cfg, nil diff --git a/config/config.yaml b/config/config.yaml index c36fe0d2..529631f2 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -44,6 +44,10 @@ providers: type: "gemini" api_key: "${GEMINI_API_KEY}" + xai-primary: + type: "xai" + api_key: "${XAI_API_KEY}" + # Example: Groq (OpenAI-compatible) # groq: # type: "openai" diff --git a/internal/core/responses.go b/internal/core/responses.go index cc22211c..c1ee4d03 100644 --- a/internal/core/responses.go +++ b/internal/core/responses.go @@ -13,6 +13,21 @@ type ResponsesRequest struct { Metadata map[string]string `json:"metadata,omitempty"` } +// WithStreaming returns a shallow copy of the request with Stream set to true. +// This avoids mutating the caller's request object. +func (r *ResponsesRequest) WithStreaming() *ResponsesRequest { + return &ResponsesRequest{ + Model: r.Model, + Input: r.Input, + Instructions: r.Instructions, + Tools: r.Tools, + Temperature: r.Temperature, + MaxOutputTokens: r.MaxOutputTokens, + Stream: true, + Metadata: r.Metadata, + } +} + // ResponsesInputItem represents an input item when Input is an array. type ResponsesInputItem struct { Role string `json:"role"` diff --git a/internal/core/types.go b/internal/core/types.go index d522bc4d..6294e2d1 100644 --- a/internal/core/types.go +++ b/internal/core/types.go @@ -9,6 +9,18 @@ type ChatRequest struct { Stream bool `json:"stream,omitempty"` } +// WithStreaming returns a shallow copy of the request with Stream set to true. +// This avoids mutating the caller's request object. +func (r *ChatRequest) WithStreaming() *ChatRequest { + return &ChatRequest{ + Temperature: r.Temperature, + MaxTokens: r.MaxTokens, + Model: r.Model, + Messages: r.Messages, + Stream: true, + } +} + // Message represents a single message in the chat type Message struct { Role string `json:"role"` diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index 9cc0f778..f863dff9 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -88,11 +88,10 @@ func (p *Provider) ChatCompletion(ctx context.Context, req *core.ChatRequest) (* // StreamChatCompletion returns a raw response body for streaming (caller must close) func (p *Provider) StreamChatCompletion(ctx context.Context, req *core.ChatRequest) (io.ReadCloser, error) { - req.Stream = true stream, err := p.client.DoStream(ctx, llmclient.Request{ Method: http.MethodPost, Endpoint: "/chat/completions", - Body: req, + Body: req.WithStreaming(), }) if err != nil { return nil, err @@ -302,7 +301,6 @@ func (p *Provider) Responses(ctx context.Context, req *core.ResponsesRequest) (* func (p *Provider) StreamResponses(ctx context.Context, req *core.ResponsesRequest) (io.ReadCloser, error) { // Convert ResponsesRequest to ChatRequest chatReq := convertResponsesRequestToChat(req) - chatReq.Stream = true // Get the streaming response from chat completions stream, err := p.StreamChatCompletion(ctx, chatReq) diff --git a/internal/providers/openai/openai.go b/internal/providers/openai/openai.go index 5f811315..7c8df723 100644 --- a/internal/providers/openai/openai.go +++ b/internal/providers/openai/openai.go @@ -72,11 +72,10 @@ func (p *Provider) ChatCompletion(ctx context.Context, req *core.ChatRequest) (* // StreamChatCompletion returns a raw response body for streaming (caller must close) func (p *Provider) StreamChatCompletion(ctx context.Context, req *core.ChatRequest) (io.ReadCloser, error) { - req.Stream = true return p.client.DoStream(ctx, llmclient.Request{ Method: http.MethodPost, Endpoint: "/chat/completions", - Body: req, + Body: req.WithStreaming(), }) } @@ -109,10 +108,9 @@ func (p *Provider) Responses(ctx context.Context, req *core.ResponsesRequest) (* // StreamResponses returns a raw response body for streaming Responses API (caller must close) func (p *Provider) StreamResponses(ctx context.Context, req *core.ResponsesRequest) (io.ReadCloser, error) { - req.Stream = true return p.client.DoStream(ctx, llmclient.Request{ Method: http.MethodPost, Endpoint: "/responses", - Body: req, + Body: req.WithStreaming(), }) } diff --git a/internal/providers/xai/xai.go b/internal/providers/xai/xai.go new file mode 100644 index 00000000..e4b4f324 --- /dev/null +++ b/internal/providers/xai/xai.go @@ -0,0 +1,116 @@ +// Package xai provides xAI (Grok) API integration for the LLM gateway. +package xai + +import ( + "context" + "io" + "net/http" + + "gomodel/internal/core" + "gomodel/internal/pkg/llmclient" + "gomodel/internal/providers" +) + +const ( + defaultBaseURL = "https://api.x.ai/v1" +) + +func init() { + // Self-register with the factory + providers.RegisterProvider("xai", New) +} + +// Provider implements the core.Provider interface for xAI +type Provider struct { + client *llmclient.Client + apiKey string +} + +// New creates a new xAI provider +func New(apiKey string) *Provider { + p := &Provider{apiKey: apiKey} + cfg := llmclient.DefaultConfig("xai", defaultBaseURL) + // Apply global hooks if available + cfg.Hooks = providers.GetGlobalHooks() + p.client = llmclient.New(cfg, p.setHeaders) + return p +} + +// NewWithHTTPClient creates a new xAI provider with a custom HTTP client +func NewWithHTTPClient(apiKey string, httpClient *http.Client) *Provider { + p := &Provider{apiKey: apiKey} + cfg := llmclient.DefaultConfig("xai", defaultBaseURL) + // Apply global hooks if available + cfg.Hooks = providers.GetGlobalHooks() + p.client = llmclient.NewWithHTTPClient(httpClient, cfg, p.setHeaders) + return p +} + +// SetBaseURL allows configuring a custom base URL for the provider +func (p *Provider) SetBaseURL(url string) { + p.client.SetBaseURL(url) +} + +// setHeaders sets the required headers for xAI API requests +func (p *Provider) setHeaders(req *http.Request) { + req.Header.Set("Authorization", "Bearer "+p.apiKey) +} + +// ChatCompletion sends a chat completion request to xAI +func (p *Provider) ChatCompletion(ctx context.Context, req *core.ChatRequest) (*core.ChatResponse, error) { + var resp core.ChatResponse + err := p.client.Do(ctx, llmclient.Request{ + Method: http.MethodPost, + Endpoint: "/chat/completions", + Body: req, + }, &resp) + if err != nil { + return nil, err + } + return &resp, nil +} + +// StreamChatCompletion returns a raw response body for streaming (caller must close) +func (p *Provider) StreamChatCompletion(ctx context.Context, req *core.ChatRequest) (io.ReadCloser, error) { + return p.client.DoStream(ctx, llmclient.Request{ + Method: http.MethodPost, + Endpoint: "/chat/completions", + Body: req.WithStreaming(), + }) +} + +// ListModels retrieves the list of available models from xAI +func (p *Provider) ListModels(ctx context.Context) (*core.ModelsResponse, error) { + var resp core.ModelsResponse + err := p.client.Do(ctx, llmclient.Request{ + Method: http.MethodGet, + Endpoint: "/models", + }, &resp) + if err != nil { + return nil, err + } + return &resp, nil +} + +// Responses sends a Responses API request to xAI +func (p *Provider) Responses(ctx context.Context, req *core.ResponsesRequest) (*core.ResponsesResponse, error) { + var resp core.ResponsesResponse + err := p.client.Do(ctx, llmclient.Request{ + Method: http.MethodPost, + Endpoint: "/responses", + Body: req, + }, &resp) + if err != nil { + return nil, err + } + return &resp, nil +} + +// StreamResponses returns a raw response body for streaming Responses API (caller must close) +func (p *Provider) StreamResponses(ctx context.Context, req *core.ResponsesRequest) (io.ReadCloser, error) { + return p.client.DoStream(ctx, llmclient.Request{ + Method: http.MethodPost, + Endpoint: "/responses", + Body: req.WithStreaming(), + }) +} diff --git a/internal/providers/xai/xai_test.go b/internal/providers/xai/xai_test.go new file mode 100644 index 00000000..90e1b7b3 --- /dev/null +++ b/internal/providers/xai/xai_test.go @@ -0,0 +1,760 @@ +package xai + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "gomodel/internal/core" +) + +func TestNew(t *testing.T) { + apiKey := "test-api-key" + provider := New(apiKey) + + if provider.apiKey != apiKey { + t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) + } + if provider.client == nil { + t.Error("client should not be nil") + } +} + +// customHeaderRoundTripper is a RoundTripper that injects a custom header +type customHeaderRoundTripper struct { + transport http.RoundTripper + headerKey string + headerVal string +} + +func (c *customHeaderRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + req.Header.Set(c.headerKey, c.headerVal) + return c.transport.RoundTrip(req) +} + +func TestNewWithHTTPClient(t *testing.T) { + const customHeaderKey = "X-Custom-Test-Header" + const customHeaderVal = "custom-test-value" + + var receivedHeader string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Capture the custom header to verify custom client was used + receivedHeader = r.Header.Get(customHeaderKey) + + // Return a valid chat completion response + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "grok-2", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + }`)) + })) + defer server.Close() + + // Create a custom HTTP client with a RoundTripper that injects a header + customClient := &http.Client{ + Transport: &customHeaderRoundTripper{ + transport: http.DefaultTransport, + headerKey: customHeaderKey, + headerVal: customHeaderVal, + }, + } + + // Create provider with custom HTTP client + provider := NewWithHTTPClient("test-api-key", customClient) + + // Verify provider is non-nil + if provider == nil { + t.Fatal("provider should not be nil") + } + if provider.client == nil { + t.Fatal("provider.client should not be nil") + } + if provider.apiKey != "test-api-key" { + t.Errorf("apiKey = %q, want %q", provider.apiKey, "test-api-key") + } + + // Set base URL to our test server + provider.SetBaseURL(server.URL) + + // Make a request to verify custom client is wired correctly + req := &core.ChatRequest{ + Model: "grok-2", + Messages: []core.Message{ + {Role: "user", Content: "Hello"}, + }, + } + + resp, err := provider.ChatCompletion(context.Background(), req) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Verify response is valid (server was hit) + if resp == nil { + t.Fatal("response should not be nil") + } + if resp.ID != "chatcmpl-123" { + t.Errorf("response ID = %q, want %q", resp.ID, "chatcmpl-123") + } + + // Verify the custom header was injected by our custom RoundTripper + if receivedHeader != customHeaderVal { + t.Errorf("custom header = %q, want %q (custom HTTP client not wired correctly)", receivedHeader, customHeaderVal) + } +} + +func TestChatCompletion(t *testing.T) { + tests := []struct { + name string + statusCode int + responseBody string + expectedError bool + checkResponse func(*testing.T, *core.ChatResponse) + }{ + { + name: "successful request", + statusCode: http.StatusOK, + responseBody: `{ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "grok-2", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": "Hello! How can I help you today?" + }, + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30 + } + }`, + expectedError: false, + checkResponse: func(t *testing.T, resp *core.ChatResponse) { + if resp.ID != "chatcmpl-123" { + t.Errorf("ID = %q, want %q", resp.ID, "chatcmpl-123") + } + if resp.Model != "grok-2" { + t.Errorf("Model = %q, want %q", resp.Model, "grok-2") + } + if len(resp.Choices) != 1 { + t.Fatalf("len(Choices) = %d, want 1", len(resp.Choices)) + } + if resp.Choices[0].Message.Content != "Hello! How can I help you today?" { + t.Errorf("Message content = %q, want %q", resp.Choices[0].Message.Content, "Hello! How can I help you today?") + } + if resp.Usage.PromptTokens != 10 { + t.Errorf("PromptTokens = %d, want 10", resp.Usage.PromptTokens) + } + if resp.Usage.CompletionTokens != 20 { + t.Errorf("CompletionTokens = %d, want 20", resp.Usage.CompletionTokens) + } + if resp.Usage.TotalTokens != 30 { + t.Errorf("TotalTokens = %d, want 30", resp.Usage.TotalTokens) + } + }, + }, + { + name: "API error", + statusCode: http.StatusUnauthorized, + responseBody: `{"error": {"message": "Invalid API key"}}`, + expectedError: true, + }, + { + name: "rate limit error", + statusCode: http.StatusTooManyRequests, + responseBody: `{"error": {"message": "Rate limit exceeded"}}`, + expectedError: true, + }, + { + name: "server error", + statusCode: http.StatusInternalServerError, + responseBody: `{"error": {"message": "Internal server error"}}`, + expectedError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Verify request headers + if r.Header.Get("Content-Type") != "application/json" { + t.Errorf("Content-Type = %q, want %q", r.Header.Get("Content-Type"), "application/json") + } + authHeader := r.Header.Get("Authorization") + if !strings.HasPrefix(authHeader, "Bearer ") { + t.Errorf("Authorization header should start with 'Bearer '") + } + + // Verify request body + body, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("failed to read request body: %v", err) + } + var req core.ChatRequest + if err := json.Unmarshal(body, &req); err != nil { + t.Fatalf("failed to unmarshal request: %v", err) + } + + w.WriteHeader(tt.statusCode) + _, _ = w.Write([]byte(tt.responseBody)) + })) + defer server.Close() + + provider := New("test-api-key") + provider.SetBaseURL(server.URL) + + req := &core.ChatRequest{ + Model: "grok-2", + Messages: []core.Message{ + {Role: "user", Content: "Hello"}, + }, + } + + resp, err := provider.ChatCompletion(context.Background(), req) + + if tt.expectedError { + if err == nil { + t.Error("expected error, got nil") + } + } else { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if tt.checkResponse != nil { + tt.checkResponse(t, resp) + } + } + }) + } +} + +func TestStreamChatCompletion(t *testing.T) { + tests := []struct { + name string + statusCode int + responseBody string + expectedError bool + }{ + { + name: "successful streaming request", + statusCode: http.StatusOK, + responseBody: `data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1677652288,"model":"grok-2","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]} + +data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1677652288,"model":"grok-2","choices":[{"index":0,"delta":{"content":"!"},"finish_reason":null}]} + +data: [DONE] +`, + expectedError: false, + }, + { + name: "API error", + statusCode: http.StatusUnauthorized, + responseBody: `{"error": {"message": "Invalid API key"}}`, + expectedError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Verify request headers + if r.Header.Get("Content-Type") != "application/json" { + t.Errorf("Content-Type = %q, want %q", r.Header.Get("Content-Type"), "application/json") + } + authHeader := r.Header.Get("Authorization") + if !strings.HasPrefix(authHeader, "Bearer ") { + t.Errorf("Authorization header should start with 'Bearer '") + } + + // Verify stream is set in request body + body, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("failed to read request body: %v", err) + } + var req core.ChatRequest + if err := json.Unmarshal(body, &req); err != nil { + t.Fatalf("failed to unmarshal request: %v", err) + } + if !req.Stream { + t.Error("Stream should be true in request") + } + + w.WriteHeader(tt.statusCode) + _, _ = w.Write([]byte(tt.responseBody)) + })) + defer server.Close() + + provider := New("test-api-key") + provider.SetBaseURL(server.URL) + + req := &core.ChatRequest{ + Model: "grok-2", + Messages: []core.Message{ + {Role: "user", Content: "Hello"}, + }, + } + + body, err := provider.StreamChatCompletion(context.Background(), req) + + if tt.expectedError { + if err == nil { + t.Error("expected error, got nil") + } + } else { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if body == nil { + t.Fatal("body should not be nil") + } + defer func() { _ = body.Close() }() + + // Read and verify the streaming response + respBody, err := io.ReadAll(body) + if err != nil { + t.Fatalf("failed to read response body: %v", err) + } + if string(respBody) != tt.responseBody { + t.Errorf("response body = %q, want %q", string(respBody), tt.responseBody) + } + } + }) + } +} + +func TestListModels(t *testing.T) { + tests := []struct { + name string + statusCode int + responseBody string + expectedError bool + checkResponse func(*testing.T, *core.ModelsResponse) + }{ + { + name: "successful request", + statusCode: http.StatusOK, + responseBody: `{ + "object": "list", + "data": [ + { + "id": "grok-2", + "object": "model", + "created": 1687882411, + "owned_by": "xai" + }, + { + "id": "grok-2-mini", + "object": "model", + "created": 1687882410, + "owned_by": "xai" + } + ] + }`, + expectedError: false, + checkResponse: func(t *testing.T, resp *core.ModelsResponse) { + if resp.Object != "list" { + t.Errorf("Object = %q, want %q", resp.Object, "list") + } + if len(resp.Data) != 2 { + t.Fatalf("len(Data) = %d, want 2", len(resp.Data)) + } + if resp.Data[0].ID != "grok-2" { + t.Errorf("Data[0].ID = %q, want %q", resp.Data[0].ID, "grok-2") + } + if resp.Data[0].OwnedBy != "xai" { + t.Errorf("Data[0].OwnedBy = %q, want %q", resp.Data[0].OwnedBy, "xai") + } + }, + }, + { + name: "API error", + statusCode: http.StatusUnauthorized, + responseBody: `{"error": {"message": "Invalid API key"}}`, + expectedError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Verify request method and path + if r.Method != http.MethodGet { + t.Errorf("Method = %q, want %q", r.Method, http.MethodGet) + } + if r.URL.Path != "/models" { + t.Errorf("Path = %q, want %q", r.URL.Path, "/models") + } + + // Verify authorization header + authHeader := r.Header.Get("Authorization") + if !strings.HasPrefix(authHeader, "Bearer ") { + t.Errorf("Authorization header should start with 'Bearer '") + } + + w.WriteHeader(tt.statusCode) + _, _ = w.Write([]byte(tt.responseBody)) + })) + defer server.Close() + + provider := New("test-api-key") + provider.SetBaseURL(server.URL) + + resp, err := provider.ListModels(context.Background()) + + if tt.expectedError { + if err == nil { + t.Error("expected error, got nil") + } + } else { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if tt.checkResponse != nil { + tt.checkResponse(t, resp) + } + } + }) + } +} + +func TestChatCompletionWithContext(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Simulate a slow response + <-r.Context().Done() + w.WriteHeader(http.StatusRequestTimeout) + })) + defer server.Close() + + provider := New("test-api-key") + provider.SetBaseURL(server.URL) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // Cancel immediately + + req := &core.ChatRequest{ + Model: "grok-2", + Messages: []core.Message{ + {Role: "user", Content: "Hello"}, + }, + } + + _, err := provider.ChatCompletion(ctx, req) + if err == nil { + t.Error("expected error when context is cancelled, got nil") + } +} + +func TestResponses(t *testing.T) { + tests := []struct { + name string + statusCode int + responseBody string + expectedError bool + checkResponse func(*testing.T, *core.ResponsesResponse) + }{ + { + name: "successful request with string input", + statusCode: http.StatusOK, + responseBody: `{ + "id": "resp_123", + "object": "response", + "created_at": 1677652288, + "model": "grok-2", + "status": "completed", + "output": [{ + "id": "msg_123", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{ + "type": "output_text", + "text": "Hello! How can I help you today?" + }] + }], + "usage": { + "input_tokens": 10, + "output_tokens": 20, + "total_tokens": 30 + } + }`, + expectedError: false, + checkResponse: func(t *testing.T, resp *core.ResponsesResponse) { + if resp.ID != "resp_123" { + t.Errorf("ID = %q, want %q", resp.ID, "resp_123") + } + if resp.Object != "response" { + t.Errorf("Object = %q, want %q", resp.Object, "response") + } + if resp.Model != "grok-2" { + t.Errorf("Model = %q, want %q", resp.Model, "grok-2") + } + if resp.Status != "completed" { + t.Errorf("Status = %q, want %q", resp.Status, "completed") + } + if len(resp.Output) != 1 { + t.Fatalf("len(Output) = %d, want 1", len(resp.Output)) + } + if len(resp.Output[0].Content) != 1 { + t.Fatalf("len(Output[0].Content) = %d, want 1", len(resp.Output[0].Content)) + } + if resp.Output[0].Content[0].Text != "Hello! How can I help you today?" { + t.Errorf("Output text = %q, want %q", resp.Output[0].Content[0].Text, "Hello! How can I help you today?") + } + if resp.Usage == nil { + t.Fatal("Usage should not be nil") + } + if resp.Usage.InputTokens != 10 { + t.Errorf("InputTokens = %d, want 10", resp.Usage.InputTokens) + } + if resp.Usage.OutputTokens != 20 { + t.Errorf("OutputTokens = %d, want 20", resp.Usage.OutputTokens) + } + if resp.Usage.TotalTokens != 30 { + t.Errorf("TotalTokens = %d, want 30", resp.Usage.TotalTokens) + } + }, + }, + { + name: "API error - unauthorized", + statusCode: http.StatusUnauthorized, + responseBody: `{"error": {"message": "Invalid API key"}}`, + expectedError: true, + }, + { + name: "rate limit error", + statusCode: http.StatusTooManyRequests, + responseBody: `{"error": {"message": "Rate limit exceeded"}}`, + expectedError: true, + }, + { + name: "server error", + statusCode: http.StatusInternalServerError, + responseBody: `{"error": {"message": "Internal server error"}}`, + expectedError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Verify request headers + if r.Header.Get("Content-Type") != "application/json" { + t.Errorf("Content-Type = %q, want %q", r.Header.Get("Content-Type"), "application/json") + } + authHeader := r.Header.Get("Authorization") + if !strings.HasPrefix(authHeader, "Bearer ") { + t.Errorf("Authorization header should start with 'Bearer '") + } + + // Verify request path + if r.URL.Path != "/responses" { + t.Errorf("Path = %q, want %q", r.URL.Path, "/responses") + } + + // Verify request body + body, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("failed to read request body: %v", err) + } + var req core.ResponsesRequest + if err := json.Unmarshal(body, &req); err != nil { + t.Fatalf("failed to unmarshal request: %v", err) + } + + w.WriteHeader(tt.statusCode) + _, _ = w.Write([]byte(tt.responseBody)) + })) + defer server.Close() + + provider := New("test-api-key") + provider.SetBaseURL(server.URL) + + req := &core.ResponsesRequest{ + Model: "grok-2", + Input: "Hello", + } + + resp, err := provider.Responses(context.Background(), req) + + if tt.expectedError { + if err == nil { + t.Error("expected error, got nil") + } + } else { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if tt.checkResponse != nil { + tt.checkResponse(t, resp) + } + } + }) + } +} + +func TestStreamResponses(t *testing.T) { + tests := []struct { + name string + statusCode int + responseBody string + expectedError bool + checkStream func(*testing.T, io.ReadCloser) + }{ + { + name: "successful streaming request", + statusCode: http.StatusOK, + responseBody: `event: response.created +data: {"type":"response.created","response":{"id":"resp_123","object":"response","status":"in_progress","model":"grok-2"}} + +event: response.output_text.delta +data: {"type":"response.output_text.delta","delta":"Hello"} + +event: response.output_text.delta +data: {"type":"response.output_text.delta","delta":"!"} + +event: response.done +data: {"type":"response.done","response":{"id":"resp_123","object":"response","status":"completed","model":"grok-2"}} + +data: [DONE] +`, + expectedError: false, + checkStream: func(t *testing.T, body io.ReadCloser) { + if body == nil { + t.Fatal("body should not be nil") + } + defer func() { _ = body.Close() }() + + // Read and verify the streaming response + respBody, err := io.ReadAll(body) + if err != nil { + t.Fatalf("failed to read response body: %v", err) + } + + responseStr := string(respBody) + if !strings.Contains(responseStr, "response.created") { + t.Error("response should contain response.created event") + } + if !strings.Contains(responseStr, "response.output_text.delta") { + t.Error("response should contain response.output_text.delta event") + } + if !strings.Contains(responseStr, "[DONE]") { + t.Error("response should end with [DONE]") + } + }, + }, + { + name: "API error - unauthorized", + statusCode: http.StatusUnauthorized, + responseBody: `{"error": {"message": "Invalid API key"}}`, + expectedError: true, + }, + { + name: "rate limit error", + statusCode: http.StatusTooManyRequests, + responseBody: `{"error": {"message": "Rate limit exceeded"}}`, + expectedError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Verify request headers + if r.Header.Get("Content-Type") != "application/json" { + t.Errorf("Content-Type = %q, want %q", r.Header.Get("Content-Type"), "application/json") + } + authHeader := r.Header.Get("Authorization") + if !strings.HasPrefix(authHeader, "Bearer ") { + t.Errorf("Authorization header should start with 'Bearer '") + } + + // Verify request path + if r.URL.Path != "/responses" { + t.Errorf("Path = %q, want %q", r.URL.Path, "/responses") + } + + // Verify stream is set in request body + body, err := io.ReadAll(r.Body) + if err != nil { + t.Fatalf("failed to read request body: %v", err) + } + var req core.ResponsesRequest + if err := json.Unmarshal(body, &req); err != nil { + t.Fatalf("failed to unmarshal request: %v", err) + } + if !req.Stream { + t.Error("Stream should be true in request") + } + + w.WriteHeader(tt.statusCode) + _, _ = w.Write([]byte(tt.responseBody)) + })) + defer server.Close() + + provider := New("test-api-key") + provider.SetBaseURL(server.URL) + + req := &core.ResponsesRequest{ + Model: "grok-2", + Input: "Hello", + } + + body, err := provider.StreamResponses(context.Background(), req) + + if tt.expectedError { + if err == nil { + t.Error("expected error, got nil") + } + } else { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if tt.checkStream != nil { + tt.checkStream(t, body) + } + } + }) + } +} + +func TestResponsesWithContext(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Simulate a slow response + <-r.Context().Done() + w.WriteHeader(http.StatusRequestTimeout) + })) + defer server.Close() + + provider := New("test-api-key") + provider.SetBaseURL(server.URL) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // Cancel immediately + + req := &core.ResponsesRequest{ + Model: "grok-2", + Input: "Hello", + } + + _, err := provider.Responses(ctx, req) + if err == nil { + t.Error("expected error when context is cancelled, got nil") + } +}