smart-search-back/internal/ai/openai.go

package ai

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

	"git.techease.ru/Smart-search/smart-search-back/pkg/errors"
)

type OpenAIClient struct {
	apiKey string
	client *http.Client
}

type openAIRequest struct {
	Model    string          `json:"model"`
	Messages []openAIMessage `json:"messages"`
}

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

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

type tzResponse struct {
	TZResponse string `json:"tz_response"`
}

func NewOpenAIClient(apiKey string) *OpenAIClient {
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
	}

	client := &http.Client{
		Transport: transport,
	}

	return &OpenAIClient{
		apiKey: apiKey,
		client: client,
	}
}

func (c *OpenAIClient) isMockMode() bool {
	return c.apiKey == ""
}

func (c *OpenAIClient) GenerateTZ(requestTxt string) (string, error) {
	if c.isMockMode() {
		return c.generateMockTZ(requestTxt), nil
	}

	prompt := fmt.Sprintf(`Ты — эксперт по разработке технических заданий.

ЗАДАЧА:
Преобразовать описание заказа в структурированное техническое задание.

ВХОДНЫЕ ДАННЫЕ:
Описание: %s

СТРУКТУРА ТЗ:
1. Наименование и описание
2. Количество и единицы измерения
3. Технические требования
4. Сроки выполнения/поставки
6. Требования к качеству (если есть)
7. Стоимость/бюджет (если указан)
8. Особые условия

ПРИМЕР:

Входные данные: "Нужны офисные столы деревянные, 10 шт, белого цвета, на колесиках, 50 тысяч, на ул. Пушкина"

Ответ:
{
  "tz_response": "ТЕХНИЧЕСКОЕ ЗАДАНИЕ\n\n1. ПРЕДМЕТ\nПоставка офисных столов\n\n2. КОЛИЧЕСТВО\n10 шт.\n\n3. ТРЕБОВАНИЯ\n- Материал: дерево\n- Цвет: белый\n- На колесиках\n\n4. МЕСТО ДОСТАВКИ\nг. Москва, ул. Пушкина\n\n5. БЮДЖЕТ\n50 000 руб.\n\n6. СРОКИ\nОт согласования"
}

ФОРМАТ ОТВЕТА:
{
  "tz_response": "ТЕХНИЧЕСКОЕ ЗАДАНИЕ\n\n1. РАЗДЕЛ\nДетали...\n\n2. РАЗДЕЛ\nДетали..."
}

ПРАВИЛА:
- Ответ ТОЛЬКО в формате JSON
- Используй \n для переносов строк
- Без текста до или после JSON
- Ясная нумерация разделов
- Все параметры из описания должны быть в ТЗ
- Если параметр не указан → укажи "Не указано"`, requestTxt)

	reqBody := openAIRequest{
		Model: "gpt-4o-mini",
		Messages: []openAIMessage{
			{
				Role:    "user",
				Content: prompt,
			},
		},
	}

	jsonData, err := json.Marshal(reqBody)
	if err != nil {
		return "", errors.NewInternalError(errors.AIAPIError, "failed to marshal request", err)
	}

	req, err := http.NewRequest("POST", "https://api.openai.com/v1/chat/completions", bytes.NewBuffer(jsonData))
	if err != nil {
		return "", errors.NewInternalError(errors.AIAPIError, "failed to create request", err)
	}

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

	resp, err := c.client.Do(req)
	if err != nil {
		return "", errors.NewInternalError(errors.AIAPIError, "failed to send request", err)
	}
	defer func() { _ = resp.Body.Close() }()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", errors.NewInternalError(errors.AIAPIError, "failed to read response", err)
	}

	var aiResp openAIResponse
	if err := json.Unmarshal(body, &aiResp); err != nil {
		return "", errors.NewInternalError(errors.AIAPIError, "failed to unmarshal response", err)
	}

	if aiResp.Error != nil {
		return "", errors.NewInternalError(errors.AIAPIError, aiResp.Error.Message, nil)
	}

	if len(aiResp.Choices) == 0 {
		return "", errors.NewInternalError(errors.AIAPIError, "no choices in response", nil)
	}

	content := aiResp.Choices[0].Message.Content

	content = strings.TrimPrefix(content, "```json")
	content = strings.TrimPrefix(content, "```")
	content = strings.TrimSuffix(content, "```")
	content = strings.TrimSpace(content)

	var tzResp tzResponse
	if err := json.Unmarshal([]byte(content), &tzResp); err != nil {
		re := regexp.MustCompile(`\{[\s\S]*\}`)
		match := re.FindString(content)
		if match != "" {
			if err := json.Unmarshal([]byte(match), &tzResp); err != nil {
				return "", errors.NewInternalError(errors.AIAPIError, "failed to parse TZ response", err)
			}
		} else {
			return "", errors.NewInternalError(errors.AIAPIError, "failed to parse TZ response", err)
		}
	}

	return tzResp.TZResponse, nil
}

func (c *OpenAIClient) generateMockTZ(requestTxt string) string {
	return fmt.Sprintf(`ТЕХНИЧЕСКОЕ ЗАДАНИЕ (MOCK)

1. ПРЕДМЕТ
%s

2. КОЛИЧЕСТВО
По запросу

3. ТРЕБОВАНИЯ
- Качественное исполнение
- Соответствие стандартам
- Своевременная поставка

4. МЕСТО ДОСТАВКИ
г. Москва

5. БЮДЖЕТ
Уточняется

6. СРОКИ
От согласования

7. ОСОБЫЕ УСЛОВИЯ
Mock-данные для тестирования. API ключ OpenAI не настроен.`, requestTxt)
}