prompt-generator/internal/llm/client.go

package llm

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"prompt-generator/internal/models"
)

type Client struct {
	BaseURL string
	APIKey  string
	Model   string
}

func NewClient(baseURL, apiKey, model string) *Client {
	return &Client{
		BaseURL: baseURL,
		APIKey:  apiKey,
		Model:   model,
	}
}

type chatRequest struct {
	Model    string        `json:"model"`
	Messages []chatMessage `json:"messages"`
}

type chatMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type chatResponse struct {
	Choices []struct {
		Message struct {
			Content string `json:"content"`
		} `json:"message"`
	} `json:"choices"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error,omitempty"`
}

func (c *Client) GetSuggestions(currentPrompt string, claudeMD string) ([]models.Suggestion, error) {
	systemPrompt := `你是一个 Prompt 工程专家。用户正在构建一个用于 Coding Agent 的 Prompt。
请分析当前 Prompt，找出缺失的关键约束或可以改进的地方。
以 JSON 数组返回建议，每条建议包含：
- title: 建议标题
- description: 详细说明
- constraint_text: 建议注入的约束文本

只返回 JSON 数组，不要有其他内容。`

	userContent := fmt.Sprintf("当前 Prompt 内容如下：\n---\n%s\n---", currentPrompt)
	if claudeMD != "" {
		userContent += fmt.Sprintf("\n\n用户的 CLAUDE.md 内容：\n---\n%s\n---", claudeMD)
	}

	reqBody := chatRequest{
		Model: c.Model,
		Messages: []chatMessage{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: userContent},
		},
	}

	jsonData, err := json.Marshal(reqBody)
	if err != nil {
		return nil, fmt.Errorf("序列化请求失败: %w", err)
	}

	url := strings.TrimRight(c.BaseURL, "/") + "/chat/completions"
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData))
	if err != nil {
		return nil, fmt.Errorf("创建请求失败: %w", err)
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.APIKey)

	client := &http.Client{Timeout: 30 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("调用 LLM API 失败: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取响应失败: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("LLM API 返回错误 (%d): %s", resp.StatusCode, string(body))
	}

	var chatResp chatResponse
	if err := json.Unmarshal(body, &chatResp); err != nil {
		return nil, fmt.Errorf("解析响应失败: %w", err)
	}

	if chatResp.Error != nil {
		return nil, fmt.Errorf("LLM 错误: %s", chatResp.Error.Message)
	}

	if len(chatResp.Choices) == 0 {
		return nil, fmt.Errorf("LLM 返回空结果")
	}

	content := chatResp.Choices[0].Message.Content
	suggestions := parseSuggestions(content)
	if suggestions == nil {
		return nil, fmt.Errorf("AI 返回格式异常，请重试")
	}

	return suggestions, nil
}

func parseSuggestions(data string) []models.Suggestion {
	// Try to find JSON array in the response
	start := -1
	end := -1
	for i := 0; i < len(data); i++ {
		if data[i] == '[' && start == -1 {
			start = i
		}
		if data[i] == ']' {
			end = i + 1
		}
	}

	if start == -1 || end == -1 {
		return nil
	}

	var suggestions []models.Suggestion
	if err := json.Unmarshal([]byte(data[start:end]), &suggestions); err != nil {
		return nil
	}
	return suggestions
}