rony-chat-bot/internal/agent/runner.go

// Package agent wraps the LLM client + RAG pipeline behind a single
// streaming call the HTTP handler can drive.
//
// We use llm.LLMClient directly (not agent.Loop) because the bot is a
// straight Q&A flow: no tools, no multi-iteration reasoning. Calling the
// underlying provider keeps the prompt under our full control so we can
// inject RAG context into the system message exactly where we want it.
package agent

import (
	"context"
	"fmt"
	"iter"
	"strings"

	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
	llmpersona "github.com/VictorVargas/rony-llm-agent/pkg/persona"

	botpersona "github.com/VictorVargas/rony-chat-bot/internal/persona"
	"github.com/VictorVargas/rony-chat-bot/internal/portfolio"
)

// Message aliases keep the HTTP handler decoupled from the upstream types.
type Message = llm.Message
type Role = llm.Role

const (
	RoleSystem    = llm.RoleSystem
	RoleUser      = llm.RoleUser
	RoleAssistant = llm.RoleAssistant
)

// Runner ties together an LLM client, the system prompt, and the RAG store.
// The persona struct is kept only for the UI greeting (its name + intro);
// the system prompt itself lives in the YAML and is passed in directly.
type Runner struct {
	client       llm.LLMClient
	persona      llmpersona.Persona
	systemPrompt string
	store        *portfolio.Store
	topK         int
	usage        *Usage
}

type Usage struct {
	InputTokens  int
	OutputTokens int
}

// New returns a Runner. The store may be nil (RAG disabled).
// systemPrompt is the full hand-written prompt from configs/...yaml.
func New(client llm.LLMClient, p llmpersona.Persona, systemPrompt string, store *portfolio.Store, topK int) *Runner {
	return &Runner{client: client, persona: p, systemPrompt: systemPrompt, store: store, topK: topK, usage: &Usage{}}
}

// LastUsage returns the token usage recorded on the most recent call.
func (r *Runner) LastUsage() Usage { return *r.usage }

// BuildMessages prepares the system prompt and turns the chat history into
// the upstream message list. The system prompt includes RAG context for the
// user's last message (if RAG is enabled).
func (r *Runner) BuildMessages(ctx context.Context, history []Message) ([]Message, string, error) {
	ragContext := ""
	if r.store != nil && len(history) > 0 {
		last := history[len(history)-1]
		if last.Role == RoleUser {
			hits, err := r.store.Search(ctx, last.Content, r.topK)
			if err != nil {
				return nil, "", fmt.Errorf("rag search: %w", err)
			}
			if len(hits) > 0 {
				ragContext = formatHits(hits)
			}
		}
	}

	// The system prompt comes from the YAML, not from the persona struct.
	// Keep the persona around only for the UI greeting.
	system := botpersona.BuildSystemPrompt(r.systemPrompt, ragContext)
	msgs := make([]Message, 0, len(history)+1)
	msgs = append(msgs, Message{Role: RoleSystem, Content: system})
	msgs = append(msgs, history...)
	return msgs, ragContext, nil
}

// Stream runs the model and yields each streamed chunk. The caller
// forwards chunk.Delta to the SSE stream; chunk.Usage on the last chunk
// carries token counts.
func (r *Runner) Stream(ctx context.Context, history []Message) iter.Seq2[llm.StreamChunk, error] {
	return func(yield func(llm.StreamChunk, error) bool) {
		msgs, _, err := r.BuildMessages(ctx, history)
		if err != nil {
			yield(llm.StreamChunk{}, err)
			return
		}
		req := llm.CompletionRequest{
			Messages: msgs,
			// No tools: this is a Q&A bot, not an agent.
		}
		for chunk, err := range r.client.Stream(ctx, req) {
			if chunk.Usage.TotalTokens > 0 || chunk.Usage.InputTokens > 0 || chunk.Usage.OutputTokens > 0 {
				r.usage = &Usage{
					InputTokens:  chunk.Usage.InputTokens,
					OutputTokens: chunk.Usage.OutputTokens,
				}
			}
			if !yield(chunk, err) {
				return
			}
			if err != nil {
				return
			}
		}
	}
}

func formatHits(hits []portfolio.SearchResult) string {
	var b strings.Builder
	for i, h := range hits {
		fmt.Fprintf(&b, "### [%d] %s — %s\n", i+1, h.ProjectID, h.Section)
		b.WriteString(h.Content)
		b.WriteString("\n\n")
	}
	return b.String()
}