gooseek/backend/internal/agent/researcher.go

package agent

import (
	"context"

	"github.com/gooseek/backend/internal/llm"
	"github.com/gooseek/backend/internal/search"
	"github.com/gooseek/backend/internal/session"
	"github.com/gooseek/backend/internal/types"
	"github.com/google/uuid"
)

type ResearchInput struct {
	ChatHistory      []llm.Message
	FollowUp         string
	Classification   *ClassificationResult
	Mode             Mode
	Sources          []string
	Locale           string
	DetectedLang     string
	IsArticleSummary bool
}

func research(
	ctx context.Context,
	sess *session.Session,
	llmClient llm.Client,
	searchClient *search.SearXNGClient,
	input ResearchInput,
) ([]types.Chunk, error) {
	maxIterations := 1
	switch input.Mode {
	case ModeBalanced:
		maxIterations = 3
	case ModeQuality:
		maxIterations = 10
	}

	researchBlockID := uuid.New().String()
	sess.EmitBlock(types.NewResearchBlock(researchBlockID))

	allResults := make([]types.Chunk, 0)
	seenURLs := make(map[string]bool)

	searchQuery := input.Classification.StandaloneFollowUp
	if searchQuery == "" {
		searchQuery = input.FollowUp
	}

	for i := 0; i < maxIterations; i++ {
		queries := generateSearchQueries(searchQuery)

		sess.UpdateBlock(researchBlockID, []session.Patch{
			{
				Op:   "replace",
				Path: "/data/subSteps",
				Value: []types.ResearchSubStep{
					{
						ID:        uuid.New().String(),
						Type:      "searching",
						Searching: queries,
					},
				},
			},
		})

		for _, q := range queries {
			resp, err := searchClient.Search(ctx, q, &search.SearchOptions{
				Categories: categoriesToSearch(input.Sources),
				PageNo:     1,
			})
			if err != nil {
				continue
			}

			for _, r := range resp.Results {
				if r.URL != "" && !seenURLs[r.URL] {
					seenURLs[r.URL] = true
					allResults = append(allResults, r.ToChunk())
				}
			}
		}

		if input.Mode == ModeSpeed {
			break
		}

		if len(allResults) >= 20 && input.Mode == ModeBalanced {
			break
		}

		if len(allResults) >= 50 {
			break
		}
	}

	return allResults, nil
}

func categoriesToSearch(sources []string) []string {
	if len(sources) == 0 {
		return []string{"general", "news"}
	}

	categories := make([]string, 0)
	for _, s := range sources {
		switch s {
		case "web":
			categories = append(categories, "general")
		case "discussions":
			categories = append(categories, "social media")
		case "academic":
			categories = append(categories, "science")
		case "news":
			categories = append(categories, "news")
		case "images":
			categories = append(categories, "images")
		case "videos":
			categories = append(categories, "videos")
		}
	}

	if len(categories) == 0 {
		return []string{"general"}
	}

	return categories
}