DeepRes_bot/internal/research/pipeline/pipeline.go

package pipeline

import (
	"context"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/emil/deepres/internal/models"
	"github.com/emil/deepres/internal/research/dedup"
	"github.com/emil/deepres/internal/research/embed"
	"github.com/emil/deepres/internal/research/search"
	"github.com/emil/deepres/internal/research/synthesize"
)

type Pipeline struct {
	searchCli  *search.SearXNGClient
	embedCli   *embed.EmbedClient
	llmCli     *synthesize.LLMClient
	reportsDir string
}

func New(searchCli *search.SearXNGClient, embedCli *embed.EmbedClient, llmCli *synthesize.LLMClient, reportsDir string) *Pipeline {
	// Создаём директорию для отчётов если нет
	if reportsDir != "" {
		_ = os.MkdirAll(reportsDir, 0755)
	}
	return &Pipeline{
		searchCli:  searchCli,
		embedCli:   embedCli,
		llmCli:     llmCli,
		reportsDir: reportsDir,
	}
}

func (p *Pipeline) RunFast(ctx context.Context, req *models.ResearchRequest) (*models.ResearchResult, error) {
	log.Printf("[fast] starting research: %s", req.Query)

	// 1. Поиск — 3-5 запросов с разными формулировками
	variations := generateQueryVariations(req.Query)
	results, err := p.searchCli.SearchWithVariations(ctx, req.Query, variations, 10)
	if err != nil {
		return nil, fmt.Errorf("search: %w", err)
	}
	if len(results) == 0 {
		return nil, fmt.Errorf("ничего не найдено по запросу")
	}
	log.Printf("[fast] found %d raw results", len(results))

	// 2. Эмбеддинги + дедупликация
	deduped, err := dedup.Deduplicate(ctx, p.embedCli, results, 0.85)
	if err != nil {
		return nil, fmt.Errorf("dedup: %w", err)
	}
	log.Printf("[fast] after dedup: %d unique sources", len(deduped.Items))

	if len(deduped.RawTexts) == 0 {
		return nil, fmt.Errorf("не удалось обработать источники")
	}

	// 3. LLM-синтез выжимки
	prompt := synthesize.BuildFastResearchPrompt(req.Query, deduped.RawTexts)
	summary, err := p.llmCli.Chat(ctx, prompt)
	if err != nil {
		return nil, fmt.Errorf("synthesize: %w", err)
	}

	// 4. Формируем источники
	var sources []models.Source
	for _, item := range deduped.Items {
		sources = append(sources, models.Source{Title: item.Title, URL: item.URL})
	}
	return &models.ResearchResult{
		ResearchID: req.ResearchID,
		Status:     "done",
		Summary:    summary,
		Sources:    sources,
	}, nil
}

func (p *Pipeline) RunDeep(ctx context.Context, req *models.ResearchRequest) (*models.ResearchResult, error) {
	log.Printf("[deep] starting research: %s", req.Query)

	// 1. Поиск — 10-15 запросов
	variations := generateQueryVariations(req.Query)
	results, err := p.searchCli.SearchWithVariations(ctx, req.Query, variations, 15)
	if err != nil {
		return nil, fmt.Errorf("search: %w", err)
	}
	if len(results) == 0 {
		return nil, fmt.Errorf("ничего не найдено по запросу")
	}

	// 2. Дедупликация
	deduped, err := dedup.Deduplicate(ctx, p.embedCli, results, 0.80)
	if err != nil {
		return nil, fmt.Errorf("dedup: %w", err)
	}
	if len(deduped.RawTexts) == 0 {
		return nil, fmt.Errorf("не удалось обработать источники")
	}

	// 3. Первый проход LLM — определяем чего не хватает
	prompt1 := synthesize.BuildGapPrompt(req.Query, deduped.RawTexts)
	gapAnswer, err := p.llmCli.Chat(ctx, prompt1)
	if err != nil {
		return nil, fmt.Errorf("gap analysis: %w", err)
	}

	// 4. Второй раунд поиска по недостающим аспектам
	gapVariations := extractFollowUpQueries(gapAnswer)
	if len(gapVariations) > 0 {
		log.Printf("[deep] follow-up queries: %v", gapVariations)
		moreResults, err := p.searchCli.SearchWithVariations(ctx, req.Query, gapVariations, 5)
		if err == nil && len(moreResults) > 0 {
			moreDeduped, err := dedup.Deduplicate(ctx, p.embedCli, moreResults, 0.80)
			if err == nil {
				// Добавляем только новые URL
				seen := make(map[string]bool)
				for _, item := range deduped.Items {
					seen[item.URL] = true
				}
				for _, item := range moreDeduped.Items {
					if !seen[item.URL] {
						deduped.Items = append(deduped.Items, item)
						deduped.RawTexts = append(deduped.RawTexts, item.Title+". "+item.Content)
					}
				}
			}
		}
	}

	// 5. Финальный синтез
	prompt2 := synthesize.BuildDeepResearchPrompt(req.Query, deduped.RawTexts, "")
	finalReport, err := p.llmCli.Chat(ctx, prompt2)
	if err != nil {
		return nil, fmt.Errorf("final pass: %w", err)
	}

	// 6. Сохраняем отчёт в файл
	reportPath := ""
	if p.reportsDir != "" {
		reportPath = fmt.Sprintf("%s/%s_%s.md", p.reportsDir, time.Now().Format("20060102"), req.ResearchID[:8])
		if err := p.saveReport(reportPath, finalReport, deduped.Items, req.Query); err != nil {
			log.Printf("[deep] save report error: %v", err)
			reportPath = ""
		}
	}

	var sources []models.Source
	for _, item := range deduped.Items {
		sources = append(sources, models.Source{Title: item.Title, URL: item.URL})
	}

	return &models.ResearchResult{
		ResearchID: req.ResearchID,
		Status:     "done",
		Summary:    truncateText(finalReport, 1500),
		ReportPath: reportPath,
		Sources:    sources,
	}, nil
}

func generateQueryVariations(base string) []string {
	return []string{
		base + " обзор",
		base + " анализ",
		base + " исследование",
		base + " последние новости",
		base + " что это такое",
		base + " проблемы и решения",
		base + " будущее",
	}
}

func extractFollowUpQueries(text string) []string {
	// Простой парсер: ищем строки, которые выглядят как вопросы или темы
	var queries []string
	lines := strings.Split(text, "\n")
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		// Ищем номерованные строки или буллеты
		if strings.HasPrefix(line, "-") || strings.HasPrefix(line, "*") || (len(line) > 10 && len(line) < 200) {
			// Убираем маркеры и добавляем в запросы
			clean := strings.TrimPrefix(line, "-")
			clean = strings.TrimPrefix(clean, "*")
			clean = strings.TrimSpace(clean)
			if clean != "" && len(clean) > 5 {
				queries = append(queries, clean)
			}
		}
	}
	// Ограничиваем до 3 запросов
	if len(queries) > 3 {
		queries = queries[:3]
	}
	return queries
}

func (p *Pipeline) saveReport(path, report string, items []dedup.DedupItem, query string) error {
	content := fmt.Sprintf("# %s\n\n%s\n\n## Источники\n\n", query, report)
	for _, item := range items {
		content += fmt.Sprintf("- [%s](%s)\n", item.Title, item.URL)
	}
	return os.WriteFile(path, []byte(content), 0644)
}

func truncateText(s string, maxLen int) string {
	runes := []rune(s)
	if len(runes) <= maxLen {
		return s
	}
	return string(runes[:maxLen]) + "..."
}