diff --git a/cmd/cmd.go b/cmd/cmd.go index b6183a8..9ff7492 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -35,7 +35,7 @@ const ( // sensitiveConfigKeys lists the config keys that should be stored in the // secure credential store rather than in the plaintext YAML config file. -var sensitiveConfigKeys = []string{"openai.api_key", "gemini.api_key"} +var sensitiveConfigKeys = []string{"openai.api_key", "gemini.api_key", "litellm.api_key"} // migrateCredentialsToStore moves any plaintext API keys found in the YAML // config into the secure credential store and clears them from the config file. diff --git a/cmd/config_list.go b/cmd/config_list.go index 6ee6572..e65f235 100644 --- a/cmd/config_list.go +++ b/cmd/config_list.go @@ -18,37 +18,41 @@ func init() { // availableKeys is a map of configuration keys and their descriptions var availableKeys = map[string]string{ - "git.diff_unified": "Number of context lines in git diff output (default: 3)", - "git.exclude_list": "Files to exclude from git diff command", - "git.template_file": "Path to template file for commit messages", - "git.template_string": "Template string for formatting commit messages", - "openai.socks": "SOCKS proxy URL for API connections", - "openai.api_key": "Authentication key for OpenAI API access", - "openai.api_key_helper": "Shell command to dynamically generate API key", - "openai.api_key_helper_refresh_interval": "Interval in seconds to refresh credentials from apiKeyHelper (default: 900)", - "openai.model": "AI model identifier to use for requests", - "openai.org_id": "Organization ID for multi-org OpenAI accounts", - "openai.proxy": "HTTP proxy URL for API connections", - "output.lang": "Language for summarization output (default: English)", - "openai.base_url": "Custom base URL for API requests", - "openai.timeout": "Maximum duration to wait for API response", - "openai.max_tokens": "Maximum token limit for generated completions", - "openai.temperature": "Randomness control parameter (0-1): lower values for focused results, higher for creative variety", - "openai.provider": "Service provider selection ('openai' or 'azure')", - "openai.skip_verify": "Option to bypass TLS certificate verification", - "openai.headers": "Additional custom HTTP headers for API requests", - "openai.api_version": "Specific API version to target", - "openai.top_p": "Nucleus sampling parameter: controls diversity by limiting to top percentage of probability mass", - "openai.frequency_penalty": "Parameter to reduce repetition by penalizing tokens based on their frequency", - "openai.presence_penalty": "Parameter to encourage topic diversity by penalizing previously used tokens", - "openai.stream": "Enable streaming output for real-time token display", - "prompt.folder": "Directory path for custom prompt templates", - "gemini.project_id": "VertexAI project for Gemini provider", - "gemini.location": "VertexAI location for Gemini provider", - "gemini.backend": "Gemini backend (BackendGeminiAPI or BackendVertexAI)", - "gemini.api_key": "API key for Gemini provider", - "gemini.api_key_helper": "Shell command to dynamically generate Gemini API key", - "gemini.api_key_helper_refresh_interval": "Interval in seconds to refresh Gemini credentials from apiKeyHelper (default: 900)", + "git.diff_unified": "Number of context lines in git diff output (default: 3)", + "git.exclude_list": "Files to exclude from git diff command", + "git.template_file": "Path to template file for commit messages", + "git.template_string": "Template string for formatting commit messages", + "openai.socks": "SOCKS proxy URL for API connections", + "openai.api_key": "Authentication key for OpenAI API access", + "openai.api_key_helper": "Shell command to dynamically generate API key", + "openai.api_key_helper_refresh_interval": "Interval in seconds to refresh credentials from apiKeyHelper (default: 900)", + "openai.model": "AI model identifier to use for requests", + "openai.org_id": "Organization ID for multi-org OpenAI accounts", + "openai.proxy": "HTTP proxy URL for API connections", + "output.lang": "Language for summarization output (default: English)", + "openai.base_url": "Custom base URL for API requests", + "openai.timeout": "Maximum duration to wait for API response", + "openai.max_tokens": "Maximum token limit for generated completions", + "openai.temperature": "Randomness control parameter (0-1): lower values for focused results, higher for creative variety", + "openai.provider": "Service provider selection ('openai' or 'azure')", + "openai.skip_verify": "Option to bypass TLS certificate verification", + "openai.headers": "Additional custom HTTP headers for API requests", + "openai.api_version": "Specific API version to target", + "openai.top_p": "Nucleus sampling parameter: controls diversity by limiting to top percentage of probability mass", + "openai.frequency_penalty": "Parameter to reduce repetition by penalizing tokens based on their frequency", + "openai.presence_penalty": "Parameter to encourage topic diversity by penalizing previously used tokens", + "openai.stream": "Enable streaming output for real-time token display", + "prompt.folder": "Directory path for custom prompt templates", + "gemini.project_id": "VertexAI project for Gemini provider", + "gemini.location": "VertexAI location for Gemini provider", + "gemini.backend": "Gemini backend (BackendGeminiAPI or BackendVertexAI)", + "gemini.api_key": "API key for Gemini provider", + "gemini.api_key_helper": "Shell command to dynamically generate Gemini API key", + "gemini.api_key_helper_refresh_interval": "Interval in seconds to refresh Gemini credentials from apiKeyHelper (default: 900)", + "litellm.api_key": "API key for LiteLLM proxy (master key or virtual key)", + "litellm.api_key_helper": "Shell command to dynamically generate LiteLLM API key", + "litellm.api_key_helper_refresh_interval": "Interval in seconds to refresh LiteLLM credentials from apiKeyHelper (default: 900)", + "litellm.base_url": "LiteLLM proxy base URL (default: http://localhost:4000/v1)", } // configListCmd represents the command to list the configuration values. diff --git a/cmd/config_set.go b/cmd/config_set.go index ad25983..fda18c4 100644 --- a/cmd/config_set.go +++ b/cmd/config_set.go @@ -50,6 +50,14 @@ func init() { configSetCmd.Flags().String("gemini.api_key_helper", "", availableKeys["gemini.api_key_helper"]) configSetCmd.Flags(). Int("gemini.api_key_helper_refresh_interval", 900, availableKeys["gemini.api_key_helper_refresh_interval"]) + // LiteLLM flags + configSetCmd.Flags().String("litellm.api_key", "", availableKeys["litellm.api_key"]) + configSetCmd.Flags(). + String("litellm.api_key_helper", "", availableKeys["litellm.api_key_helper"]) + configSetCmd.Flags(). + Int("litellm.api_key_helper_refresh_interval", 900, availableKeys["litellm.api_key_helper_refresh_interval"]) + configSetCmd.Flags(). + String("litellm.base_url", "http://localhost:4000/v1", availableKeys["litellm.base_url"]) _ = viper.BindPFlag("openai.base_url", configSetCmd.Flags().Lookup("base_url")) _ = viper.BindPFlag("openai.org_id", configSetCmd.Flags().Lookup("org_id")) @@ -87,6 +95,16 @@ func init() { "gemini.api_key_helper_refresh_interval", configSetCmd.Flags().Lookup("gemini.api_key_helper_refresh_interval"), ) + _ = viper.BindPFlag("litellm.api_key", configSetCmd.Flags().Lookup("litellm.api_key")) + _ = viper.BindPFlag( + "litellm.api_key_helper", + configSetCmd.Flags().Lookup("litellm.api_key_helper"), + ) + _ = viper.BindPFlag( + "litellm.api_key_helper_refresh_interval", + configSetCmd.Flags().Lookup("litellm.api_key_helper_refresh_interval"), + ) + _ = viper.BindPFlag("litellm.base_url", configSetCmd.Flags().Lookup("litellm.base_url")) } // configSetCmd updates the config value. diff --git a/cmd/provider.go b/cmd/provider.go index 63546b6..d415b82 100644 --- a/cmd/provider.go +++ b/cmd/provider.go @@ -8,6 +8,7 @@ import ( "github.com/appleboy/CodeGPT/core" "github.com/appleboy/CodeGPT/provider/anthropic" "github.com/appleboy/CodeGPT/provider/gemini" + "github.com/appleboy/CodeGPT/provider/litellm" "github.com/appleboy/CodeGPT/provider/openai" "github.com/appleboy/CodeGPT/util" @@ -205,6 +206,63 @@ func GetClient(ctx context.Context, p core.Platform) (core.Generative, error) { return NewOpenAI(ctx) case core.Anthropic: return NewAnthropic(ctx) + case core.LiteLLM: + return NewLiteLLM(ctx) } return nil, errors.New("invalid provider") } + +// NewLiteLLM creates a new LiteLLM client that connects to a LiteLLM proxy server, +// providing access to 100+ LLM providers through a unified OpenAI-compatible API. +func NewLiteLLM(ctx context.Context) (*litellm.Client, error) { + _ = ctx + + var apiKey string + + // Try litellm-specific key helper first + if helper := viper.GetString("litellm.api_key_helper"); helper != "" { + var refreshInterval time.Duration + if viper.IsSet("litellm.api_key_helper_refresh_interval") { + refreshInterval = time.Duration( + viper.GetInt("litellm.api_key_helper_refresh_interval"), + ) * time.Second + } else { + refreshInterval = util.DefaultRefreshInterval + } + key, err := util.GetAPIKeyFromHelperWithCache(ctx, helper, refreshInterval) + if err != nil { + return nil, err + } + apiKey = key + } else { + // Try litellm.api_key first, fall back to openai.api_key + key, err := getAPIKey("litellm.api_key") + if err != nil { + return nil, err + } + apiKey = key + if apiKey == "" { + openaiKey, err := getAPIKey("openai.api_key") + if err != nil { + return nil, err + } + apiKey = openaiKey + } + } + + return litellm.New( + litellm.WithToken(apiKey), + litellm.WithModel(viper.GetString("openai.model")), + litellm.WithBaseURL(viper.GetString("litellm.base_url")), + litellm.WithProxyURL(viper.GetString("openai.proxy")), + litellm.WithSocksURL(viper.GetString("openai.socks")), + litellm.WithTimeout(viper.GetDuration("openai.timeout")), + litellm.WithMaxTokens(viper.GetInt("openai.max_tokens")), + litellm.WithTemperature(float32(viper.GetFloat64("openai.temperature"))), + litellm.WithTopP(float32(viper.GetFloat64("openai.top_p"))), + litellm.WithFrequencyPenalty(float32(viper.GetFloat64("openai.frequency_penalty"))), + litellm.WithPresencePenalty(float32(viper.GetFloat64("openai.presence_penalty"))), + litellm.WithSkipVerify(viper.GetBool("openai.skip_verify")), + litellm.WithHeaders(viper.GetStringSlice("openai.headers")), + ) +} diff --git a/core/platform.go b/core/platform.go index 0621830..9b80b6c 100644 --- a/core/platform.go +++ b/core/platform.go @@ -12,6 +12,8 @@ const ( Gemini Platform = "gemini" // Anthropic represents the Anthropic platform. Anthropic Platform = "anthropic" + // LiteLLM represents the LiteLLM AI gateway platform. + LiteLLM Platform = "litellm" ) // String returns the string representation of the Platform. @@ -22,7 +24,7 @@ func (p Platform) String() string { // IsValid returns true if the Platform is valid. func (p Platform) IsValid() bool { switch p { - case OpenAI, Azure, Gemini, Anthropic: + case OpenAI, Azure, Gemini, Anthropic, LiteLLM: return true } return false diff --git a/provider/litellm/edge_test.go b/provider/litellm/edge_test.go new file mode 100644 index 0000000..6dcb556 --- /dev/null +++ b/provider/litellm/edge_test.go @@ -0,0 +1,607 @@ +package litellm + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// --- Auth / API key errors --- + +func TestCompletion_InvalidAPIKey(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "message": "Incorrect API key provided: sk-inva*****key.", + "type": "invalid_request_error", + "code": "invalid_api_key", + }, + }) + })) + defer server.Close() + + client, err := New( + WithToken("sk-invalid-key"), + WithModel("anthropic/claude-sonnet-4-6"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + _, err = client.Completion(context.Background(), "test") + if err == nil { + t.Fatal("expected error for invalid API key, got nil") + } + if !strings.Contains(err.Error(), "401") { + t.Errorf("expected 401 in error, got: %v", err) + } +} + +func TestCompletionStream_InvalidAPIKey(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "message": "Invalid API key", + "type": "invalid_request_error", + }, + }) + })) + defer server.Close() + + client, err := New( + WithToken("sk-bad"), + WithModel("openai/gpt-4o"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + var buf bytes.Buffer + _, err = client.CompletionStream(context.Background(), "test", &buf) + if err == nil { + t.Fatal("expected error for invalid API key on stream, got nil") + } +} + +// --- Model not found --- + +func TestCompletion_ModelNotFound(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotFound) + json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "message": "Model 'nonexistent/model' not found. Check your LiteLLM proxy config.", + "type": "invalid_request_error", + "code": "model_not_found", + }, + }) + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("nonexistent/model"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + _, err = client.Completion(context.Background(), "test") + if err == nil { + t.Fatal("expected error for model not found, got nil") + } + if !strings.Contains(err.Error(), "404") { + t.Errorf("expected 404 in error, got: %v", err) + } +} + +// --- Rate limit (429) --- + +func TestCompletion_RateLimit(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", "1") + w.WriteHeader(http.StatusTooManyRequests) + json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "message": "Rate limit reached for anthropic/claude-sonnet-4-6", + "type": "rate_limit_error", + }, + }) + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("anthropic/claude-sonnet-4-6"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + _, err = client.Completion(context.Background(), "test") + if err == nil { + t.Fatal("expected error for rate limit, got nil") + } + if !strings.Contains(err.Error(), "429") { + t.Errorf("expected 429 in error, got: %v", err) + } +} + +// --- Context timeout --- + +func TestCompletion_Timeout(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(2 * time.Second) + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{ + {"message": map[string]any{"content": "late"}}, + }, + }) + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("openai/gpt-4o"), + WithBaseURL(server.URL+"/v1"), + WithTimeout(100*time.Millisecond), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + + _, err = client.Completion(ctx, "test") + if err == nil { + t.Fatal("expected timeout error, got nil") + } +} + +// --- Empty / malformed responses --- + +func TestCompletion_EmptyContent(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(map[string]any{ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "openai/gpt-4o", + "choices": []map[string]any{ + { + "index": 0, + "message": map[string]any{"role": "assistant", "content": ""}, + "finish_reason": "stop", + }, + }, + "usage": map[string]any{"prompt_tokens": 5, "completion_tokens": 0, "total_tokens": 5}, + }) + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("openai/gpt-4o"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + resp, err := client.Completion(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "" { + t.Errorf("expected empty content, got %q", resp.Content) + } +} + +func TestCompletion_NullContent(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write( + []byte( + `{"id":"1","object":"chat.completion","created":1,"model":"test","choices":[{"index":0,"message":{"role":"assistant","content":null},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":0,"total_tokens":1}}`, + ), + ) + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("test-model"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + resp, err := client.Completion(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error on null content: %v", err) + } + if resp.Content != "" { + t.Errorf("expected empty content for null, got %q", resp.Content) + } +} + +func TestCompletion_MalformedJSON(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{invalid json`)) + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("test-model"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + _, err = client.Completion(context.Background(), "test") + if err == nil { + t.Fatal("expected error for malformed JSON, got nil") + } +} + +// --- Streaming edge cases --- + +func TestCompletionStream_AbruptTermination(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + + fmt.Fprintf( + w, + "data: %s\n\n", + `{"id":"1","object":"chat.completion.chunk","created":1,"model":"test","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}`, + ) + // Server closes connection without sending [DONE] + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("test-model"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + var buf bytes.Buffer + resp, err := client.CompletionStream(context.Background(), "test", &buf) + // Abrupt close triggers EOF which is handled gracefully + if err != nil { + t.Fatalf("unexpected error on abrupt stream: %v", err) + } + if resp.Content != "Hello" { + t.Errorf("expected partial content %q, got %q", "Hello", resp.Content) + } +} + +func TestCompletionStream_EmptyChunks(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + chunks := []string{ + `{"id":"1","object":"chat.completion.chunk","created":1,"model":"test","choices":[{"index":0,"delta":{"content":""},"finish_reason":null}]}`, + `{"id":"1","object":"chat.completion.chunk","created":1,"model":"test","choices":[{"index":0,"delta":{"content":"World"},"finish_reason":null}]}`, + `{"id":"1","object":"chat.completion.chunk","created":1,"model":"test","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":1,"total_tokens":6}}`, + } + for _, chunk := range chunks { + fmt.Fprintf(w, "data: %s\n\n", chunk) + } + fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("test-model"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + var buf bytes.Buffer + resp, err := client.CompletionStream(context.Background(), "test", &buf) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "World" { + t.Errorf("expected %q, got %q", "World", resp.Content) + } +} + +func TestCompletionStream_ServerError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "message": "Internal server error", + "type": "server_error", + }, + }) + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("test-model"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + var buf bytes.Buffer + _, err = client.CompletionStream(context.Background(), "test", &buf) + if err == nil { + t.Fatal("expected error for server error on stream, got nil") + } +} + +// --- Reasoning content fallback --- + +func TestCompletion_ReasoningContentFallback(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write( + []byte( + `{"id":"1","object":"chat.completion","created":1,"model":"openai/o3","choices":[{"index":0,"message":{"role":"assistant","content":"","reasoning_content":"The answer is 4 because 2+2=4"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":8,"total_tokens":18}}`, + ), + ) + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("openai/o3"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + resp, err := client.Completion(context.Background(), "What is 2+2?") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "The answer is 4 because 2+2=4" { + t.Errorf("expected reasoning content fallback, got %q", resp.Content) + } +} + +func TestCompletionStream_ReasoningContent(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + chunks := []string{ + `{"id":"1","object":"chat.completion.chunk","created":1,"model":"openai/o3","choices":[{"index":0,"delta":{"reasoning_content":"thinking..."},"finish_reason":null}]}`, + `{"id":"1","object":"chat.completion.chunk","created":1,"model":"openai/o3","choices":[{"index":0,"delta":{"content":"4"},"finish_reason":null}]}`, + `{"id":"1","object":"chat.completion.chunk","created":1,"model":"openai/o3","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}`, + } + for _, chunk := range chunks { + fmt.Fprintf(w, "data: %s\n\n", chunk) + } + fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("openai/o3"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + var buf bytes.Buffer + resp, err := client.CompletionStream(context.Background(), "2+2?", &buf) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "thinking...4" { + t.Errorf("expected reasoning+content, got %q", resp.Content) + } +} + +// --- Model string format verification --- + +func TestModelStringPassedCorrectly(t *testing.T) { + var receivedModel string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req map[string]any + json.NewDecoder(r.Body).Decode(&req) + receivedModel, _ = req["model"].(string) + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "id": "1", + "object": "chat.completion", + "created": 1, + "model": receivedModel, + "choices": []map[string]any{ + { + "index": 0, + "message": map[string]any{"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + }, + }, + "usage": map[string]any{"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + })) + defer server.Close() + + models := []string{ + "anthropic/claude-sonnet-4-6", + "openai/gpt-4o", + "bedrock/anthropic.claude-sonnet-4-6-v1", + "vertex_ai/gemini-2.5-flash", + "groq/llama-4-scout-17b-16e-instruct", + } + + for _, model := range models { + t.Run(model, func(t *testing.T) { + client, err := New( + WithToken("test-key"), + WithModel(model), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + _, err = client.Completion(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if receivedModel != model { + t.Errorf("model sent to proxy = %q, want %q", receivedModel, model) + } + }) + } +} + +// --- API key forwarding --- + +func TestAPIKeyForwarding(t *testing.T) { + var receivedAuth string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedAuth = r.Header.Get("Authorization") + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "id": "1", + "object": "chat.completion", + "created": 1, + "model": "test", + "choices": []map[string]any{ + { + "index": 0, + "message": map[string]any{"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + }, + }, + "usage": map[string]any{"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + })) + defer server.Close() + + client, err := New( + WithToken("sk-litellm-master-key-123"), + WithModel("anthropic/claude-sonnet-4-6"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + _, err = client.Completion(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if receivedAuth != "Bearer sk-litellm-master-key-123" { + t.Errorf("expected Bearer auth header, got %q", receivedAuth) + } +} + +// --- Context cancellation --- + +func TestCompletion_ContextCancelled(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(5 * time.Second) + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("test-model"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err = client.Completion(ctx, "test") + if err == nil { + t.Fatal("expected error for cancelled context, got nil") + } +} + +// --- GetSummaryPrefix with no tool calls in response --- + +func TestGetSummaryPrefix_NoToolCallsReturnsContent(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(map[string]any{ + "id": "1", + "object": "chat.completion", + "created": 1, + "model": "test", + "choices": []map[string]any{ + { + "index": 0, + "message": map[string]any{ + "role": "assistant", + "content": "feat(auth): add OAuth2 support", + }, + "finish_reason": "stop", + }, + }, + "usage": map[string]any{ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + }) + })) + defer server.Close() + + client, err := New( + WithToken("test-key"), + WithModel("test-model"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + resp, err := client.GetSummaryPrefix(context.Background(), "test diff") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "feat(auth): add OAuth2 support" { + t.Errorf("expected plain content fallback, got %q", resp.Content) + } +} diff --git a/provider/litellm/func.go b/provider/litellm/func.go new file mode 100644 index 0000000..ae51141 --- /dev/null +++ b/provider/litellm/func.go @@ -0,0 +1,49 @@ +package litellm + +import ( + "encoding/json" + + "github.com/appleboy/com/bytesconv" + "github.com/sashabaranov/go-openai/jsonschema" + + openai "github.com/sashabaranov/go-openai" +) + +// summaryPrefixFunc defines the OpenAI function-calling schema for extracting +// a conventional commit prefix and scope from a diff summary. +var summaryPrefixFunc = openai.FunctionDefinition{ + Name: "get_summary_prefix", + Parameters: jsonschema.Definition{ + Type: jsonschema.Object, + Properties: map[string]jsonschema.Definition{ + "prefix": { + Type: jsonschema.String, + Enum: []string{ + "build", "chore", "ci", + "docs", "feat", "fix", + "perf", "refactor", "style", + "test", + }, + }, + "scope": { + Type: jsonschema.String, + Description: "A short lowercase word identifying the module, package, or component most central to the change", + }, + }, + Required: []string{"prefix", "scope"}, + }, +} + +type summaryPrefixParams struct { + Prefix string `json:"prefix"` + Scope string `json:"scope"` +} + +func getSummaryPrefixArgs(data string) summaryPrefixParams { + var prefix summaryPrefixParams + err := json.Unmarshal(bytesconv.StrToBytes(data), &prefix) + if err != nil { + panic(err) + } + return prefix +} diff --git a/provider/litellm/func_test.go b/provider/litellm/func_test.go new file mode 100644 index 0000000..ec1e555 --- /dev/null +++ b/provider/litellm/func_test.go @@ -0,0 +1,21 @@ +package litellm + +import ( + "reflect" + "testing" +) + +func TestGetSummaryPrefixArgs(t *testing.T) { + data := `{"prefix": "feat", "scope": "provider", "param2": "value2"}` + + result := getSummaryPrefixArgs(data) + + expected := summaryPrefixParams{ + Prefix: "feat", + Scope: "provider", + } + + if !reflect.DeepEqual(result, expected) { + t.Errorf("Expected %v, but got %v", expected, result) + } +} diff --git a/provider/litellm/litellm.go b/provider/litellm/litellm.go new file mode 100644 index 0000000..5b2e7c3 --- /dev/null +++ b/provider/litellm/litellm.go @@ -0,0 +1,224 @@ +package litellm + +import ( + "context" + "errors" + "fmt" + "io" + "strings" + + "github.com/appleboy/CodeGPT/core" + "github.com/appleboy/CodeGPT/core/transport" + "github.com/appleboy/CodeGPT/proxy" + "github.com/appleboy/CodeGPT/version" + + openai "github.com/sashabaranov/go-openai" +) + +var _ core.Generative = (*Client)(nil) + +// Client wraps an OpenAI-compatible client configured to talk to a LiteLLM proxy. +// LiteLLM normalizes 100+ provider APIs (Anthropic, Bedrock, Vertex, Groq, etc.) +// into the OpenAI Chat Completions format, so the same go-openai SDK works unchanged. +type Client struct { + client *openai.Client + model string + maxTokens int + temperature float32 + topP float32 + frequencyPenalty float32 + presencePenalty float32 +} + +// newBaseRequest builds a ChatCompletionRequest with the client's model parameters. +func (c *Client) newBaseRequest(content string) openai.ChatCompletionRequest { + return openai.ChatCompletionRequest{ + Model: c.model, + MaxCompletionTokens: c.maxTokens, + Temperature: c.temperature, + TopP: c.topP, + FrequencyPenalty: c.frequencyPenalty, + PresencePenalty: c.presencePenalty, + Messages: []openai.ChatCompletionMessage{ + { + Role: openai.ChatMessageRoleAssistant, + Content: "You are a helpful assistant.", + }, + { + Role: openai.ChatMessageRoleUser, + Content: content, + }, + }, + } +} + +// convertUsage converts an openai.Usage to a core.Usage. +func convertUsage(u openai.Usage) core.Usage { + return core.Usage{ + PromptTokens: u.PromptTokens, + CompletionTokens: u.CompletionTokens, + TotalTokens: u.TotalTokens, + CompletionTokensDetails: u.CompletionTokensDetails, + PromptTokensDetails: u.PromptTokensDetails, + } +} + +// Completion generates a non-streaming completion via the LiteLLM proxy. +func (c *Client) Completion(ctx context.Context, content string) (*core.Response, error) { + resp, err := c.client.CreateChatCompletion(ctx, c.newBaseRequest(content)) + if err != nil { + return nil, err + } + if len(resp.Choices) == 0 { + return nil, errors.New("no choices returned from LiteLLM proxy") + } + + text := resp.Choices[0].Message.Content + if text == "" && resp.Choices[0].Message.ReasoningContent != "" { + text = resp.Choices[0].Message.ReasoningContent + } + + return &core.Response{ + Content: text, + Usage: convertUsage(resp.Usage), + }, nil +} + +// CompletionStream streams completion tokens to the writer as they arrive. +func (c *Client) CompletionStream( + ctx context.Context, + content string, + w io.Writer, +) (*core.Response, error) { + req := c.newBaseRequest(content) + req.Stream = true + req.StreamOptions = &openai.StreamOptions{IncludeUsage: true} + + stream, err := c.client.CreateChatCompletionStream(ctx, req) + if err != nil { + return nil, err + } + defer stream.Close() + + var sb strings.Builder + var usage openai.Usage + for { + chunk, err := stream.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + return nil, err + } + + if chunk.Usage != nil { + usage = *chunk.Usage + } + + if len(chunk.Choices) > 0 { + text := chunk.Choices[0].Delta.Content + if text == "" { + text = chunk.Choices[0].Delta.ReasoningContent + } + if text != "" { + sb.WriteString(text) + if _, err := io.WriteString(w, text); err != nil { + return nil, fmt.Errorf("writing streamed completion: %w", err) + } + } + } + } + + return &core.Response{ + Content: sb.String(), + Usage: convertUsage(usage), + }, nil +} + +// GetSummaryPrefix uses OpenAI-format function calling through the LiteLLM proxy +// to extract a conventional commit prefix. Falls back to plain completion if the +// backing model does not support tool/function calls. +func (c *Client) GetSummaryPrefix(ctx context.Context, content string) (*core.Response, error) { + req := c.newBaseRequest(content) + req.Tools = []openai.Tool{ + { + Type: openai.ToolTypeFunction, + Function: &summaryPrefixFunc, + }, + } + req.ToolChoice = openai.ToolChoice{ + Type: openai.ToolTypeFunction, + Function: openai.ToolFunction{ + Name: summaryPrefixFunc.Name, + }, + } + + resp, err := c.client.CreateChatCompletion(ctx, req) + if err != nil { + // Function calling not supported by the backing model; fall back to plain completion. + return c.Completion(ctx, content) + } + + if len(resp.Choices) == 0 { + return nil, errors.New("no choices returned from LiteLLM proxy") + } + + msg := resp.Choices[0].Message + usage := convertUsage(resp.Usage) + + if len(msg.ToolCalls) == 0 { + return &core.Response{ + Content: msg.Content, + Usage: usage, + }, nil + } + + args := getSummaryPrefixArgs(msg.ToolCalls[len(msg.ToolCalls)-1].Function.Arguments) + return &core.Response{ + Content: fmt.Sprintf("%s(%s)", args.Prefix, args.Scope), + Usage: usage, + }, nil +} + +// New creates a new LiteLLM client that connects to a LiteLLM proxy server. +func New(opts ...Option) (*Client, error) { + cfg := newConfig(opts...) + if err := cfg.valid(); err != nil { + return nil, err + } + + engine := &Client{ + model: cfg.model, + maxTokens: cfg.maxTokens, + temperature: cfg.temperature, + topP: cfg.topP, + frequencyPenalty: cfg.frequencyPenalty, + presencePenalty: cfg.presencePenalty, + } + + c := openai.DefaultConfig(cfg.token) + c.BaseURL = cfg.baseURL + + httpClient, err := proxy.New( + proxy.WithProxyURL(cfg.proxyURL), + proxy.WithSocksURL(cfg.socksURL), + proxy.WithSkipVerify(cfg.skipVerify), + proxy.WithTimeout(cfg.timeout), + proxy.WithHeaders(cfg.headers), + ) + if err != nil { + return nil, fmt.Errorf("can't create a new HTTP client: %w", err) + } + + httpClient.Transport = &transport.DefaultHeaderTransport{ + Origin: httpClient.Transport, + Header: nil, + AppName: version.App, + AppVersion: version.Version, + } + + c.HTTPClient = httpClient + engine.client = openai.NewClientWithConfig(c) + + return engine, nil +} diff --git a/provider/litellm/options.go b/provider/litellm/options.go new file mode 100644 index 0000000..4d705d2 --- /dev/null +++ b/provider/litellm/options.go @@ -0,0 +1,168 @@ +package litellm + +import ( + "errors" + "time" +) + +var ( + errorsMissingToken = errors.New("please set LITELLM_API_KEY environment variable") + errorsMissingModel = errors.New("missing model") +) + +const ( + defaultMaxTokens = 300 + defaultTemperature = float32(1.0) + defaultTopP = float32(1.0) + defaultBaseURL = "http://localhost:4000/v1" +) + +// Option is an interface that specifies instrumentation configuration options. +type Option interface { + apply(*config) +} + +// optionFunc is a type of function that can be used to implement the Option interface. +type optionFunc func(*config) + +var _ Option = (*optionFunc)(nil) + +func (o optionFunc) apply(c *config) { + o(c) +} + +// WithToken sets the API key for the LiteLLM proxy. +func WithToken(val string) Option { + return optionFunc(func(c *config) { + c.token = val + }) +} + +// WithModel sets the model identifier (e.g. "anthropic/claude-sonnet-4-6", "openai/gpt-4o"). +func WithModel(val string) Option { + return optionFunc(func(c *config) { + c.model = val + }) +} + +// WithBaseURL sets the LiteLLM proxy base URL. +func WithBaseURL(val string) Option { + return optionFunc(func(c *config) { + c.baseURL = val + }) +} + +// WithProxyURL sets the HTTP proxy URL for outbound connections. +func WithProxyURL(val string) Option { + return optionFunc(func(c *config) { + c.proxyURL = val + }) +} + +// WithSocksURL sets the SOCKS proxy URL for outbound connections. +func WithSocksURL(val string) Option { + return optionFunc(func(c *config) { + c.socksURL = val + }) +} + +// WithTimeout sets the per-request HTTP timeout. +func WithTimeout(val time.Duration) Option { + return optionFunc(func(c *config) { + c.timeout = val + }) +} + +// WithMaxTokens sets the maximum token limit for generated completions. +func WithMaxTokens(val int) Option { + if val <= 0 { + val = defaultMaxTokens + } + return optionFunc(func(c *config) { + c.maxTokens = val + }) +} + +// WithTemperature sets the sampling temperature (0-2). +func WithTemperature(val float32) Option { + if val <= 0 { + val = defaultTemperature + } + return optionFunc(func(c *config) { + c.temperature = val + }) +} + +// WithTopP sets the nucleus sampling parameter. +func WithTopP(val float32) Option { + return optionFunc(func(c *config) { + c.topP = val + }) +} + +// WithFrequencyPenalty sets the frequency penalty parameter. +func WithFrequencyPenalty(val float32) Option { + return optionFunc(func(c *config) { + c.frequencyPenalty = val + }) +} + +// WithPresencePenalty sets the presence penalty parameter. +func WithPresencePenalty(val float32) Option { + return optionFunc(func(c *config) { + c.presencePenalty = val + }) +} + +// WithSkipVerify disables TLS certificate verification. +func WithSkipVerify(val bool) Option { + return optionFunc(func(c *config) { + c.skipVerify = val + }) +} + +// WithHeaders sets additional HTTP headers for requests. +func WithHeaders(headers []string) Option { + return optionFunc(func(c *config) { + c.headers = headers + }) +} + +type config struct { + token string + model string + baseURL string + proxyURL string + socksURL string + timeout time.Duration + maxTokens int + temperature float32 + topP float32 + frequencyPenalty float32 + presencePenalty float32 + skipVerify bool + headers []string +} + +func (cfg *config) valid() error { + if cfg.token == "" { + return errorsMissingToken + } + if cfg.model == "" { + return errorsMissingModel + } + return nil +} + +func newConfig(opts ...Option) *config { + c := &config{ + maxTokens: defaultMaxTokens, + temperature: defaultTemperature, + topP: defaultTopP, + baseURL: defaultBaseURL, + } + for _, opt := range opts { + opt.apply(c) + } + return c +} diff --git a/provider/litellm/options_test.go b/provider/litellm/options_test.go new file mode 100644 index 0000000..74ef302 --- /dev/null +++ b/provider/litellm/options_test.go @@ -0,0 +1,148 @@ +package litellm + +import ( + "testing" + "time" +) + +func Test_config_valid(t *testing.T) { + tests := []struct { + name string + cfg *config + wantErr error + }{ + { + name: "valid config", + cfg: newConfig( + WithToken("test-key"), + WithModel("anthropic/claude-sonnet-4-6"), + ), + wantErr: nil, + }, + { + name: "missing token", + cfg: newConfig(), + wantErr: errorsMissingToken, + }, + { + name: "missing model", + cfg: newConfig( + WithToken("test-key"), + WithModel(""), + ), + wantErr: errorsMissingModel, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if err := tt.cfg.valid(); err != tt.wantErr { + t.Errorf("config.valid() error = %v, wantErr %v", err, tt.wantErr) + } + }) + } +} + +func Test_config_defaults(t *testing.T) { + cfg := newConfig() + if cfg.baseURL != defaultBaseURL { + t.Errorf("expected default baseURL %q, got %q", defaultBaseURL, cfg.baseURL) + } + if cfg.maxTokens != defaultMaxTokens { + t.Errorf("expected default maxTokens %d, got %d", defaultMaxTokens, cfg.maxTokens) + } + if cfg.temperature != defaultTemperature { + t.Errorf("expected default temperature %f, got %f", defaultTemperature, cfg.temperature) + } + if cfg.topP != defaultTopP { + t.Errorf("expected default topP %f, got %f", defaultTopP, cfg.topP) + } +} + +func Test_config_options(t *testing.T) { + cfg := newConfig( + WithToken("sk-litellm-key"), + WithModel("openai/gpt-4o"), + WithBaseURL("http://myproxy:8000/v1"), + WithProxyURL("http://proxy:8080"), + WithSocksURL("socks5://proxy:1080"), + WithTimeout(30*time.Second), + WithMaxTokens(500), + WithTemperature(0.7), + WithTopP(0.9), + WithFrequencyPenalty(0.5), + WithPresencePenalty(0.3), + WithSkipVerify(true), + WithHeaders([]string{"X-Custom: value"}), + ) + + if cfg.token != "sk-litellm-key" { + t.Errorf("expected token %q, got %q", "sk-litellm-key", cfg.token) + } + if cfg.model != "openai/gpt-4o" { + t.Errorf("expected model %q, got %q", "openai/gpt-4o", cfg.model) + } + if cfg.baseURL != "http://myproxy:8000/v1" { + t.Errorf("expected baseURL %q, got %q", "http://myproxy:8000/v1", cfg.baseURL) + } + if cfg.proxyURL != "http://proxy:8080" { + t.Errorf("expected proxyURL %q, got %q", "http://proxy:8080", cfg.proxyURL) + } + if cfg.socksURL != "socks5://proxy:1080" { + t.Errorf("expected socksURL %q, got %q", "socks5://proxy:1080", cfg.socksURL) + } + if cfg.timeout != 30*time.Second { + t.Errorf("expected timeout %v, got %v", 30*time.Second, cfg.timeout) + } + if cfg.maxTokens != 500 { + t.Errorf("expected maxTokens %d, got %d", 500, cfg.maxTokens) + } + if cfg.temperature != 0.7 { + t.Errorf("expected temperature %f, got %f", float32(0.7), cfg.temperature) + } + if cfg.topP != 0.9 { + t.Errorf("expected topP %f, got %f", float32(0.9), cfg.topP) + } + if cfg.frequencyPenalty != 0.5 { + t.Errorf("expected frequencyPenalty %f, got %f", float32(0.5), cfg.frequencyPenalty) + } + if cfg.presencePenalty != 0.3 { + t.Errorf("expected presencePenalty %f, got %f", float32(0.3), cfg.presencePenalty) + } + if !cfg.skipVerify { + t.Error("expected skipVerify true") + } + if len(cfg.headers) != 1 || cfg.headers[0] != "X-Custom: value" { + t.Errorf("expected headers [X-Custom: value], got %v", cfg.headers) + } +} + +func TestNew_missingToken(t *testing.T) { + _, err := New(WithModel("test-model")) + if err != errorsMissingToken { + t.Errorf("expected errorsMissingToken, got %v", err) + } +} + +func TestNew_missingModel(t *testing.T) { + _, err := New(WithToken("test-key"), WithModel("")) + if err != errorsMissingModel { + t.Errorf("expected errorsMissingModel, got %v", err) + } +} + +func TestNew_success(t *testing.T) { + client, err := New( + WithToken("test-key"), + WithModel("anthropic/claude-sonnet-4-6"), + WithBaseURL("http://localhost:4000/v1"), + ) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if client == nil { + t.Fatal("expected non-nil client") + } + if client.model != "anthropic/claude-sonnet-4-6" { + t.Errorf("expected model %q, got %q", "anthropic/claude-sonnet-4-6", client.model) + } +} diff --git a/provider/litellm/stream_test.go b/provider/litellm/stream_test.go new file mode 100644 index 0000000..9872770 --- /dev/null +++ b/provider/litellm/stream_test.go @@ -0,0 +1,278 @@ +package litellm + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" +) + +func TestCompletionStream(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + + chunks := []string{ + `{"id":"1","object":"chat.completion.chunk","created":1,"model":"anthropic/claude-sonnet-4-6","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}`, + `{"id":"1","object":"chat.completion.chunk","created":1,"model":"anthropic/claude-sonnet-4-6","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}`, + `{"id":"1","object":"chat.completion.chunk","created":1,"model":"anthropic/claude-sonnet-4-6","choices":[{"index":0,"delta":{"content":" from"},"finish_reason":null}]}`, + `{"id":"1","object":"chat.completion.chunk","created":1,"model":"anthropic/claude-sonnet-4-6","choices":[{"index":0,"delta":{"content":" LiteLLM"},"finish_reason":null}]}`, + `{"id":"1","object":"chat.completion.chunk","created":1,"model":"anthropic/claude-sonnet-4-6","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":15,"completion_tokens":3,"total_tokens":18}}`, + } + + for _, chunk := range chunks { + fmt.Fprintf(w, "data: %s\n\n", chunk) + } + fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer server.Close() + + client, err := New( + WithToken("test-token"), + WithModel("anthropic/claude-sonnet-4-6"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + var buf bytes.Buffer + resp, err := client.CompletionStream(context.Background(), "test prompt", &buf) + if err != nil { + t.Fatalf("CompletionStream failed: %v", err) + } + + expectedContent := "Hello from LiteLLM" + if resp.Content != expectedContent { + t.Errorf("expected content %q, got %q", expectedContent, resp.Content) + } + + if buf.String() != expectedContent { + t.Errorf("expected writer output %q, got %q", expectedContent, buf.String()) + } + + if resp.Usage.TotalTokens != 18 { + t.Errorf("expected total tokens 18, got %d", resp.Usage.TotalTokens) + } + if resp.Usage.PromptTokens != 15 { + t.Errorf("expected prompt tokens 15, got %d", resp.Usage.PromptTokens) + } + if resp.Usage.CompletionTokens != 3 { + t.Errorf("expected completion tokens 3, got %d", resp.Usage.CompletionTokens) + } +} + +func TestCompletion(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + resp := map[string]any{ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "anthropic/claude-sonnet-4-6", + "choices": []map[string]any{ + { + "index": 0, + "message": map[string]any{ + "role": "assistant", + "content": "4", + }, + "finish_reason": "stop", + }, + }, + "usage": map[string]any{ + "prompt_tokens": 10, + "completion_tokens": 1, + "total_tokens": 11, + }, + } + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + client, err := New( + WithToken("test-token"), + WithModel("anthropic/claude-sonnet-4-6"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + resp, err := client.Completion(context.Background(), "What is 2+2?") + if err != nil { + t.Fatalf("Completion failed: %v", err) + } + + if resp.Content != "4" { + t.Errorf("expected content %q, got %q", "4", resp.Content) + } + + if resp.Usage.TotalTokens != 11 { + t.Errorf("expected total tokens 11, got %d", resp.Usage.TotalTokens) + } +} + +func TestGetSummaryPrefix(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + resp := map[string]any{ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "anthropic/claude-sonnet-4-6", + "choices": []map[string]any{ + { + "index": 0, + "message": map[string]any{ + "role": "assistant", + "content": "", + "tool_calls": []map[string]any{ + { + "id": "call_1", + "type": "function", + "function": map[string]any{ + "name": "get_summary_prefix", + "arguments": `{"prefix":"feat","scope":"litellm"}`, + }, + }, + }, + }, + "finish_reason": "tool_calls", + }, + }, + "usage": map[string]any{ + "prompt_tokens": 20, + "completion_tokens": 5, + "total_tokens": 25, + }, + } + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + client, err := New( + WithToken("test-token"), + WithModel("anthropic/claude-sonnet-4-6"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + resp, err := client.GetSummaryPrefix(context.Background(), "Added LiteLLM provider") + if err != nil { + t.Fatalf("GetSummaryPrefix failed: %v", err) + } + + expected := "feat(litellm)" + if resp.Content != expected { + t.Errorf("expected content %q, got %q", expected, resp.Content) + } + + if resp.Usage.TotalTokens != 25 { + t.Errorf("expected total tokens 25, got %d", resp.Usage.TotalTokens) + } +} + +func TestGetSummaryPrefix_fallback(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Decode request to check if it has tools + var req map[string]any + json.NewDecoder(r.Body).Decode(&req) + if _, hasTools := req["tools"]; hasTools { + // Simulate function calling not supported + w.WriteHeader(http.StatusBadRequest) + json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "message": "tools not supported", + "type": "invalid_request_error", + }, + }) + return + } + + // Plain completion fallback + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "openai/o1-mini", + "choices": []map[string]any{ + { + "index": 0, + "message": map[string]any{ + "role": "assistant", + "content": "feat(litellm): add LiteLLM provider", + }, + "finish_reason": "stop", + }, + }, + "usage": map[string]any{ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + }) + })) + defer server.Close() + + client, err := New( + WithToken("test-token"), + WithModel("openai/o1-mini"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + resp, err := client.GetSummaryPrefix(context.Background(), "Added LiteLLM provider") + if err != nil { + t.Fatalf("GetSummaryPrefix fallback failed: %v", err) + } + + if resp.Content != "feat(litellm): add LiteLLM provider" { + t.Errorf("expected fallback content, got %q", resp.Content) + } +} + +func TestCompletion_noChoices(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(map[string]any{ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "test-model", + "choices": []map[string]any{}, + "usage": map[string]any{ + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + }, + }) + })) + defer server.Close() + + client, err := New( + WithToken("test-token"), + WithModel("test-model"), + WithBaseURL(server.URL+"/v1"), + ) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + _, err = client.Completion(context.Background(), "test") + if err == nil { + t.Fatal("expected error for empty choices") + } + if err.Error() != "no choices returned from LiteLLM proxy" { + t.Errorf("expected 'no choices returned from LiteLLM proxy', got %q", err.Error()) + } +}