From 27e0056f2926bae8635d6677585491db2d19a442 Mon Sep 17 00:00:00 2001 From: Rodin Date: Fri, 1 May 2026 12:31:41 -0700 Subject: [PATCH] feat: add context.Context + unexport client fields REVIEW.md findings 1-4, 14: - All Gitea client methods now accept context.Context as first param - All LLM client methods now accept context.Context as first param - Use http.NewRequestWithContext for cancellation/timeout support - Main uses 3-minute timeout context for all operations - Unexport Client struct fields (baseURL, token, apiKey, etc.) - Use bytes.NewReader instead of strings.NewReader(string(...)) --- cmd/review-bot/main.go | 32 +++++++++------- gitea/client.go | 84 +++++++++++++++++++++--------------------- gitea/client_test.go | 25 +++++++------ integration_test.go | 11 ++++-- llm/client.go | 36 +++++++++--------- llm/client_test.go | 23 ++++++------ 6 files changed, 112 insertions(+), 99 deletions(-) diff --git a/cmd/review-bot/main.go b/cmd/review-bot/main.go index 8715b6e..3e7c447 100644 --- a/cmd/review-bot/main.go +++ b/cmd/review-bot/main.go @@ -1,12 +1,14 @@ package main import ( + "context" "flag" "fmt" "log" "os" "strconv" "strings" + "time" "gitea.weiker.me/rodin/review-bot/gitea" "gitea.weiker.me/rodin/review-bot/llm" @@ -64,17 +66,21 @@ func main() { llmClient.WithTemperature(*llmTemp) } + // Create a top-level context with a 3-minute timeout for all operations + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute) + defer cancel() + log.Printf("Reviewing PR #%d on %s/%s", prNumber, owner, repoName) // Step 1: Fetch PR metadata - pr, err := giteaClient.GetPullRequest(owner, repoName, prNumber) + pr, err := giteaClient.GetPullRequest(ctx, owner, repoName, prNumber) if err != nil { log.Fatalf("Failed to fetch PR: %v", err) } log.Printf("PR: %s", pr.Title) // Step 2: Fetch diff - diff, err := giteaClient.GetPullRequestDiff(owner, repoName, prNumber) + diff, err := giteaClient.GetPullRequestDiff(ctx, owner, repoName, prNumber) if err != nil { log.Fatalf("Failed to fetch diff: %v", err) } @@ -82,11 +88,11 @@ func main() { // Step 3: Fetch full file content for modified files fileContext := "" - files, err := giteaClient.GetPullRequestFiles(owner, repoName, prNumber) + files, err := giteaClient.GetPullRequestFiles(ctx, owner, repoName, prNumber) if err != nil { log.Printf("Warning: could not fetch PR files list: %v", err) } else { - fileContext = fetchFileContext(giteaClient, owner, repoName, pr.Head.Ref, files) + fileContext = fetchFileContext(ctx, giteaClient, owner, repoName, pr.Head.Ref, files) log.Printf("Fetched full context for %d files", len(files)) } @@ -94,7 +100,7 @@ func main() { ciPassed := true ciDetails := "" if pr.Head.Sha != "" { - statuses, err := giteaClient.GetCommitStatuses(owner, repoName, pr.Head.Sha) + statuses, err := giteaClient.GetCommitStatuses(ctx, owner, repoName, pr.Head.Sha) if err != nil { log.Printf("Warning: could not fetch CI status: %v", err) } else { @@ -106,7 +112,7 @@ func main() { // Step 5: Load conventions file if specified conventions := "" if *conventionsFile != "" { - content, err := giteaClient.GetFileContent(owner, repoName, *conventionsFile) + content, err := giteaClient.GetFileContent(ctx, owner, repoName, *conventionsFile) if err != nil { log.Printf("Warning: could not load conventions file %q: %v", *conventionsFile, err) } else { @@ -118,7 +124,7 @@ func main() { // Step 6: Load patterns from external repo if specified patterns := "" if *patternsRepo != "" { - patterns = fetchPatterns(giteaClient, *patternsRepo, *patternsFiles) + patterns = fetchPatterns(ctx, giteaClient, *patternsRepo, *patternsFiles) log.Printf("Loaded patterns from %s (%d bytes)", *patternsRepo, len(patterns)) } @@ -133,7 +139,7 @@ func main() { {Role: "user", Content: userPrompt}, } - response, err := llmClient.Complete(messages) + response, err := llmClient.Complete(ctx, messages) if err != nil { log.Fatalf("LLM request failed: %v", err) } @@ -158,20 +164,20 @@ func main() { } log.Printf("Posting review (event=%s)...", event) - if err := giteaClient.PostReview(owner, repoName, prNumber, event, reviewBody); err != nil { + if err := giteaClient.PostReview(ctx, owner, repoName, prNumber, event, reviewBody); err != nil { log.Fatalf("Failed to post review: %v", err) } log.Printf("Review posted successfully!") } // fetchFileContext fetches the full content of modified files from the PR branch. -func fetchFileContext(client *gitea.Client, owner, repo, ref string, files []gitea.ChangedFile) string { +func fetchFileContext(ctx context.Context, client *gitea.Client, owner, repo, ref string, files []gitea.ChangedFile) string { var sb strings.Builder for _, f := range files { if f.Status == "removed" { continue // Skip deleted files } - content, err := client.GetFileContentRef(owner, repo, f.Filename, ref) + content, err := client.GetFileContentRef(ctx, owner, repo, f.Filename, ref) if err != nil { log.Printf("Warning: could not fetch %s: %v", f.Filename, err) continue @@ -188,7 +194,7 @@ func fetchFileContext(client *gitea.Client, owner, repo, ref string, files []git // patternsRepo is comma-separated list of owner/name repos. // patternsFiles is comma-separated list of file paths or directories. // If a path ends with / or is a directory, all files within it are fetched recursively. -func fetchPatterns(client *gitea.Client, patternsRepo, patternsFiles string) string { +func fetchPatterns(ctx context.Context, client *gitea.Client, patternsRepo, patternsFiles string) string { var sb strings.Builder repos := strings.Split(patternsRepo, ",") @@ -212,7 +218,7 @@ func fetchPatterns(client *gitea.Client, patternsRepo, patternsFiles string) str continue } - files, err := client.GetAllFilesInPath(owner, repo, path) + files, err := client.GetAllFilesInPath(ctx, owner, repo, path) if err != nil { log.Printf("Warning: could not fetch %s from %s: %v", path, repoRef, err) continue diff --git a/gitea/client.go b/gitea/client.go index 71607f4..7bed748 100644 --- a/gitea/client.go +++ b/gitea/client.go @@ -1,6 +1,8 @@ package gitea import ( + "bytes" + "context" "encoding/json" "fmt" "io" @@ -10,17 +12,17 @@ import ( // Client interacts with the Gitea API. type Client struct { - BaseURL string - Token string - HTTP *http.Client + baseURL string + token string + http *http.Client } // NewClient creates a new Gitea API client. func NewClient(baseURL, token string) *Client { return &Client{ - BaseURL: strings.TrimRight(baseURL, "/"), - Token: token, - HTTP: &http.Client{}, + baseURL: strings.TrimRight(baseURL, "/"), + token: token, + http: &http.Client{}, } } @@ -49,9 +51,9 @@ type ChangedFile struct { } // GetPullRequest fetches PR metadata. -func (c *Client) GetPullRequest(owner, repo string, number int) (*PullRequest, error) { - url := fmt.Sprintf("%s/api/v1/repos/%s/%s/pulls/%d", c.BaseURL, owner, repo, number) - body, err := c.doGet(url) +func (c *Client) GetPullRequest(ctx context.Context, owner, repo string, number int) (*PullRequest, error) { + url := fmt.Sprintf("%s/api/v1/repos/%s/%s/pulls/%d", c.baseURL, owner, repo, number) + body, err := c.doGet(ctx, url) if err != nil { return nil, fmt.Errorf("fetch PR: %w", err) } @@ -63,9 +65,9 @@ func (c *Client) GetPullRequest(owner, repo string, number int) (*PullRequest, e } // GetPullRequestDiff fetches the unified diff for a PR. -func (c *Client) GetPullRequestDiff(owner, repo string, number int) (string, error) { - url := fmt.Sprintf("%s/api/v1/repos/%s/%s/pulls/%d.diff", c.BaseURL, owner, repo, number) - body, err := c.doGet(url) +func (c *Client) GetPullRequestDiff(ctx context.Context, owner, repo string, number int) (string, error) { + url := fmt.Sprintf("%s/api/v1/repos/%s/%s/pulls/%d.diff", c.baseURL, owner, repo, number) + body, err := c.doGet(ctx, url) if err != nil { return "", fmt.Errorf("fetch diff: %w", err) } @@ -73,9 +75,9 @@ func (c *Client) GetPullRequestDiff(owner, repo string, number int) (string, err } // GetPullRequestFiles fetches the list of files changed in a PR. -func (c *Client) GetPullRequestFiles(owner, repo string, number int) ([]ChangedFile, error) { - url := fmt.Sprintf("%s/api/v1/repos/%s/%s/pulls/%d/files", c.BaseURL, owner, repo, number) - body, err := c.doGet(url) +func (c *Client) GetPullRequestFiles(ctx context.Context, owner, repo string, number int) ([]ChangedFile, error) { + url := fmt.Sprintf("%s/api/v1/repos/%s/%s/pulls/%d/files", c.baseURL, owner, repo, number) + body, err := c.doGet(ctx, url) if err != nil { return nil, fmt.Errorf("fetch PR files: %w", err) } @@ -87,9 +89,9 @@ func (c *Client) GetPullRequestFiles(owner, repo string, number int) ([]ChangedF } // GetCommitStatuses fetches CI statuses for a commit SHA. -func (c *Client) GetCommitStatuses(owner, repo, sha string) ([]CommitStatus, error) { - url := fmt.Sprintf("%s/api/v1/repos/%s/%s/commits/%s/statuses", c.BaseURL, owner, repo, sha) - body, err := c.doGet(url) +func (c *Client) GetCommitStatuses(ctx context.Context, owner, repo, sha string) ([]CommitStatus, error) { + url := fmt.Sprintf("%s/api/v1/repos/%s/%s/commits/%s/statuses", c.baseURL, owner, repo, sha) + body, err := c.doGet(ctx, url) if err != nil { return nil, fmt.Errorf("fetch commit statuses: %w", err) } @@ -101,9 +103,9 @@ func (c *Client) GetCommitStatuses(owner, repo, sha string) ([]CommitStatus, err } // GetFileContent fetches a file from the default branch of a repo. -func (c *Client) GetFileContent(owner, repo, filepath string) (string, error) { - url := fmt.Sprintf("%s/api/v1/repos/%s/%s/raw/%s", c.BaseURL, owner, repo, filepath) - body, err := c.doGet(url) +func (c *Client) GetFileContent(ctx context.Context, owner, repo, filepath string) (string, error) { + url := fmt.Sprintf("%s/api/v1/repos/%s/%s/raw/%s", c.baseURL, owner, repo, filepath) + body, err := c.doGet(ctx, url) if err != nil { return "", fmt.Errorf("fetch file %s: %w", filepath, err) } @@ -111,9 +113,9 @@ func (c *Client) GetFileContent(owner, repo, filepath string) (string, error) { } // GetFileContentRef fetches a file from a specific ref (branch/tag/sha) in a repo. -func (c *Client) GetFileContentRef(owner, repo, filepath, ref string) (string, error) { - url := fmt.Sprintf("%s/api/v1/repos/%s/%s/raw/%s?ref=%s", c.BaseURL, owner, repo, filepath, ref) - body, err := c.doGet(url) +func (c *Client) GetFileContentRef(ctx context.Context, owner, repo, filepath, ref string) (string, error) { + url := fmt.Sprintf("%s/api/v1/repos/%s/%s/raw/%s?ref=%s", c.baseURL, owner, repo, filepath, ref) + body, err := c.doGet(ctx, url) if err != nil { return "", fmt.Errorf("fetch file %s@%s: %w", filepath, ref, err) } @@ -122,8 +124,8 @@ func (c *Client) GetFileContentRef(owner, repo, filepath, ref string) (string, e // PostReview submits a review to a PR. // event should be "APPROVED" or "REQUEST_CHANGES". -func (c *Client) PostReview(owner, repo string, number int, event, body string) error { - url := fmt.Sprintf("%s/api/v1/repos/%s/%s/pulls/%d/reviews", c.BaseURL, owner, repo, number) +func (c *Client) PostReview(ctx context.Context, owner, repo string, number int, event, body string) error { + url := fmt.Sprintf("%s/api/v1/repos/%s/%s/pulls/%d/reviews", c.baseURL, owner, repo, number) payload := struct { Body string `json:"body"` @@ -138,14 +140,14 @@ func (c *Client) PostReview(owner, repo string, number int, event, body string) return fmt.Errorf("marshal review payload: %w", err) } - req, err := http.NewRequest("POST", url, strings.NewReader(string(data))) + req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(data)) if err != nil { return fmt.Errorf("create review request: %w", err) } - req.Header.Set("Authorization", "token "+c.Token) + req.Header.Set("Authorization", "token "+c.token) req.Header.Set("Content-Type", "application/json") - resp, err := c.HTTP.Do(req) + resp, err := c.http.Do(req) if err != nil { return fmt.Errorf("post review: %w", err) } @@ -158,14 +160,14 @@ func (c *Client) PostReview(owner, repo string, number int, event, body string) return nil } -func (c *Client) doGet(url string) ([]byte, error) { - req, err := http.NewRequest("GET", url, nil) +func (c *Client) doGet(ctx context.Context, url string) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, "GET", url, nil) if err != nil { return nil, err } - req.Header.Set("Authorization", "token "+c.Token) + req.Header.Set("Authorization", "token "+c.token) - resp, err := c.HTTP.Do(req) + resp, err := c.http.Do(req) if err != nil { return nil, err } @@ -186,9 +188,9 @@ type ContentEntry struct { } // ListContents lists files and directories at a given path in a repo. -func (c *Client) ListContents(owner, repo, path string) ([]ContentEntry, error) { - url := fmt.Sprintf("%s/api/v1/repos/%s/%s/contents/%s", c.BaseURL, owner, repo, path) - body, err := c.doGet(url) +func (c *Client) ListContents(ctx context.Context, owner, repo, path string) ([]ContentEntry, error) { + url := fmt.Sprintf("%s/api/v1/repos/%s/%s/contents/%s", c.baseURL, owner, repo, path) + body, err := c.doGet(ctx, url) if err != nil { return nil, fmt.Errorf("list contents %s: %w", path, err) } @@ -202,14 +204,14 @@ func (c *Client) ListContents(owner, repo, path string) ([]ContentEntry, error) // GetAllFilesInPath recursively fetches all file contents under a path. // If the path is a file, returns just that file's content. // If the path is a directory, recursively fetches all files within it. -func (c *Client) GetAllFilesInPath(owner, repo, path string) (map[string]string, error) { +func (c *Client) GetAllFilesInPath(ctx context.Context, owner, repo, path string) (map[string]string, error) { results := make(map[string]string) // Try listing as directory first - entries, err := c.ListContents(owner, repo, path) + entries, err := c.ListContents(ctx, owner, repo, path) if err != nil { // Might be a file, try fetching directly - content, fileErr := c.GetFileContent(owner, repo, path) + content, fileErr := c.GetFileContent(ctx, owner, repo, path) if fileErr != nil { return nil, fmt.Errorf("path %q is neither a file nor directory: %w", path, err) } @@ -220,13 +222,13 @@ func (c *Client) GetAllFilesInPath(owner, repo, path string) (map[string]string, for _, entry := range entries { switch entry.Type { case "file": - content, err := c.GetFileContent(owner, repo, entry.Path) + content, err := c.GetFileContent(ctx, owner, repo, entry.Path) if err != nil { continue // Skip files we can't read } results[entry.Path] = content case "dir": - subResults, err := c.GetAllFilesInPath(owner, repo, entry.Path) + subResults, err := c.GetAllFilesInPath(ctx, owner, repo, entry.Path) if err != nil { continue } diff --git a/gitea/client_test.go b/gitea/client_test.go index 94ff233..0ec067c 100644 --- a/gitea/client_test.go +++ b/gitea/client_test.go @@ -1,6 +1,7 @@ package gitea import ( + "context" "encoding/json" "fmt" "net/http" @@ -28,7 +29,7 @@ func TestGetPullRequest(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-token") - got, err := client.GetPullRequest("owner", "repo", 1) + got, err := client.GetPullRequest(context.Background(), "owner", "repo", 1) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -55,7 +56,7 @@ func TestGetPullRequestDiff(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-token") - got, err := client.GetPullRequestDiff("owner", "repo", 5) + got, err := client.GetPullRequestDiff(context.Background(), "owner", "repo", 5) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -80,7 +81,7 @@ func TestGetCommitStatuses(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-token") - got, err := client.GetCommitStatuses("owner", "repo", "abc123") + got, err := client.GetCommitStatuses(context.Background(), "owner", "repo", "abc123") if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -127,7 +128,7 @@ func TestPostReview(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-token") - err := client.PostReview("owner", "repo", 3, "APPROVED", "LGTM") + err := client.PostReview(context.Background(), "owner", "repo", 3, "APPROVED", "LGTM") if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -141,7 +142,7 @@ func TestGetPullRequest_Non200(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-token") - _, err := client.GetPullRequest("owner", "repo", 999) + _, err := client.GetPullRequest(context.Background(), "owner", "repo", 999) if err == nil { t.Fatal("expected error for 404, got nil") } @@ -154,7 +155,7 @@ func TestGetPullRequest_BadJSON(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-token") - _, err := client.GetPullRequest("owner", "repo", 1) + _, err := client.GetPullRequest(context.Background(), "owner", "repo", 1) if err == nil { t.Fatal("expected error for bad JSON, got nil") } @@ -168,7 +169,7 @@ func TestPostReview_Non200(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-token") - err := client.PostReview("owner", "repo", 1, "APPROVED", "test") + err := client.PostReview(context.Background(), "owner", "repo", 1, "APPROVED", "test") if err == nil { t.Fatal("expected error for 403, got nil") } @@ -186,7 +187,7 @@ func TestGetFileContent(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-token") - got, err := client.GetFileContent("owner", "repo", "CONVENTIONS.md") + got, err := client.GetFileContent(context.Background(), "owner", "repo", "CONVENTIONS.md") if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -206,7 +207,7 @@ func TestGetPullRequestFiles(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-token") - files, err := client.GetPullRequestFiles("owner", "repo", 1) + files, err := client.GetPullRequestFiles(context.Background(), "owner", "repo", 1) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -231,7 +232,7 @@ func TestGetFileContentRef(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-token") - content, err := client.GetFileContentRef("owner", "repo", "main.go", "feature-branch") + content, err := client.GetFileContentRef(context.Background(), "owner", "repo", "main.go", "feature-branch") if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -251,7 +252,7 @@ func TestListContents(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-token") - entries, err := client.ListContents("owner", "repo", "docs") + entries, err := client.ListContents(context.Background(), "owner", "repo", "docs") if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -282,7 +283,7 @@ func TestGetAllFilesInPath_File(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-token") - files, err := client.GetAllFilesInPath("owner", "repo", "README.md") + files, err := client.GetAllFilesInPath(context.Background(), "owner", "repo", "README.md") if err != nil { t.Fatalf("unexpected error: %v", err) } diff --git a/integration_test.go b/integration_test.go index 4c3fc28..948c301 100644 --- a/integration_test.go +++ b/integration_test.go @@ -3,6 +3,7 @@ package main import ( + "context" "os" "strconv" "testing" @@ -44,7 +45,7 @@ func TestIntegration_FullReviewFlow(t *testing.T) { // Parse owner/repo owner, repoName := "", "" for i, c := range giteaRepo { - if c == / { + if c == '/' { owner = giteaRepo[:i] repoName = giteaRepo[i+1:] break @@ -54,16 +55,18 @@ func TestIntegration_FullReviewFlow(t *testing.T) { t.Fatalf("Invalid repo format %q", giteaRepo) } + ctx := context.Background() + // Step 1: Fetch PR giteaClient := gitea.NewClient(giteaURL, giteaToken) - pr, err := giteaClient.GetPullRequest(owner, repoName, prNumber) + pr, err := giteaClient.GetPullRequest(ctx, owner, repoName, prNumber) if err != nil { t.Fatalf("GetPullRequest: %v", err) } t.Logf("PR: %s (sha: %s)", pr.Title, pr.Head.Sha) // Step 2: Fetch diff - diff, err := giteaClient.GetPullRequestDiff(owner, repoName, prNumber) + diff, err := giteaClient.GetPullRequestDiff(ctx, owner, repoName, prNumber) if err != nil { t.Fatalf("GetPullRequestDiff: %v", err) } @@ -78,7 +81,7 @@ func TestIntegration_FullReviewFlow(t *testing.T) { // Step 4: Call LLM llmClient := llm.NewClient(llmBaseURL, llmAPIKey, llmModel) - response, err := llmClient.Complete([]llm.Message{ + response, err := llmClient.Complete(ctx, []llm.Message{ {Role: "system", Content: systemPrompt}, {Role: "user", Content: userPrompt}, }) diff --git a/llm/client.go b/llm/client.go index 1f3e580..9117a3d 100644 --- a/llm/client.go +++ b/llm/client.go @@ -2,6 +2,7 @@ package llm import ( "bytes" + "context" "encoding/json" "fmt" "io" @@ -11,26 +12,26 @@ import ( // Client calls an OpenAI-compatible chat completion API. type Client struct { - BaseURL string - APIKey string - Model string - Temperature float64 - HTTP *http.Client + baseURL string + apiKey string + model string + temperature float64 + http *http.Client } // NewClient creates a new LLM client. func NewClient(baseURL, apiKey, model string) *Client { return &Client{ - BaseURL: strings.TrimRight(baseURL, "/"), - APIKey: apiKey, - Model: model, - HTTP: &http.Client{}, + baseURL: strings.TrimRight(baseURL, "/"), + apiKey: apiKey, + model: model, + http: &http.Client{}, } } // WithTemperature sets the temperature for LLM requests (0 = omit, uses server default). func (c *Client) WithTemperature(t float64) *Client { - c.Temperature = t + c.temperature = t return c } @@ -57,12 +58,11 @@ type ChatResponse struct { } // Complete sends a chat completion request and returns the assistant's response content. -func (c *Client) Complete(messages []Message) (string, error) { +func (c *Client) Complete(ctx context.Context, messages []Message) (string, error) { reqBody := ChatRequest{ - Model: c.Model, - Temperature: c.Temperature, + Model: c.model, + Temperature: c.temperature, Messages: messages, - } data, err := json.Marshal(reqBody) @@ -70,15 +70,15 @@ func (c *Client) Complete(messages []Message) (string, error) { return "", fmt.Errorf("marshal request: %w", err) } - url := c.BaseURL + "/chat/completions" - req, err := http.NewRequest("POST", url, bytes.NewReader(data)) + url := c.baseURL + "/chat/completions" + req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(data)) if err != nil { return "", fmt.Errorf("create request: %w", err) } - req.Header.Set("Authorization", "Bearer "+c.APIKey) + req.Header.Set("Authorization", "Bearer "+c.apiKey) req.Header.Set("Content-Type", "application/json") - resp, err := c.HTTP.Do(req) + resp, err := c.http.Do(req) if err != nil { return "", fmt.Errorf("LLM request: %w", err) } diff --git a/llm/client_test.go b/llm/client_test.go index 2487e4a..0b7dc91 100644 --- a/llm/client_test.go +++ b/llm/client_test.go @@ -1,6 +1,7 @@ package llm import ( + "context" "encoding/json" "net/http" "net/http/httptest" @@ -51,7 +52,7 @@ func TestComplete_Success(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-key", "gpt-4") - got, err := client.Complete([]Message{{Role: "user", Content: "Hi"}}) + got, err := client.Complete(context.Background(), []Message{{Role: "user", Content: "Hi"}}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -68,7 +69,7 @@ func TestComplete_APIError(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-key", "gpt-4") - _, err := client.Complete([]Message{{Role: "user", Content: "Hi"}}) + _, err := client.Complete(context.Background(), []Message{{Role: "user", Content: "Hi"}}) if err == nil { t.Fatal("expected error for 429, got nil") } @@ -82,7 +83,7 @@ func TestComplete_NoChoices(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-key", "gpt-4") - _, err := client.Complete([]Message{{Role: "user", Content: "Hi"}}) + _, err := client.Complete(context.Background(), []Message{{Role: "user", Content: "Hi"}}) if err == nil { t.Fatal("expected error for no choices, got nil") } @@ -95,7 +96,7 @@ func TestComplete_BadJSON(t *testing.T) { defer server.Close() client := NewClient(server.URL, "test-key", "gpt-4") - _, err := client.Complete([]Message{{Role: "user", Content: "Hi"}}) + _, err := client.Complete(context.Background(), []Message{{Role: "user", Content: "Hi"}}) if err == nil { t.Fatal("expected error for bad JSON, got nil") } @@ -103,7 +104,7 @@ func TestComplete_BadJSON(t *testing.T) { func TestComplete_ServerDown(t *testing.T) { client := NewClient("http://127.0.0.1:1", "test-key", "gpt-4") - _, err := client.Complete([]Message{{Role: "user", Content: "Hi"}}) + _, err := client.Complete(context.Background(), []Message{{Role: "user", Content: "Hi"}}) if err == nil { t.Fatal("expected error for connection refused, got nil") } @@ -111,16 +112,16 @@ func TestComplete_ServerDown(t *testing.T) { func TestWithTemperature(t *testing.T) { client := NewClient("http://example.com", "key", "model") - if client.Temperature != 0 { - t.Errorf("expected initial temperature 0, got %f", client.Temperature) + if client.temperature != 0 { + t.Errorf("expected initial temperature 0, got %f", client.temperature) } result := client.WithTemperature(0.7) if result != client { t.Error("WithTemperature should return the same client for chaining") } - if client.Temperature != 0.7 { - t.Errorf("expected temperature 0.7, got %f", client.Temperature) + if client.temperature != 0.7 { + t.Errorf("expected temperature 0.7, got %f", client.temperature) } } @@ -147,7 +148,7 @@ func TestComplete_TemperatureOmittedWhenZero(t *testing.T) { defer server.Close() client := NewClient(server.URL, "key", "model") - _, err := client.Complete([]Message{{Role: "user", Content: "Hi"}}) + _, err := client.Complete(context.Background(), []Message{{Role: "user", Content: "Hi"}}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -180,7 +181,7 @@ func TestComplete_TemperatureIncludedWhenSet(t *testing.T) { defer server.Close() client := NewClient(server.URL, "key", "model").WithTemperature(0.7) - _, err := client.Complete([]Message{{Role: "user", Content: "Hi"}}) + _, err := client.Complete(context.Background(), []Message{{Role: "user", Content: "Hi"}}) if err != nil { t.Fatalf("unexpected error: %v", err) }