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