package grok import ( "bufio" "bytes" "context" "encoding/json" "fmt" "net/http" "os" "regexp" "strings" "time" "github.com/fatih/color" "gmgauthier.com/grokkit/internal/logger" ) // Client represents a client for interacting with the XAI API. // It holds the API key and base URL for making requests. type Client struct { APIKey string BaseURL string } // NewClient creates and returns a new Client instance. // It retrieves the API key from the XAI_API_KEY environment variable. // If the variable is not set, it prints an error and exits the program. func NewClient() *Client { key := os.Getenv("XAI_API_KEY") if key == "" { color.Red("Error: XAI_API_KEY environment variable not set") os.Exit(1) } return &Client{ APIKey: key, BaseURL: "https://api.x.ai/v1", } } // Stream sends a streaming chat completion request to the API and prints the response live to the terminal. // It uses a default temperature of 0.7. This is intended for non-TUI commands. // Returns the full response text. func (c *Client) Stream(messages []map[string]string, model string) string { return c.StreamWithTemp(messages, model, 0.7) } // StreamWithTemp sends a streaming chat completion request to the API with a specified temperature // and prints the response live to the terminal. // Returns the full response text. func (c *Client) StreamWithTemp(messages []map[string]string, model string, temperature float64) string { return c.streamInternal(messages, model, temperature, true) } // StreamSilent sends a streaming chat completion request to the API and returns the full response text // without printing it live. It uses a default temperature of 0.7. // This is intended for TUI and agent use. func (c *Client) StreamSilent(messages []map[string]string, model string) string { return c.streamInternal(messages, model, 0.7, false) } func (c *Client) streamInternal(messages []map[string]string, model string, temperature float64, printLive bool) string { startTime := time.Now() url := c.BaseURL + "/chat/completions" logger.Debug("preparing API request", "model", model, "temperature", temperature, "message_count", len(messages), "stream", true) payload := map[string]interface{}{ "model": model, "messages": messages, "temperature": temperature, "stream": true, } body, err := json.Marshal(payload) if err != nil { logger.Error("failed to marshal API request", "error", err) color.Red("Failed to marshal request: %v", err) os.Exit(1) } ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) defer cancel() req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body)) if err != nil { logger.Error("failed to create HTTP request", "error", err, "url", url) color.Red("Failed to create request: %v", err) os.Exit(1) } req.Header.Set("Authorization", "Bearer "+c.APIKey) req.Header.Set("Content-Type", "application/json") logger.Info("sending API request", "model", model, "url", url) resp, err := http.DefaultClient.Do(req) if err != nil { logger.Error("API request failed", "error", err, "model", model, "duration_ms", time.Since(startTime).Milliseconds()) color.Red("Request failed: %v", err) os.Exit(1) } defer resp.Body.Close() logger.Debug("API response received", "status", resp.Status, "status_code", resp.StatusCode, "duration_ms", time.Since(startTime).Milliseconds()) var fullReply strings.Builder scanner := bufio.NewScanner(resp.Body) chunkCount := 0 for scanner.Scan() { line := scanner.Text() if strings.HasPrefix(line, "data: ") { data := strings.TrimPrefix(line, "data: ") if data == "[DONE]" { break } var chunk map[string]any if json.Unmarshal([]byte(data), &chunk) == nil { if choices, ok := chunk["choices"].([]any); ok && len(choices) > 0 { if delta, ok := choices[0].(map[string]any)["delta"].(map[string]any); ok { if content, ok := delta["content"].(string); ok && content != "" { chunkCount++ fullReply.WriteString(content) if printLive { fmt.Print(content) } } } } } } } if printLive { fmt.Println() } responseLength := fullReply.Len() duration := time.Since(startTime) logger.Info("API request completed", "model", model, "response_length", responseLength, "chunks_received", chunkCount, "duration_ms", duration.Milliseconds(), "duration", duration.String()) return fullReply.String() } // CleanCodeResponse removes Markdown code fences from the input text and returns the pure code content. // It removes opening fences (with optional language tags) and closing fences, then trims leading and trailing whitespace. // Internal blank lines are preserved as they may be intentional in code. func CleanCodeResponse(text string) string { fence := "```" // Remove any line that starts with the fence (opening fence, possibly with language tag) text = regexp.MustCompile(`(?m)^`+regexp.QuoteMeta(fence)+`.*$`).ReplaceAllString(text, "") // Remove any line that is just the fence (closing fence) text = regexp.MustCompile(`(?m)^`+regexp.QuoteMeta(fence)+`\s*$`).ReplaceAllString(text, "") // Trim only leading and trailing whitespace. // Do NOT collapse internal blank lines — they are intentional in code. text = strings.TrimSpace(text) return text }