yoohoo
This commit is contained in:
347
internal/ai/client.go
Normal file
347
internal/ai/client.go
Normal file
@@ -0,0 +1,347 @@
|
||||
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,
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user