rony-chat-bot/cmd/chat-bot/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/spf13/cobra"

	"github.com/VictorVargas/rony-chat-bot/internal/agent"
	"github.com/VictorVargas/rony-chat-bot/internal/config"
	"github.com/VictorVargas/rony-chat-bot/internal/embed"
	"github.com/VictorVargas/rony-chat-bot/internal/persona"
	"github.com/VictorVargas/rony-chat-bot/internal/portfolio"
	"github.com/VictorVargas/rony-chat-bot/internal/server"
)

const version = "0.2.0-dev"

var (
	cfgFile   string
	reindexOn bool
	askNoStrm bool
)

func main() {
	root := &cobra.Command{
		Use:   "chat-bot",
		Short: "Portfolio chatbot HTTP server (rony-llm-agent + SQLite FTS5)",
		Long:  "Rony Chat Bot — HTTP server that answers questions about your portfolio using a local LLM and SQLite FTS5 RAG.",
	}
	root.PersistentFlags().StringVar(&cfgFile, "config", "configs/portfolio-bot.yaml", "Path to YAML config")

	root.AddCommand(serveCmd())
	root.AddCommand(reindexCmd())
	root.AddCommand(askCmd())
	root.AddCommand(configCmd())
	root.AddCommand(healthCmd())
	root.AddCommand(versionCmd())

	if err := root.Execute(); err != nil {
		os.Exit(1)
	}
}

func loadConfig() (*config.Config, error) {
	return config.Load(cfgFile)
}

func setupLogging(c *config.Config) {
	level := slog.LevelInfo
	switch strings.ToLower(c.Logging.Level) {
	case "debug":
		level = slog.LevelDebug
	case "warn":
		level = slog.LevelWarn
	case "error":
		level = slog.LevelError
	}
	opts := &slog.HandlerOptions{Level: level}
	var h slog.Handler
	if strings.ToLower(c.Logging.Format) == "text" {
		h = slog.NewTextHandler(os.Stderr, opts)
	} else {
		h = slog.NewJSONHandler(os.Stderr, opts)
	}
	slog.SetDefault(slog.New(h))
}

func serveCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "serve",
		Short: "Start the HTTP server",
		RunE: func(cmd *cobra.Command, _ []string) error {
			cfg, err := loadConfig()
			if err != nil {
				return err
			}
			setupLogging(cfg)

			slog.Info("starting chat-bot",
				"version", version,
				"provider", cfg.DefaultProvider().Name,
				"rag_enabled", cfg.RAG.Enabled,
				"addr", cfg.Addr(),
			)

			if reindexOn {
				if err := runReindex(cfg); err != nil {
					slog.Error("reindex-on-start failed", "err", err)
				}
			}

			store, err := portfolio.OpenStore(cfg.RAG.DBPath)
			if err != nil {
				return fmt.Errorf("open rag store: %w", err)
			}
			defer store.Close()

			provider := cfg.DefaultProvider()
			client, err := agent.NewClient(*provider)
			if err != nil {
				return fmt.Errorf("init provider %s: %w", provider.Name, err)
			}
			slog.Info("provider ready", "name", client.Name(), "type", provider.Type, "model", provider.Model)

			p, err := persona.FromConfig(cfg)
			if err != nil {
				return err
			}
			runner := agent.New(client, p, cfg.SystemPrompt, store, cfg.RAG.TopK).
				WithCatalog(cfg.RAG.IncludeCatalog).
				WithLocalizedPrompt("es", cfg.SystemPromptES).
				WithEmbedder(newEmbedder(cfg)).
				WithCompaction(agent.CompactionConfig{
					Enabled:             cfg.Compaction.Enabled,
					ThresholdRatio:      cfg.Compaction.ThresholdRatio,
					KeepRecentTurns:     cfg.Compaction.KeepRecentTurns,
					SummarySystemPrompt: cfg.Compaction.SummarySystemPrompt,
				})

			h := server.NewHandlers(cfg, runner, store, version).WithEmbedder(newEmbedder(cfg))
			srv := server.New(cfg, h)

			ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
			defer stop()

			errCh := make(chan error, 1)
			go func() { errCh <- srv.Start() }()

			select {
			case <-ctx.Done():
				slog.Info("shutdown signal received")
				shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
				defer cancel()
				return srv.Shutdown(shutdownCtx)
			case err := <-errCh:
				return err
			}
		},
	}
	cmd.Flags().BoolVar(&reindexOn, "reindex-on-start", false, "Re-index RAG before serving")
	return cmd
}

func reindexCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "reindex",
		Short: "Rebuild the SQLite FTS5 index from data/projects/",
		RunE: func(_ *cobra.Command, _ []string) error {
			cfg, err := loadConfig()
			if err != nil {
				return err
			}
			setupLogging(cfg)
			return runReindex(cfg)
		},
	}
}

// newEmbedder returns the embeddings client, or nil when the feature is off.
// Nil is a supported value everywhere downstream: retrieval falls back to
// keyword-only rather than failing.
func newEmbedder(cfg *config.Config) portfolio.Embedder {
	if !cfg.Embeddings.Enabled || cfg.Embeddings.Endpoint == "" {
		return nil
	}
	return embed.New(embed.Config{
		BaseURL:   cfg.Embeddings.Endpoint,
		Model:     cfg.Embeddings.Model,
		BatchSize: cfg.Embeddings.BatchSize,
		TimeoutMS: cfg.Embeddings.TimeoutMS,
	})
}

func runReindex(cfg *config.Config) error {
	dur, files, chunks, err := portfolio.ReindexOnDisk(cfg.RAG.DBPath, portfolio.SourcesFor(cfg.RAG.DataPath, cfg.RAG.DocsPath), portfolio.DefaultChunkerConfig(), newEmbedder(cfg))
	if err != nil {
		return err
	}
	slog.Info("reindex complete",
		"files", files,
		"chunks", chunks,
		"duration_ms", dur.Milliseconds(),
		"db", cfg.RAG.DBPath,
	)
	fmt.Printf("Indexed %d files → %d chunks in %s (%dms)\n", files, chunks, cfg.RAG.DBPath, dur.Milliseconds())

	return nil
}

func askCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "ask <question>",
		Short: "Ask a single question (no HTTP server, useful for smoke tests)",
		Args:  cobra.MinimumNArgs(1),
		RunE: func(_ *cobra.Command, args []string) error {
			cfg, err := loadConfig()
			if err != nil {
				return err
			}
			setupLogging(cfg)
			question := strings.Join(args, " ")
			return runAsk(cfg, question, askNoStrm)
		},
	}
	cmd.Flags().BoolVar(&askNoStrm, "no-stream", false, "Disable streaming output")
	return cmd
}

func runAsk(cfg *config.Config, question string, noStream bool) error {
	provider := cfg.DefaultProvider()
	fmt.Fprintf(os.Stderr, "[%s via %s] %s\n", provider.Name, provider.Type, version)

	store, err := portfolio.OpenStore(cfg.RAG.DBPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "warning: rag store unavailable (%v); answering without RAG\n", err)
	}
	defer func() {
		if store != nil {
			_ = store.Close()
		}
	}()

	client, err := agent.NewClient(*provider)
	if err != nil {
		return err
	}
	p, err := persona.FromConfig(cfg)
	if err != nil {
		return err
	}
	runner := agent.New(client, p, cfg.SystemPrompt, store, cfg.RAG.TopK).
		WithCatalog(cfg.RAG.IncludeCatalog).
		WithLocalizedPrompt("es", cfg.SystemPromptES).
		WithEmbedder(newEmbedder(cfg)).
		WithCompaction(agent.CompactionConfig{
			Enabled:             cfg.Compaction.Enabled,
			ThresholdRatio:      cfg.Compaction.ThresholdRatio,
			KeepRecentTurns:     cfg.Compaction.KeepRecentTurns,
			SummarySystemPrompt: cfg.Compaction.SummarySystemPrompt,
		})
	history := []agent.Message{{Role: agent.RoleUser, Content: question}}

	if noStream {
		var full strings.Builder
		for chunk, err := range runner.Stream(context.Background(), history) {
			if err != nil {
				return err
			}
			full.WriteString(chunk.Delta)
		}
		fmt.Println(full.String())
		return nil
	}
	for chunk, err := range runner.Stream(context.Background(), history) {
		if err != nil {
			return err
		}
		fmt.Print(chunk.Delta)
	}
	fmt.Println()
	return nil
}

func configCmd() *cobra.Command {
	cmd := &cobra.Command{Use: "config", Short: "Config utilities"}
	cmd.AddCommand(&cobra.Command{
		Use:   "validate",
		Short: "Validate the YAML config file",
		RunE: func(_ *cobra.Command, _ []string) error {
			cfg, err := loadConfig()
			if err != nil {
				return err
			}
			fmt.Printf("OK — server=%s, provider=%s (%s), rag=%v\n",
				cfg.Addr(), cfg.DefaultProvider().Name, cfg.DefaultProvider().Type, cfg.RAG.Enabled)
			return nil
		},
	})
	return cmd
}

func healthCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "health",
		Short: "Hit /api/health on a running server",
		RunE: func(_ *cobra.Command, _ []string) error {
			cfg, err := loadConfig()
			if err != nil {
				return err
			}
			url := fmt.Sprintf("http://%s/api/health", cfg.Addr())
			resp, err := http.Get(url)
			if err != nil {
				return fmt.Errorf("connect %s: %w", url, err)
			}
			defer resp.Body.Close()
			body, _ := io.ReadAll(resp.Body)
			if resp.StatusCode != 200 {
				return fmt.Errorf("health %d: %s", resp.StatusCode, string(body))
			}
			fmt.Println(string(body))
			return nil
		},
	}
}

func versionCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "version",
		Short: "Print version",
		Run: func(_ *cobra.Command, _ []string) {
			out := map[string]string{"version": version}
			_ = json.NewEncoder(os.Stdout).Encode(out)
		},
	}
}