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

package logseq

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"log/slog"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/firebase/genkit/go/ai"
	"github.com/firebase/genkit/go/genkit"

	"mimi/internal/bot/llm/agent"
	"mimi/internal/provider/logseq"
)

const (
	retrievePrompt = "logseq-retrieve"
	evalPrompt     = "logseq-eval"
)

type LogseqAgent struct {
	g              *genkit.Genkit
	graph          logseq.RegexGraph
	retrievePrompt ai.Prompt
	evalPrompt     ai.Prompt
}

func New(g *genkit.Genkit, graph logseq.RegexGraph) LogseqAgent {
	// Fail fast if prompt wasn't found
	retrieve := genkit.LookupPrompt(g, retrievePrompt)
	if retrieve == nil {
		log.Fatalf("no prompt named '%s' found", retrievePrompt)
	}
	eval := genkit.LookupPrompt(g, evalPrompt)
	if eval == nil {
		log.Fatalf("no prompt named '%s' found", evalPrompt)
	}

	return LogseqAgent{
		g:              g,
		graph:          graph,
		retrievePrompt: retrieve,
		evalPrompt:     eval,
	}
}

func (a LogseqAgent) GetInfo() agent.Info {
	return agent.Info{
		Name: "logseq",
		Description: `Knows all about cyber valley's / cvland / cyberia history.
		Capable of answering to questions about flora and fauna, main goals and mindsets.
		Gets info from Logseq database and know how to find any stuff in it.`,
	}
}

// sanitizeTitle ensures a title is valid for JSON serialization and LLM processing
func sanitizeTitle(title string) string {
	// Check if valid UTF-8
	if !utf8.ValidString(title) {
		// Convert to valid UTF-8 by replacing invalid sequences
		title = strings.ToValidUTF8(title, "")
	}

	// Remove control characters (except tab and newline which we'll handle separately)
	title = strings.Map(func(r rune) rune {
		if unicode.IsControl(r) && r != '\t' && r != '\n' && r != '\r' {
			return -1 // remove
		}
		return r
	}, title)

	// Replace tabs and newlines with spaces to keep titles on one line
	title = strings.ReplaceAll(title, "\t", " ")
	title = strings.ReplaceAll(title, "\n", " ")
	title = strings.ReplaceAll(title, "\r", " ")

	// Collapse multiple spaces into one
	title = strings.Join(strings.Fields(title), " ")

	// Verify it's JSON-serializable by attempting to marshal
	if _, err := json.Marshal(title); err != nil {
		slog.Warn("Title failed JSON serialization, using fallback", "original", title, "error", err)
		return "untitled"
	}

	// Truncate extremely long titles (>500 chars)
	const maxTitleLength = 500
	if len(title) > maxTitleLength {
		slog.Warn("Truncating overly long title", "original_length", len(title))
		title = title[:maxTitleLength] + "..."
	}

	return title
}

// validateTitles ensures all titles are properly JSON-serializable
func validateTitles(titles []string) []string {
	validated := make([]string, 0, len(titles))
	seen := make(map[string]bool)

	for _, title := range titles {
		sanitized := sanitizeTitle(title)

		// Skip empty titles after sanitization
		if sanitized == "" {
			slog.Warn("Skipping empty title after sanitization")
			continue
		}

		// Skip duplicates (can happen after sanitization)
		if seen[sanitized] {
			continue
		}
		seen[sanitized] = true

		validated = append(validated, sanitized)
	}

	slog.Info("Validated titles", "original_count", len(titles), "validated_count", len(validated))
	return validated
}

func (a LogseqAgent) Run(ctx context.Context, query string, msgs ...*ai.Message) (agent.Response, error) {
	var result agent.Response

	// Get all page titles from the graph
	titles := a.graph.GetAllTitles()
	slog.Info("retrieved titles from graph", "count", len(titles))

	// Validate and sanitize titles for JSON serialization
	titles = validateTitles(titles)
	if len(titles) == 0 {
		return result, fmt.Errorf("no valid titles found in knowledge graph after sanitization")
	}

	titleDocs := make([]*ai.Document, len(titles))
	for i, t := range titles {
		titleDocs[i] = ai.DocumentFromText(t, map[string]any{})
	}

	// Ask LLM to filter only relevant pages with retry mechanism
	var relevantPages map[string][]string
	var resp *ai.ModelResponse
	var err error
	const maxRetries = 3

	for attempt := 1; attempt <= maxRetries; attempt++ {
		resp, err = a.retrievePrompt.Execute(
			ctx,
			ai.WithDocs(titleDocs...),
			ai.WithMessages(msgs...),
			ai.WithInput(map[string]any{"query": query}),
		)

		if err != nil {
			slog.Error("LLM request failed", "attempt", attempt, "error", err)
			if attempt == maxRetries {
				return result, fmt.Errorf("LLM request failed after %d attempts with %w", maxRetries, err)
			}
			continue
		}

		// Log the raw response for debugging
		slog.Info("LLM retrieve response", "attempt", attempt, "text", resp.Text())

		if err := resp.Output(&relevantPages); err != nil {
			slog.Error("Failed to parse LLM output", "attempt", attempt, "response", resp.Text(), "error", err)

			if attempt == maxRetries {
				// On final attempt, return empty titles rather than failing completely
				slog.Warn("Using fallback: returning empty titles list after parse failures")
				relevantPages = map[string][]string{"titles": {}}
				break
			}

			// Add error feedback to help the model correct itself on retry
			errorMsg := fmt.Sprintf("Previous response had parsing error: %v. Please ensure you return valid JSON with format: {\"titles\": [\"title1\", \"title2\"]} or {\"titles\": []} for no matches.", err)
			msgs = append(msgs, ai.NewTextMessage(ai.RoleModel, resp.Text()))
			msgs = append(msgs, ai.NewTextMessage(ai.RoleUser, errorMsg))
			continue
		}

		// Success - break out of retry loop
		break
	}

	slog.Info("relevant pages", "titles", relevantPages["titles"])

	// Fetch relevant docs
	var docs []*ai.Document
	selectedTitles := relevantPages["titles"]

	// Sanitize selected titles to ensure they're valid
	if len(selectedTitles) > 0 {
		selectedTitles = validateTitles(selectedTitles)
	}

	// Fallback: if no titles selected, try using a default set
	if len(selectedTitles) == 0 {
		slog.Warn("No relevant pages found, using fallback to 'cyber valley' and 'cyberia'")
		// Try common fallback titles
		fallbackTitles := []string{"cyber valley", "cyberia", "cvland"}
		for _, fb := range fallbackTitles {
			// Check if title exists in our graph
			for _, title := range titles {
				if strings.EqualFold(title, fb) {
					selectedTitles = append(selectedTitles, title)
					break
				}
			}
		}
		// If still empty, just take first few pages as context
		if len(selectedTitles) == 0 && len(titles) > 0 {
			slog.Warn("Fallback titles not found, using first 3 pages from graph")
			maxFallback := 3
			if len(titles) < maxFallback {
				maxFallback = len(titles)
			}
			selectedTitles = titles[:maxFallback]
		}
	}

	// Track which pages we've already added to avoid duplicates
	addedPages := make(map[string]bool)

	for _, title := range selectedTitles {
		// First, add the selected page itself
		page, found := a.graph.GetPageByTitle(title)
		if !found {
			slog.Warn("Selected title not found in graph", "title", title)
			continue
		}

		// Read the page content
		content, err := page.Read()
		if err != nil {
			slog.Warn("Failed to read page content", "title", title, "error", err)
			continue
		}

		pageTitleLower := strings.ToLower(page.Title())
		if !addedPages[pageTitleLower] {
			docs = append(docs, ai.DocumentFromText(content, map[string]any{"title": page.Title()}))
			addedPages[pageTitleLower] = true
			slog.Info("Added main page to documents", "title", page.Title())
		}

		// Then add related pages (linked from this page)
		relatedPages := a.graph.GetRelatedPages(title, 2) // Reduced depth to 2 for performance
		slog.Info("Found related pages", "for_page", title, "count", len(relatedPages))

		for _, relPage := range relatedPages {
			relTitleLower := strings.ToLower(relPage.Title())
			if addedPages[relTitleLower] {
				continue
			}

			relContent, err := relPage.Read()
			if err != nil {
				slog.Warn("Failed to read related page content", "title", relPage.Title(), "error", err)
				continue
			}

			docs = append(docs, ai.DocumentFromText(relContent, map[string]any{"title": relPage.Title()}))
			addedPages[relTitleLower] = true
		}
	}

	if len(docs) == 0 {
		return result, fmt.Errorf("failed to fetch any relevant pages")
	}
	slog.Info("relevant documents", "length", len(docs))

	// Evaluate final prompt
	resp, err = a.evalPrompt.Execute(
		ctx,
		ai.WithDocs(docs...),
		ai.WithInput(map[string]any{"query": query}),
	)
	if err != nil {
		return result, fmt.Errorf("failed to evaluate final step with %w", err)
	}

	result = agent.NewResponse(agent.DataText{Text: resp.Text()}, resp)
	return result, nil
}