005b8ac · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

initial working agent

005b8acKujtim Hoxha2 years ago

cmd/root.go+3−4

compare with main

unifiedsplit

@@ -109,8 +109,6 @@ func setupSubscriptions(app *app.App) (chan tea.Msg, func()) {

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() {

112func Execute() {

115 err := rootCmd.Execute()

113 err := rootCmd.Execute()

116 if err != nil {

114 if err != nil {

@@ -131,13 +129,14 @@ func loadConfig() {

131

129

132 // LLM

130 // LLM

133 viper.SetDefault("models.big", string(models.DefaultBigModel))

131 viper.SetDefault("models.big", string(models.DefaultBigModel))

134− viper.SetDefault("models.little", string(models.DefaultLittleModel))

132+ viper.SetDefault("models.small", string(models.DefaultLittleModel))

135 viper.SetDefault("providers.openai.key", os.Getenv("OPENAI_API_KEY"))

133 viper.SetDefault("providers.openai.key", os.Getenv("OPENAI_API_KEY"))

136 viper.SetDefault("providers.anthropic.key", os.Getenv("ANTHROPIC_API_KEY"))

134 viper.SetDefault("providers.anthropic.key", os.Getenv("ANTHROPIC_API_KEY"))

135+ viper.SetDefault("providers.groq.key", os.Getenv("GROQ_API_KEY"))

137 viper.SetDefault("providers.common.max_tokens", 4000)

136 viper.SetDefault("providers.common.max_tokens", 4000)

138

137

139 viper.SetDefault("agents.default", "coder")

138 viper.SetDefault("agents.default", "coder")

140− //

139+

141 viper.ReadInConfig()

140 viper.ReadInConfig()

142

141

143 workdir, err := os.Getwd()

142 workdir, err := os.Getwd()

internal/llm/agent/title.gonew file+31

@@ -0,0 +1,31 @@

1+package agent

3+import (

4+ "context"

6+ "github.com/cloudwego/eino/schema"

7+ "github.com/kujtimiihoxha/termai/internal/llm/models"

8+ "github.com/spf13/viper"

9+)

10+

11+func GenerateTitle(ctx context.Context, content string) (string, error) {

12+ model, err := models.GetModel(ctx, models.ModelID(viper.GetString("models.small")))

13+ if err != nil {

14+ return "", err

15+ }

16+ out, err := model.Generate(

17+ ctx,

18+ []*schema.Message{

19+ schema.SystemMessage(`- you will generate a short title based on the first message a user begins a conversation with

20+ - ensure it is not more than 80 characters long

21+ - the title should be a summary of the user's message

22+ - do not use quotes or colons

23+ - the entire text you return will be used as the title`),

24+ schema.UserMessage(content),

25+ },

26+ )

27+ if err != nil {

28+ return "", err

29+ }

30+ return out.Content, nil

31+}

internal/llm/llm.go+28−3

@@ -11,6 +11,7 @@ import (

11 "github.com/cloudwego/eino/schema"

12 "github.com/google/uuid"

13 "github.com/kujtimiihoxha/termai/internal/llm/agent"

14+ "github.com/kujtimiihoxha/termai/internal/llm/models"

14 "github.com/kujtimiihoxha/termai/internal/logging"

15 "github.com/kujtimiihoxha/termai/internal/logging"

15 "github.com/kujtimiihoxha/termai/internal/message"

16 "github.com/kujtimiihoxha/termai/internal/message"

16 "github.com/kujtimiihoxha/termai/internal/pubsub"

17 "github.com/kujtimiihoxha/termai/internal/pubsub"

@@ -88,7 +89,7 @@ func (s *service) handleRequest(id string, sessionID string, content string) {

88 }

89 }

90 log.Printf("Request: %s", content)

91 log.Printf("Request: %s", content)

91− agent, systemMessage, err := agent.GetAgent(s.ctx, viper.GetString("agents.default"))

92+ currentAgent, systemMessage, err := agent.GetAgent(s.ctx, viper.GetString("agents.default"))

92 if err != nil {

93 if err != nil {

93 s.Publish(AgentErrorEvent, AgentEvent{

94 s.Publish(AgentErrorEvent, AgentEvent{

94 ID: id,

95 ID: id,

@@ -110,6 +111,7 @@ func (s *service) handleRequest(id string, sessionID string, content string) {

110 for _, m := range history {

111 for _, m := range history {

111 messages = append(messages, &m.MessageData)

112 messages = append(messages, &m.MessageData)

112 }

113 }

114+

113 builder := callbacks.NewHandlerBuilder()

115 builder := callbacks.NewHandlerBuilder()

114 builder.OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {

116 builder.OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {

115 i, ok := input.(*eModel.CallbackInput)

117 i, ok := input.(*eModel.CallbackInput)

@@ -140,7 +142,7 @@ func (s *service) handleRequest(id string, sessionID string, content string) {

140 return ctx

142 return ctx

141 })

143 })

142

144

143− out, err := agent.Generate(s.ctx, messages, enioAgent.WithComposeOptions(compose.WithCallbacks(builder.Build())))

145+ out, err := currentAgent.Generate(s.ctx, messages, enioAgent.WithComposeOptions(compose.WithCallbacks(builder.Build())))

144 if err != nil {

146 if err != nil {

145 s.Publish(AgentErrorEvent, AgentEvent{

147 s.Publish(AgentErrorEvent, AgentEvent{

146 ID: id,

148 ID: id,

@@ -153,6 +155,7 @@ func (s *service) handleRequest(id string, sessionID string, content string) {

153 return

155 return

154 }

156 }

155 usage := out.ResponseMeta.Usage

157 usage := out.ResponseMeta.Usage

158+ s.messages.Create(sessionID, *out)

156 if usage != nil {

159 if usage != nil {

157 log.Printf("Prompt Tokens: %d, Completion Tokens: %d, Total Tokens: %d", usage.PromptTokens, usage.CompletionTokens, usage.TotalTokens)

160 log.Printf("Prompt Tokens: %d, Completion Tokens: %d, Total Tokens: %d", usage.PromptTokens, usage.CompletionTokens, usage.TotalTokens)

158 session, err := s.sessions.Get(sessionID)

161 session, err := s.sessions.Get(sessionID)

@@ -170,6 +173,29 @@ func (s *service) handleRequest(id string, sessionID string, content string) {

170 session.PromptTokens += int64(usage.PromptTokens)

173 session.PromptTokens += int64(usage.PromptTokens)

171 session.CompletionTokens += int64(usage.CompletionTokens)

174 session.CompletionTokens += int64(usage.CompletionTokens)

172 // TODO: calculate cost

175 // TODO: calculate cost

176+ model := models.SupportedModels[models.ModelID(viper.GetString("models.big"))]

177+ session.Cost += float64(usage.PromptTokens)*(model.CostPer1MIn/1_000_000) +

178+ float64(usage.CompletionTokens)*(model.CostPer1MOut/1_000_000)

179+ var newTitle string

180+ if len(history) == 1 {

181+ // first message generate the title

182+ newTitle, err = agent.GenerateTitle(s.ctx, content)

183+ if err != nil {

184+ s.Publish(AgentErrorEvent, AgentEvent{

185+ ID: id,

186+ Type: AgentMessageTypeError,

187+ AgentID: RootAgent,

188+ MessageID: "",

189+ SessionID: sessionID,

190+ Content: err.Error(),

191+ })

192+ return

193+ }

194+ }

195+ if newTitle != "" {

196+ session.Title = newTitle

197+ }

198+

173 _, err = s.sessions.Save(session)

199 _, err = s.sessions.Save(session)

174 if err != nil {

200 if err != nil {

175 s.Publish(AgentErrorEvent, AgentEvent{

201 s.Publish(AgentErrorEvent, AgentEvent{

@@ -183,7 +209,6 @@ func (s *service) handleRequest(id string, sessionID string, content string) {

183 return

209 return

184 }

210 }

185 }

211 }

186− s.messages.Create(sessionID, *out)

187}

212}

188

213

189func (s *service) SendRequest(sessionID string, content string) {

214func (s *service) SendRequest(sessionID string, content string) {

internal/llm/models/models.go+40−8

@@ -3,6 +3,7 @@ package models

3import (

4 "context"

5 "errors"

6+ "log"

7 "github.com/cloudwego/eino-ext/components/model/claude"

8 "github.com/cloudwego/eino-ext/components/model/claude"

8 "github.com/cloudwego/eino-ext/components/model/openai"

9 "github.com/cloudwego/eino-ext/components/model/openai"

@@ -16,10 +17,12 @@ type (

16)

17)

18type Model struct {

19type Model struct {

19− ID ModelID `json:"id"`

20+ ID ModelID `json:"id"`

20− Name string `json:"name"`

21+ Name string `json:"name"`

21− Provider ModelProvider `json:"provider"`

22+ Provider ModelProvider `json:"provider"`

22− APIModel string `json:"api_model"` // Actual value used when calling the API

23+ APIModel string `json:"api_model"`

24+ CostPer1MIn float64 `json:"cost_per_1m_in"`

25+ CostPer1MOut float64 `json:"cost_per_1m_out"`

23}

26}

25const (

28const (

@@ -52,6 +55,9 @@ const (

52 // Meta

55 // Meta

53 Llama3 ModelID = "llama-3"

56 Llama3 ModelID = "llama-3"

54 Llama270B ModelID = "llama-2-70b"

57 Llama270B ModelID = "llama-2-70b"

58+ // GROQ

59+ GroqLlama3SpecDec ModelID = "groq-llama-3-spec-dec"

60+ GroqQwen32BCoder ModelID = "qwen-2.5-coder-32b"

55)

61)

57const (

63const (

@@ -61,6 +67,7 @@ const (

61 ProviderXAI ModelProvider = "xai"

67 ProviderXAI ModelProvider = "xai"

62 ProviderDeepSeek ModelProvider = "deepseek"

68 ProviderDeepSeek ModelProvider = "deepseek"

63 ProviderMeta ModelProvider = "meta"

69 ProviderMeta ModelProvider = "meta"

70+ ProviderGroq ModelProvider = "groq"

64)

71)

66var SupportedModels = map[ModelID]Model{

73var SupportedModels = map[ModelID]Model{

@@ -72,10 +79,12 @@ var SupportedModels = map[ModelID]Model{

72 APIModel: "gpt-4o",

79 APIModel: "gpt-4o",

73 },

80 },

74 GPT4oMini: {

81 GPT4oMini: {

75− ID: GPT4oMini,

82+ ID: GPT4oMini,

76− Name: "GPT-4o Mini",

83+ Name: "GPT-4o Mini",

77− Provider: ProviderOpenAI,

84+ Provider: ProviderOpenAI,

78− APIModel: "gpt-4o-mini",

85+ APIModel: "gpt-4o-mini",

86+ CostPer1MIn: 0.150,

87+ CostPer1MOut: 0.600,

79 },

88 },

80 GPT45: {

89 GPT45: {

81 ID: GPT45,

90 ID: GPT45,

@@ -172,10 +181,25 @@ var SupportedModels = map[ModelID]Model{

172 Provider: ProviderMeta,

181 Provider: ProviderMeta,

173 APIModel: "llama-2-70b",

182 APIModel: "llama-2-70b",

174 },

183 },

184+

185+ // GROQ

186+ GroqLlama3SpecDec: {

187+ ID: GroqLlama3SpecDec,

188+ Name: "GROQ LLaMA 3 SpecDec",

189+ Provider: ProviderGroq,

190+ APIModel: "llama-3.3-70b-specdec",

191+ },

192+ GroqQwen32BCoder: {

193+ ID: GroqQwen32BCoder,

194+ Name: "GROQ Qwen 2.5 Coder 32B",

195+ Provider: ProviderGroq,

196+ APIModel: "qwen-2.5-coder-32b",

197+ },

175}

198}

176

199

177func GetModel(ctx context.Context, model ModelID) (model.ChatModel, error) {

200func GetModel(ctx context.Context, model ModelID) (model.ChatModel, error) {

178 provider := SupportedModels[model].Provider

201 provider := SupportedModels[model].Provider

202+ log.Printf("Provider: %s", provider)

179 maxTokens := viper.GetInt("providers.common.max_tokens")

203 maxTokens := viper.GetInt("providers.common.max_tokens")

180 switch provider {

204 switch provider {

181 case ProviderOpenAI:

205 case ProviderOpenAI:

@@ -191,6 +215,14 @@ func GetModel(ctx context.Context, model ModelID) (model.ChatModel, error) {

191 MaxTokens: maxTokens,

215 MaxTokens: maxTokens,

192 })

216 })

193

217

218+ case ProviderGroq:

219+ return openai.NewChatModel(ctx, &openai.ChatModelConfig{

220+ BaseURL: "https://api.groq.com/openai/v1",

221+ APIKey: viper.GetString("providers.groq.key"),

222+ Model: string(SupportedModels[model].APIModel),

223+ MaxTokens: &maxTokens,

224+ })

225+

194 }

226 }

195 return nil, errors.New("unsupported provider")

227 return nil, errors.New("unsupported provider")

196}

228}

internal/tui/components/repl/messages.go+70−5

@@ -1,22 +1,33 @@

1package repl

3import (

4+ "github.com/charmbracelet/bubbles/key"

5+ "github.com/charmbracelet/bubbles/viewport"

4 tea "github.com/charmbracelet/bubbletea"

6 tea "github.com/charmbracelet/bubbletea"

5 "github.com/charmbracelet/lipgloss"

7 "github.com/charmbracelet/lipgloss"

6 "github.com/kujtimiihoxha/termai/internal/app"

8 "github.com/kujtimiihoxha/termai/internal/app"

7 "github.com/kujtimiihoxha/termai/internal/message"

9 "github.com/kujtimiihoxha/termai/internal/message"

8 "github.com/kujtimiihoxha/termai/internal/pubsub"

10 "github.com/kujtimiihoxha/termai/internal/pubsub"

9 "github.com/kujtimiihoxha/termai/internal/session"

11 "github.com/kujtimiihoxha/termai/internal/session"

12+ "github.com/kujtimiihoxha/termai/internal/tui/layout"

10)

13)

15+type MessagesCmp interface {

16+ tea.Model

17+ layout.Focusable

18+ layout.Bordered

19+ layout.Sizeable

20+ layout.Bindings

21+}

22+

12type messagesCmp struct {

23type messagesCmp struct {

13 app *app.App

24 app *app.App

14 messages []message.Message

25 messages []message.Message

15 session session.Session

26 session session.Session

16−}

27+ viewport viewport.Model

17−

28+ width int

18−func (m *messagesCmp) Init() tea.Cmd {

29+ height int

19− return nil

30+ focused bool

20}

31}

22func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

33func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

@@ -25,6 +36,12 @@ func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

25 if msg.Type == pubsub.CreatedEvent {

36 if msg.Type == pubsub.CreatedEvent {

26 m.messages = append(m.messages, msg.Payload)

37 m.messages = append(m.messages, msg.Payload)

27 }

38 }

39+ case pubsub.Event[session.Session]:

40+ if msg.Type == pubsub.UpdatedEvent {

41+ if m.session.ID == msg.Payload.ID {

42+ m.session = msg.Payload

43+ }

44+ }

28 case SelectedSessionMsg:

45 case SelectedSessionMsg:

29 m.session, _ = m.app.Sessions.Get(msg.SessionID)

46 m.session, _ = m.app.Sessions.Get(msg.SessionID)

30 m.messages, _ = m.app.Messages.List(m.session.ID)

47 m.messages, _ = m.app.Messages.List(m.session.ID)

@@ -40,7 +57,55 @@ func (i *messagesCmp) View() string {

40 return lipgloss.JoinVertical(lipgloss.Top, stringMessages...)

57 return lipgloss.JoinVertical(lipgloss.Top, stringMessages...)

41}

58}

60+// BindingKeys implements MessagesCmp.

61+func (m *messagesCmp) BindingKeys() []key.Binding {

62+ return []key.Binding{}

63+}

64+

65+// Blur implements MessagesCmp.

66+func (m *messagesCmp) Blur() tea.Cmd {

67+ m.focused = false

68+ return nil

69+}

70+

71+// BorderText implements MessagesCmp.

72+func (m *messagesCmp) BorderText() map[layout.BorderPosition]string {

73+ title := m.session.Title

74+ if len(title) > 20 {

75+ title = title[:20] + "..."

76+ }

77+ return map[layout.BorderPosition]string{

78+ layout.TopLeftBorder: title,

79+ }

80+}

81+

82+// Focus implements MessagesCmp.

83+func (m *messagesCmp) Focus() tea.Cmd {

84+ m.focused = true

85+ return nil

86+}

87+

88+// GetSize implements MessagesCmp.

89+func (m *messagesCmp) GetSize() (int, int) {

90+ return m.width, m.height

91+}

92+

93+// IsFocused implements MessagesCmp.

94+func (m *messagesCmp) IsFocused() bool {

95+ return m.focused

96+}

97+

98+// SetSize implements MessagesCmp.

99+func (m *messagesCmp) SetSize(width int, height int) {

100+ m.width = width

101+ m.height = height

102+}

103+

104+func (m *messagesCmp) Init() tea.Cmd {

105+ return nil

106+}

107+

43−func NewMessagesCmp(app *app.App) tea.Model {

108+func NewMessagesCmp(app *app.App) MessagesCmp {

44 return &messagesCmp{

109 return &messagesCmp{

45 app: app,

110 app: app,

46 messages: []message.Message{},

111 messages: []message.Message{},

internal/tui/components/repl/sessions.go+29−2

@@ -2,6 +2,7 @@ package repl

3import (

4 "fmt"

5+ "strings"

6 "github.com/charmbracelet/bubbles/key"

7 "github.com/charmbracelet/bubbles/key"

7 "github.com/charmbracelet/bubbles/list"

8 "github.com/charmbracelet/bubbles/list"

@@ -82,7 +83,7 @@ func (i *sessionsCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

82 items[i] = listItem{

83 items[i] = listItem{

83 id: s.ID,

84 id: s.ID,

84 title: s.Title,

85 title: s.Title,

85− desc: fmt.Sprintf("Tokens: %d, Cost: %.2f", s.PromptTokens+s.CompletionTokens, s.Cost),

86+ desc: formatTokensAndCost(s.PromptTokens+s.CompletionTokens, s.Cost),

86 }

87 }

88 }

88 return i, i.list.SetItems(items)

89 return i, i.list.SetItems(items)

@@ -94,7 +95,7 @@ func (i *sessionsCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

94 s := item.(listItem)

95 s := item.(listItem)

95 if s.id == msg.Payload.ID {

96 if s.id == msg.Payload.ID {

96 s.title = msg.Payload.Title

97 s.title = msg.Payload.Title

97− s.desc = fmt.Sprintf("Tokens: %d, Cost: %.2f", msg.Payload.PromptTokens+msg.Payload.CompletionTokens, msg.Payload.Cost)

98+ s.desc = formatTokensAndCost(msg.Payload.PromptTokens+msg.Payload.CompletionTokens, msg.Payload.Cost)

98 items[idx] = s

99 items[idx] = s

99 break

100 break

100 }

101 }

@@ -169,6 +170,32 @@ func (i *sessionsCmp) BindingKeys() []key.Binding {

169 return append(layout.KeyMapToSlice(i.list.KeyMap), sessionKeyMapValue.Select)

170 return append(layout.KeyMapToSlice(i.list.KeyMap), sessionKeyMapValue.Select)

170}

171}

171

172

173+func formatTokensAndCost(tokens int64, cost float64) string {

174+ // Format tokens in human-readable format (e.g., 110K, 1.2M)

175+ var formattedTokens string

176+ switch {

177+ case tokens >= 1_000_000:

178+ formattedTokens = fmt.Sprintf("%.1fM", float64(tokens)/1_000_000)

179+ case tokens >= 1_000:

180+ formattedTokens = fmt.Sprintf("%.1fK", float64(tokens)/1_000)

181+ default:

182+ formattedTokens = fmt.Sprintf("%d", tokens)

183+ }

184+

185+ // Remove .0 suffix if present

186+ if strings.HasSuffix(formattedTokens, ".0K") {

187+ formattedTokens = strings.Replace(formattedTokens, ".0K", "K", 1)

188+ }

189+ if strings.HasSuffix(formattedTokens, ".0M") {

190+ formattedTokens = strings.Replace(formattedTokens, ".0M", "M", 1)

191+ }

192+

193+ // Format cost with $ symbol and 2 decimal places

194+ formattedCost := fmt.Sprintf("$%.2f", cost)

195+

196+ return fmt.Sprintf("Tokens: %s, Cost: %s", formattedTokens, formattedCost)

197+}

198+

172func NewSessionsCmp(app *app.App) SessionsCmp {

199func NewSessionsCmp(app *app.App) SessionsCmp {

173 listDelegate := list.NewDefaultDelegate()

200 listDelegate := list.NewDefaultDelegate()

174 defaultItemStyle := list.NewDefaultItemStyles()

201 defaultItemStyle := list.NewDefaultItemStyles()