package ai import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "strings" "time" "git.gostacks.org/iwasforcedtobehere/WiseTLP/autotlp/internal/config" "git.gostacks.org/iwasforcedtobehere/WiseTLP/autotlp/pkg/types" "git.gostacks.org/iwasforcedtobehere/WiseTLP/autotlp/pkg/utils" ) type Client struct { config *types.AIConfig httpClient *http.Client logger *utils.Logger } func NewClient(config *types.AIConfig, logger *utils.Logger) *Client { return &Client{ config: config, httpClient: &http.Client{ Timeout: 30 * time.Second, }, logger: logger.WithComponent("ai"), } } func ConfigureClient(logger *utils.Logger) (*Client, error) { fmt.Println("\n" + strings.Repeat("=", 50)) fmt.Println("AI SERVICE CONFIGURATION") fmt.Println(strings.Repeat("=", 50)) config := &types.AIConfig{} fmt.Println("\nAvailable AI providers:") fmt.Println("1. Groq (Fast inference)") fmt.Println("2. OpenRouter (Multiple models)") fmt.Println("3. Gemini (Google)") fmt.Println("4. Custom (OpenAI-compatible)") choice := utils.GetUserInput("Select provider (1-4)", "1") switch choice { case "1": config.Provider = types.AIProviderGroq config.Endpoint = "https://api.groq.com/openai/v1/chat/completions" config.Model = "openai/gpt-oss-20b" case "2": config.Provider = types.AIProviderOpenRouter config.Endpoint = "https://openrouter.ai/api/v1/chat/completions" config.Model = utils.GetUserInput("Model name", "meta-llama/llama-3.1-8b-instruct:free") case "3": config.Provider = types.AIProviderGemini config.Endpoint = "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent" config.Model = "gemini-pro" case "4": config.Provider = types.AIProviderCustom config.Endpoint = utils.GetUserInput("API endpoint", "") config.Model = utils.GetUserInput("Model name", "") default: return nil, fmt.Errorf("invalid provider selection") } fmt.Printf("\nEnter your API key for %s: ", config.Provider) var apiKey string fmt.Scanln(&apiKey) config.APIKey = strings.TrimSpace(apiKey) if config.APIKey == "" { return nil, fmt.Errorf("API key is required") } config.MaxTokens = 1500 config.Temperature = 0.3 client := NewClient(config, logger) if err := client.validateConnection(context.Background()); err != nil { return nil, fmt.Errorf("failed to validate AI connection: %w", err) } fmt.Println("✓ AI service configured successfully!") return client, nil } func (c *Client) GenerateConfig(ctx context.Context, sysInfo *types.SystemInfo, preferences *types.UserPreferences) (*types.TLPConfiguration, error) { c.logger.Info("Generating TLP configuration using AI", "provider", c.config.Provider, "model", c.config.Model) prompt := c.buildPrompt(sysInfo, preferences) response, err := c.makeRequest(ctx, prompt) if err != nil { return nil, fmt.Errorf("AI request failed: %w", err) } config, err := c.parseResponse(response, sysInfo, preferences) if err != nil { return nil, fmt.Errorf("failed to parse AI response: %w", err) } c.logger.Info("TLP configuration generated successfully", "settings_count", len(config.Settings)) return config, nil } func (c *Client) buildPrompt(sysInfo *types.SystemInfo, preferences *types.UserPreferences) string { var prompt strings.Builder prompt.WriteString("Generate TLP configuration in JSON format.\n\n") prompt.WriteString("System: ") if sysInfo.Battery != nil && sysInfo.Battery.Present { prompt.WriteString("Laptop") } else { prompt.WriteString("Desktop") } prompt.WriteString(fmt.Sprintf(", %s, %d cores\n", sysInfo.Distribution.ID, sysInfo.CPU.Cores)) prompt.WriteString(fmt.Sprintf("Profile: %s, Use: %s, Mode: %s\n\n", preferences.PowerProfile, preferences.UseCase, preferences.PerformanceMode)) prompt.WriteString("Return JSON with this structure:\n") prompt.WriteString(`{"description": "Config description", "settings": {"TLP_ENABLE": "1", "CPU_SCALING_GOVERNOR_ON_AC": "performance", "CPU_SCALING_GOVERNOR_ON_BAT": "powersave", "DISK_APM_LEVEL_ON_AC": "254", "DISK_APM_LEVEL_ON_BAT": "128", "WIFI_PWR_ON_AC": "off", "WIFI_PWR_ON_BAT": "on", "USB_AUTOSUSPEND": "1"}, "rationale": {"CPU_SCALING_GOVERNOR_ON_AC": "Max performance on AC"}, "warnings": ["Monitor temperatures"]}` + "\n\n") prompt.WriteString("Generate 8-10 TLP settings for this system.") return prompt.String() } func (c *Client) makeRequest(ctx context.Context, prompt string) (string, error) { c.logger.Debug("Making AI request", "prompt_length", len(prompt)) c.logger.Debug("Full prompt being sent", "prompt", prompt) var requestBody interface{} var endpoint string switch c.config.Provider { case types.AIProviderGemini: requestBody = map[string]interface{}{ "contents": []map[string]interface{}{ { "parts": []map[string]string{ {"text": prompt}, }, }, }, "generationConfig": map[string]interface{}{ "temperature": c.config.Temperature, "maxOutputTokens": c.config.MaxTokens, }, } endpoint = c.config.Endpoint + "?key=" + c.config.APIKey default: requestBody = map[string]interface{}{ "model": c.config.Model, "messages": []map[string]interface{}{ { "role": "user", "content": prompt, }, }, "max_tokens": c.config.MaxTokens, "temperature": c.config.Temperature, "response_format": map[string]interface{}{ "type": "json_object", }, } endpoint = c.config.Endpoint } jsonBody, err := json.Marshal(requestBody) if err != nil { return "", fmt.Errorf("failed to marshal request: %w", err) } c.logger.Debug("Request body", "json", string(jsonBody)) req, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewBuffer(jsonBody)) if err != nil { return "", fmt.Errorf("failed to create request: %w", err) } req.Header.Set("Content-Type", "application/json") if c.config.Provider != types.AIProviderGemini { req.Header.Set("Authorization", "Bearer "+c.config.APIKey) } resp, err := c.httpClient.Do(req) if err != nil { return "", fmt.Errorf("request failed: %w", err) } defer resp.Body.Close() body, err := io.ReadAll(resp.Body) if err != nil { return "", fmt.Errorf("failed to read response: %w", err) } if resp.StatusCode != http.StatusOK { return "", fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } return c.extractContent(body) } func (c *Client) extractContent(responseBody []byte) (string, error) { var response map[string]interface{} if err := json.Unmarshal(responseBody, &response); err != nil { return "", fmt.Errorf("failed to parse response JSON: %w", err) } switch c.config.Provider { case types.AIProviderGemini: if candidates, ok := response["candidates"].([]interface{}); ok && len(candidates) > 0 { if candidate, ok := candidates[0].(map[string]interface{}); ok { if content, ok := candidate["content"].(map[string]interface{}); ok { if parts, ok := content["parts"].([]interface{}); ok && len(parts) > 0 { if part, ok := parts[0].(map[string]interface{}); ok { if text, ok := part["text"].(string); ok { return text, nil } } } } } } default: if choices, ok := response["choices"].([]interface{}); ok && len(choices) > 0 { if choice, ok := choices[0].(map[string]interface{}); ok { if message, ok := choice["message"].(map[string]interface{}); ok { if content, ok := message["content"].(string); ok { return content, nil } } } } } return "", fmt.Errorf("unexpected response format") } func (c *Client) parseResponse(response string, sysInfo *types.SystemInfo, preferences *types.UserPreferences) (*types.TLPConfiguration, error) { c.logger.Debug("Raw AI response", "response", response) start := strings.Index(response, "{") end := strings.LastIndex(response, "}") + 1 if start == -1 || end == 0 { c.logger.Error("No JSON found in AI response", "response_length", len(response), "response_preview", response[:min(200, len(response))]) c.logger.Info("Generating fallback TLP configuration") return c.generateFallbackConfig(sysInfo, preferences), nil } jsonStr := response[start:end] var aiResponse struct { Description string `json:"description"` Settings map[string]string `json:"settings"` Rationale map[string]string `json:"rationale"` Warnings []string `json:"warnings"` } if err := json.Unmarshal([]byte(jsonStr), &aiResponse); err != nil { c.logger.Error("Failed to parse AI JSON response", "error", err, "json_str", jsonStr) c.logger.Info("Generating fallback TLP configuration due to JSON parse error") return c.generateFallbackConfig(sysInfo, preferences), nil } if len(aiResponse.Settings) == 0 { return nil, fmt.Errorf("AI response contains no settings") } config := &types.TLPConfiguration{ Settings: aiResponse.Settings, Description: aiResponse.Description, Rationale: aiResponse.Rationale, Warnings: aiResponse.Warnings, Generated: time.Now(), SystemInfo: sysInfo, Preferences: preferences, } if _, exists := config.Settings["TLP_ENABLE"]; !exists { config.Settings["TLP_ENABLE"] = "1" } if config.Rationale == nil { config.Rationale = make(map[string]string) } return config, nil } func (c *Client) validateConnection(ctx context.Context) error { testPrompt := "Respond with a JSON object containing 'status': 'OK' to confirm the connection is working." response, err := c.makeRequest(ctx, testPrompt) if err != nil { return fmt.Errorf("connection test failed: %w", err) } if !strings.Contains(strings.ToUpper(response), "OK") { return fmt.Errorf("unexpected response from AI service: %s", response) } return nil } func min(a, b int) int { if a < b { return a } return b } func (c *Client) generateFallbackConfig(sysInfo *types.SystemInfo, preferences *types.UserPreferences) *types.TLPConfiguration { c.logger.Info("Generating fallback configuration based on user preferences") settings := config.GetRecommendedSettings(preferences) rationale := make(map[string]string) rationale["TLP_ENABLE"] = "Enable TLP power management" rationale["CPU_SCALING_GOVERNOR_ON_AC"] = "CPU governor for AC power based on user preferences" rationale["CPU_SCALING_GOVERNOR_ON_BAT"] = "CPU governor for battery power based on user preferences" warnings := []string{ "This is a fallback configuration generated when AI service was unavailable", "Configuration is based on user preferences and common best practices", "Consider running WiseTLP again when AI service is available for optimized settings", } description := fmt.Sprintf("Fallback TLP configuration for %s use case with %s power profile", preferences.UseCase, preferences.PowerProfile) return &types.TLPConfiguration{ Settings: settings, Description: description, Rationale: rationale, Warnings: warnings, Generated: time.Now(), SystemInfo: sysInfo, Preferences: preferences, } }