root.go

  1package cmd
  2
  3import (
  4	"context"
  5	"os"
  6	"sync"
  7
  8	tea "github.com/charmbracelet/bubbletea"
  9	"github.com/kujtimiihoxha/termai/internal/app"
 10	"github.com/kujtimiihoxha/termai/internal/db"
 11	"github.com/kujtimiihoxha/termai/internal/llm/models"
 12	"github.com/kujtimiihoxha/termai/internal/tui"
 13	"github.com/spf13/cobra"
 14	"github.com/spf13/viper"
 15)
 16
 17var rootCmd = &cobra.Command{
 18	Use:   "termai",
 19	Short: "A terminal ai assistant",
 20	Long:  `A terminal ai assistant`,
 21	RunE: func(cmd *cobra.Command, args []string) error {
 22		if cmd.Flag("help").Changed {
 23			cmd.Help()
 24			return nil
 25		}
 26		debug, _ := cmd.Flags().GetBool("debug")
 27		viper.Set("debug", debug)
 28		if debug {
 29			viper.Set("log.level", "debug")
 30		}
 31
 32		conn, err := db.Connect()
 33		if err != nil {
 34			return err
 35		}
 36		ctx := context.Background()
 37
 38		app := app.New(ctx, conn)
 39		app.Logger.Info("Starting termai...")
 40		tui := tea.NewProgram(
 41			tui.New(app),
 42			tea.WithAltScreen(),
 43		)
 44		app.Logger.Info("Setting up subscriptions...")
 45		ch, unsub := setupSubscriptions(app)
 46		defer unsub()
 47
 48		go func() {
 49			for msg := range ch {
 50				tui.Send(msg)
 51			}
 52		}()
 53		if _, err := tui.Run(); err != nil {
 54			return err
 55		}
 56		return nil
 57	},
 58}
 59
 60func setupSubscriptions(app *app.App) (chan tea.Msg, func()) {
 61	ch := make(chan tea.Msg)
 62	wg := sync.WaitGroup{}
 63	ctx, cancel := context.WithCancel(app.Context)
 64
 65	{
 66		sub := app.Logger.Subscribe(ctx)
 67		wg.Add(1)
 68		go func() {
 69			for ev := range sub {
 70				ch <- ev
 71			}
 72			wg.Done()
 73		}()
 74	}
 75	{
 76		sub := app.Sessions.Subscribe(ctx)
 77		wg.Add(1)
 78		go func() {
 79			for ev := range sub {
 80				ch <- ev
 81			}
 82			wg.Done()
 83		}()
 84	}
 85	{
 86		sub := app.Messages.Subscribe(ctx)
 87		wg.Add(1)
 88		go func() {
 89			for ev := range sub {
 90				ch <- ev
 91			}
 92			wg.Done()
 93		}()
 94	}
 95	{
 96		sub := app.LLM.Subscribe(ctx)
 97		wg.Add(1)
 98		go func() {
 99			for ev := range sub {
100				ch <- ev
101			}
102			wg.Done()
103		}()
104	}
105	return ch, func() {
106		cancel()
107		wg.Wait()
108		close(ch)
109	}
110}
111
112// Execute adds all child commands to the root command and sets flags appropriately.
113// This is called by main.main(). It only needs to happen once to the rootCmd.
114func Execute() {
115	err := rootCmd.Execute()
116	if err != nil {
117		os.Exit(1)
118	}
119}
120
121func loadConfig() {
122	viper.SetConfigName(".termai")
123	viper.SetConfigType("yaml")
124	viper.AddConfigPath("$HOME")
125	viper.AddConfigPath("$XDG_CONFIG_HOME/termai")
126	viper.AddConfigPath(".")
127	viper.SetEnvPrefix("TERMAI")
128	// SET DEFAULTS
129	viper.SetDefault("log.level", "info")
130	viper.SetDefault("data.dir", ".termai")
131
132	// LLM
133	viper.SetDefault("models.big", string(models.DefaultBigModel))
134	viper.SetDefault("models.little", string(models.DefaultLittleModel))
135	viper.SetDefault("providers.openai.key", os.Getenv("OPENAI_API_KEY"))
136	viper.SetDefault("providers.anthropic.key", os.Getenv("ANTHROPIC_API_KEY"))
137	viper.SetDefault("providers.common.max_tokens", 4000)
138
139	viper.SetDefault("agents.default", "coder")
140	//
141	viper.ReadInConfig()
142
143	workdir, err := os.Getwd()
144	if err != nil {
145		panic(err)
146	}
147	viper.Set("wd", workdir)
148}
149
150func init() {
151	loadConfig()
152
153	rootCmd.Flags().BoolP("help", "h", false, "Help")
154	rootCmd.Flags().BoolP("debug", "d", false, "Help")
155}