package grok import ( "bufio" "bytes" "encoding/json" "fmt" "net/http" "os" "regexp" "strings" "github.com/fatih/color" ) type Client struct { APIKey string BaseURL string } 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 prints live to terminal (used by non-TUI commands) func (c *Client) Stream(messages []map[string]string, model string) string { return c.streamInternal(messages, model, true) } // StreamSilent returns the full text without printing (used by TUI and agent) func (c *Client) StreamSilent(messages []map[string]string, model string) string { return c.streamInternal(messages, model, false) } func (c *Client) streamInternal(messages []map[string]string, model string, printLive bool) string { url := c.BaseURL + "/chat/completions" payload := map[string]interface{}{ "model": model, "messages": messages, "temperature": 0.7, "stream": true, } body, _ := json.Marshal(payload) req, _ := http.NewRequest("POST", url, bytes.NewReader(body)) req.Header.Set("Authorization", "Bearer "+c.APIKey) req.Header.Set("Content-Type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { color.Red("Request failed: %v", err) os.Exit(1) } defer resp.Body.Close() var fullReply strings.Builder scanner := bufio.NewScanner(resp.Body) 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 != "" { fullReply.WriteString(content) if printLive { fmt.Print(content) } } } } } } } if printLive { fmt.Println() } return fullReply.String() } // CleanCodeResponse removes markdown fences and returns pure code content 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 }