cyberia/research/mimi/internal/bot/llm/agent/summary/summary.go

package summary

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"log/slog"
	"strings"
	"sync"
	"time"

	"github.com/firebase/genkit/go/ai"
	"github.com/firebase/genkit/go/genkit"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/jackc/pgx/v5/pgxpool"

	"mimi/internal/bot/llm/agent"
	"mimi/internal/persist"
	"mimi/internal/provider/git"
	"mimi/internal/provider/github/db"
	"mimi/internal/provider/github/projects"
)

const (
	evalPrompt   = "summary"
	periodPrompt = "period-extractor"
)

type SummaryAgent struct {
	evalPrompt      ai.Prompt
	periodExtractor ai.Prompt
	ghOrg           string
	pgPool          *pgxpool.Pool
	logseqRepoPath  string
	projectService  *projects.Service
}

func New(ctx context.Context, g *genkit.Genkit, pgPool *pgxpool.Pool, ghOrg, logseqRepoPath string) SummaryAgent {
	// Fail fast if prompt wasn't found
	eval := genkit.LookupPrompt(g, evalPrompt)
	if eval == nil {
		log.Fatalf("no prompt named '%s' found", evalPrompt)
	}

	periodExtractor := genkit.LookupPrompt(g, periodPrompt)
	if periodExtractor == nil {
		log.Fatalf("no prompt named '%s' found", periodPrompt)
	}

	// Project service handles its own genkit/prompt initialization
	projectService := projects.New(ctx, pgPool)

	return SummaryAgent{
		pgPool:          pgPool,
		ghOrg:           ghOrg,
		evalPrompt:      eval,
		periodExtractor: periodExtractor,
		logseqRepoPath:  logseqRepoPath,
		projectService:  projectService,
	}
}

func (a SummaryAgent) GetInfo() agent.Info {
	return agent.Info{
		Name:        "summary",
		Description: `Provides overall summary across all available resources`,
	}
}

func (a SummaryAgent) Run(ctx context.Context, query string, msgs ...*ai.Message) (agent.Response, error) {
	var result agent.Response
	resp, err := a.periodExtractor.Execute(ctx, ai.WithInput(map[string]any{"query": query}))
	if err != nil {
		return result, fmt.Errorf("failed to extract period from query '%s' with %w", query, err)
	}
	period := strings.TrimSuffix(resp.Text(), "\n")
	slog.Info("generating summary", "period", period)

	since := time.Now()
	switch period {
	default:
		return result, fmt.Errorf("unexpected period '%s'", period)
	case "month":
		since = since.AddDate(0, 0, -30)
	case "week":
		since = since.AddDate(0, 0, -7)
	case "day":
		since = since.AddDate(0, 0, -1)
	}

	docChan := make(chan *ai.Document, 3)
	errChan := make(chan error, 3)
	var wg sync.WaitGroup
	wg.Add(3)
	startT := time.Now()

	// Retrieve GitHub projects with issues
	go func() {
		defer wg.Done()

		projectsWithIssues, err := a.projectService.FetchActiveProjects(ctx, a.ghOrg, since)
		if err != nil {
			errChan <- fmt.Errorf("failed to fetch GitHub projects: %w", err)
			return
		}

		if len(projectsWithIssues) == 0 {
			slog.Warn("no active projects found")
			docChan <- ai.DocumentFromText("{}", map[string]any{"info": "GitHub projects issues (none found)"})
			return
		}

		slog.Info("fetched GitHub projects with issues", "count", len(projectsWithIssues))

		// Format for LLM
		type projectIssues struct {
			Project  string     `json:"project"`
			Category string     `json:"category"`
			Issues   []db.Issue `json:"issues"`
		}
		allIssues := make([]projectIssues, 0, len(projectsWithIssues))
		for _, pwi := range projectsWithIssues {
			allIssues = append(allIssues, projectIssues{
				Project:  pwi.Title,
				Category: pwi.Category,
				Issues:   pwi.Issues,
			})
		}

		blob, err := json.Marshal(allIssues)
		if err != nil {
			errChan <- fmt.Errorf("failed to marshal GitHub projects info: %w", err)
			return
		}
		docChan <- ai.DocumentFromText(string(blob), map[string]any{"info": "GitHub projects issues"})
	}()

	// Retrieve Telegram info
	go func() {
		defer wg.Done()
		q := persist.New(a.pgPool)
		messages, err := q.FindTelegramMessages(ctx, pgtype.Timestamptz{Time: since, Valid: true})
		if err != nil {
			errChan <- fmt.Errorf("failed to retrieve Telegram message from DB: %w", err)
			return
		}
		slog.Info("retrieved Telegram messages", "length", len(messages))
		blob, err := json.Marshal(messages)
		if err != nil {
			errChan <- fmt.Errorf("failed to marshal Telegram messages: %w", err)
			return
		}
		docChan <- ai.DocumentFromText(string(blob), map[string]any{"info": "Related telegram messages"})
	}()

	// Retrieve LogSeq diff
	go func() {
		defer wg.Done()
		diff, err := git.DiffInterval(a.logseqRepoPath, since)
		if err != nil {
			errChan <- err
			return
		}
		slog.Info("retrieved LogSeq diff", "length", len(diff))
		docChan <- ai.DocumentFromText(diff, map[string]any{"info": "LogSeq git diff"})
	}()

	wg.Wait()
	slog.Info("summary data retrieved", "elapsed", time.Since(startT))
	close(errChan)
	close(docChan)

	// Process errors
	var errs []error
	for err := range errChan {
		errs = append(errs, err)
	}
	if len(errs) > 0 {
		slog.Warn("some data retrieval errors occurred", "errors", errors.Join(errs...))
	}

	// Collect documents
	var docs []*ai.Document
	for doc := range docChan {
		docs = append(docs, doc)
	}

	// If we have no documents at all, return error
	if len(docs) == 0 {
		return result, fmt.Errorf("failed to retrieve any data for summary: %w", errors.Join(errs...))
	}

	resp, err = a.evalPrompt.Execute(ctx, ai.WithDocs(docs...), ai.WithInput(map[string]any{"period": period}))
	if err != nil {
		return result, err
	}
	slog.Info("generated summary", "text", resp.Text())
	result = agent.NewResponse(agent.DataText{Text: resp.Text()}, resp)
	return result, nil
}