rework llm

3 files changed, 330 insertions, 98 deletions

diff --git a/internal/tui/components/repl/editor.go b/internal/tui/components/repl/editor.go
index d0af8d2c5..cce966ca7 100644
--- a/internal/tui/components/repl/editor.go
+++ b/internal/tui/components/repl/editor.go
@@ -6,10 +6,11 @@ import (
 	"github.com/charmbracelet/bubbles/key"
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
-	"github.com/cloudwego/eino/schema"
 	"github.com/kujtimiihoxha/termai/internal/app"
+	"github.com/kujtimiihoxha/termai/internal/llm/agent"
 	"github.com/kujtimiihoxha/termai/internal/tui/layout"
 	"github.com/kujtimiihoxha/termai/internal/tui/styles"
+	"github.com/kujtimiihoxha/termai/internal/tui/util"
 	"github.com/kujtimiihoxha/vimtea"
 )
 
@@ -112,7 +113,7 @@ func (m *editorCmp) BorderText() map[layout.BorderPosition]string {
 		title = lipgloss.NewStyle().Foreground(styles.Primary).Render(title)
 	}
 	return map[layout.BorderPosition]string{
-		layout.TopLeftBorder: title,
+		layout.BottomLeftBorder: title,
 	}
 }
 
@@ -137,9 +138,15 @@ func (m *editorCmp) SetSize(width int, height int) {
 
 func (m *editorCmp) Send() tea.Cmd {
 	return func() tea.Msg {
+		messages, _ := m.app.Messages.List(m.sessionID)
+		if hasUnfinishedMessages(messages) {
+			return util.InfoMsg("Assistant is still working on the previous message")
+		}
+		a, _ := agent.NewCoderAgent(m.app)
+
 		content := strings.Join(m.editor.GetBuffer().Lines(), "\n")
-		m.app.Messages.Create(m.sessionID, *schema.UserMessage(content))
-		m.app.LLM.SendRequest(m.sessionID, content)
+		go a.Generate(m.sessionID, content)
+
 		return m.editor.Reset()
 	}
 }
@@ -153,10 +160,11 @@ func (m *editorCmp) BindingKeys() []key.Binding {
 }
 
 func NewEditorCmp(app *app.App) EditorCmp {
+	editor := vimtea.NewEditor(
+		vimtea.WithFileName("message.md"),
+	)
 	return &editorCmp{
-		app: app,
-		editor: vimtea.NewEditor(
-			vimtea.WithFileName("message.md"),
-		),
+		app:    app,
+		editor: editor,
 	}
 }
diff --git a/internal/tui/components/repl/messages.go b/internal/tui/components/repl/messages.go
index 9b3c5bde8..7956867ce 100644
--- a/internal/tui/components/repl/messages.go
+++ b/internal/tui/components/repl/messages.go
@@ -1,8 +1,9 @@
 package repl
 
 import (
+	"encoding/json"
 	"fmt"
-	"slices"
+	"sort"
 	"strings"
 
 	"github.com/charmbracelet/bubbles/key"
@@ -10,8 +11,8 @@ import (
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/glamour"
 	"github.com/charmbracelet/lipgloss"
-	"github.com/cloudwego/eino/schema"
 	"github.com/kujtimiihoxha/termai/internal/app"
+	"github.com/kujtimiihoxha/termai/internal/llm/agent"
 	"github.com/kujtimiihoxha/termai/internal/message"
 	"github.com/kujtimiihoxha/termai/internal/pubsub"
 	"github.com/kujtimiihoxha/termai/internal/session"
@@ -28,30 +29,50 @@ type MessagesCmp interface {
 }
 
 type messagesCmp struct {
-	app        *app.App
-	messages   []message.Message
-	session    session.Session
-	viewport   viewport.Model
-	mdRenderer *glamour.TermRenderer
-	width      int
-	height     int
-	focused    bool
-	cachedView string
+	app            *app.App
+	messages       []message.Message
+	selectedMsgIdx int // Index of the selected message
+	session        session.Session
+	viewport       viewport.Model
+	mdRenderer     *glamour.TermRenderer
+	width          int
+	height         int
+	focused        bool
+	cachedView     string
 }
 
 func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	switch msg := msg.(type) {
 	case pubsub.Event[message.Message]:
 		if msg.Type == pubsub.CreatedEvent {
-			m.messages = append(m.messages, msg.Payload)
-			m.renderView()
-			m.viewport.GotoBottom()
+			if msg.Payload.SessionID == m.session.ID {
+				m.messages = append(m.messages, msg.Payload)
+				m.renderView()
+				m.viewport.GotoBottom()
+			}
+			for _, v := range m.messages {
+				for _, c := range v.ToolCalls {
+					if c.ID == msg.Payload.SessionID {
+						m.renderView()
+						m.viewport.GotoBottom()
+					}
+				}
+			}
+		} else if msg.Type == pubsub.UpdatedEvent && msg.Payload.SessionID == m.session.ID {
+			for i, v := range m.messages {
+				if v.ID == msg.Payload.ID {
+					m.messages[i] = msg.Payload
+					m.renderView()
+					if i == len(m.messages)-1 {
+						m.viewport.GotoBottom()
+					}
+					break
+				}
+			}
 		}
 	case pubsub.Event[session.Session]:
-		if msg.Type == pubsub.UpdatedEvent {
-			if m.session.ID == msg.Payload.ID {
-				m.session = msg.Payload
-			}
+		if msg.Type == pubsub.UpdatedEvent && m.session.ID == msg.Payload.ID {
+			m.session = msg.Payload
 		}
 	case SelectedSessionMsg:
 		m.session, _ = m.app.Sessions.Get(msg.SessionID)
@@ -67,26 +88,24 @@ func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	return m, nil
 }
 
-func borderColor(role schema.RoleType) lipgloss.TerminalColor {
+func borderColor(role message.MessageRole) lipgloss.TerminalColor {
 	switch role {
-	case schema.Assistant:
+	case message.Assistant:
 		return styles.Mauve
-	case schema.User:
+	case message.User:
 		return styles.Rosewater
-	case schema.Tool:
-		return styles.Peach
 	}
 	return styles.Blue
 }
 
-func borderText(msgRole schema.RoleType, currentMessage int) map[layout.BorderPosition]string {
+func borderText(msgRole message.MessageRole, currentMessage int) map[layout.BorderPosition]string {
 	role := ""
 	icon := ""
 	switch msgRole {
-	case schema.Assistant:
+	case message.Assistant:
 		role = "Assistant"
 		icon = styles.BotIcon
-	case schema.User:
+	case message.User:
 		role = "User"
 		icon = styles.UserIcon
 	}
@@ -106,81 +125,259 @@ func borderText(msgRole schema.RoleType, currentMessage int) map[layout.BorderPo
 	}
 }
 
+func hasUnfinishedMessages(messages []message.Message) bool {
+	if len(messages) == 0 {
+		return false
+	}
+	for _, msg := range messages {
+		if !msg.Finished {
+			return true
+		}
+	}
+	lastMessage := messages[len(messages)-1]
+	return lastMessage.Role != message.Assistant
+}
+
+func (m *messagesCmp) renderMessageWithToolCall(content string, tools []message.ToolCall, futureMessages []message.Message) string {
+	allParts := []string{content}
+
+	leftPaddingValue := 4
+	connectorStyle := lipgloss.NewStyle().
+		Foreground(styles.Peach).
+		Bold(true)
+
+	toolCallStyle := lipgloss.NewStyle().
+		Border(lipgloss.RoundedBorder()).
+		BorderForeground(styles.Peach).
+		Width(m.width-leftPaddingValue-5).
+		Padding(0, 1)
+
+	toolResultStyle := lipgloss.NewStyle().
+		Border(lipgloss.RoundedBorder()).
+		BorderForeground(styles.Green).
+		Width(m.width-leftPaddingValue-5).
+		Padding(0, 1)
+
+	leftPadding := lipgloss.NewStyle().Padding(0, 0, 0, leftPaddingValue)
+
+	runningStyle := lipgloss.NewStyle().
+		Foreground(styles.Peach).
+		Bold(true)
+
+	renderTool := func(toolCall message.ToolCall) string {
+		toolHeader := lipgloss.NewStyle().
+			Bold(true).
+			Foreground(styles.Blue).
+			Render(fmt.Sprintf("%s %s", styles.ToolIcon, toolCall.Name))
+
+		var paramLines []string
+		var args map[string]interface{}
+		var paramOrder []string
+
+		json.Unmarshal([]byte(toolCall.Input), &args)
+
+		for key := range args {
+			paramOrder = append(paramOrder, key)
+		}
+		sort.Strings(paramOrder)
+
+		for _, name := range paramOrder {
+			value := args[name]
+			paramName := lipgloss.NewStyle().
+				Foreground(styles.Peach).
+				Bold(true).
+				Render(name)
+
+			truncate := m.width - leftPaddingValue*2 - 10
+			if len(fmt.Sprintf("%v", value)) > truncate {
+				value = fmt.Sprintf("%v", value)[:truncate] + lipgloss.NewStyle().Foreground(styles.Blue).Render("... (truncated)")
+			}
+			paramValue := fmt.Sprintf("%v", value)
+			paramLines = append(paramLines, fmt.Sprintf("  %s: %s", paramName, paramValue))
+		}
+
+		paramBlock := lipgloss.JoinVertical(lipgloss.Left, paramLines...)
+
+		toolContent := lipgloss.JoinVertical(lipgloss.Left, toolHeader, paramBlock)
+		return toolCallStyle.Render(toolContent)
+	}
+
+	findToolResult := func(toolCallID string, messages []message.Message) *message.ToolResult {
+		for _, msg := range messages {
+			if msg.Role == message.Tool {
+				for _, result := range msg.ToolResults {
+					if result.ToolCallID == toolCallID {
+						return &result
+					}
+				}
+			}
+		}
+		return nil
+	}
+
+	renderToolResult := func(result message.ToolResult) string {
+		resultHeader := lipgloss.NewStyle().
+			Bold(true).
+			Foreground(styles.Green).
+			Render(fmt.Sprintf("%s Result", styles.CheckIcon))
+		if result.IsError {
+			resultHeader = lipgloss.NewStyle().
+				Bold(true).
+				Foreground(styles.Red).
+				Render(fmt.Sprintf("%s Error", styles.ErrorIcon))
+		}
+
+		truncate := 200
+		content := result.Content
+		if len(content) > truncate {
+			content = content[:truncate] + lipgloss.NewStyle().Foreground(styles.Blue).Render("... (truncated)")
+		}
+
+		resultContent := lipgloss.JoinVertical(lipgloss.Left, resultHeader, content)
+		return toolResultStyle.Render(resultContent)
+	}
+
+	connector := connectorStyle.Render("└─> Tool Calls:")
+	allParts = append(allParts, connector)
+
+	for _, toolCall := range tools {
+		toolOutput := renderTool(toolCall)
+		allParts = append(allParts, leftPadding.Render(toolOutput))
+
+		result := findToolResult(toolCall.ID, futureMessages)
+		if result != nil {
+
+			resultOutput := renderToolResult(*result)
+			allParts = append(allParts, leftPadding.Render(resultOutput))
+
+		} else if toolCall.Name == agent.AgentToolName {
+
+			runningIndicator := runningStyle.Render(fmt.Sprintf("%s Running...", styles.SpinnerIcon))
+			allParts = append(allParts, leftPadding.Render(runningIndicator))
+			taskSessionMessages, _ := m.app.Messages.List(toolCall.ID)
+			for _, msg := range taskSessionMessages {
+				if msg.Role == message.Assistant {
+					for _, toolCall := range msg.ToolCalls {
+						toolHeader := lipgloss.NewStyle().
+							Bold(true).
+							Foreground(styles.Blue).
+							Render(fmt.Sprintf("%s %s", styles.ToolIcon, toolCall.Name))
+
+						var paramLines []string
+						var args map[string]interface{}
+						var paramOrder []string
+
+						json.Unmarshal([]byte(toolCall.Input), &args)
+
+						for key := range args {
+							paramOrder = append(paramOrder, key)
+						}
+						sort.Strings(paramOrder)
+
+						for _, name := range paramOrder {
+							value := args[name]
+							paramName := lipgloss.NewStyle().
+								Foreground(styles.Peach).
+								Bold(true).
+								Render(name)
+
+							truncate := 50
+							if len(fmt.Sprintf("%v", value)) > truncate {
+								value = fmt.Sprintf("%v", value)[:truncate] + lipgloss.NewStyle().Foreground(styles.Blue).Render("... (truncated)")
+							}
+							paramValue := fmt.Sprintf("%v", value)
+							paramLines = append(paramLines, fmt.Sprintf("  %s: %s", paramName, paramValue))
+						}
+
+						paramBlock := lipgloss.JoinVertical(lipgloss.Left, paramLines...)
+						toolContent := lipgloss.JoinVertical(lipgloss.Left, toolHeader, paramBlock)
+						toolOutput := toolCallStyle.BorderForeground(styles.Teal).MaxWidth(m.width - leftPaddingValue*2 - 2).Render(toolContent)
+						allParts = append(allParts, lipgloss.NewStyle().Padding(0, 0, 0, leftPaddingValue*2).Render(toolOutput))
+					}
+				}
+			}
+
+		} else {
+			runningIndicator := runningStyle.Render(fmt.Sprintf("%s Running...", styles.SpinnerIcon))
+			allParts = append(allParts, "    "+runningIndicator)
+		}
+	}
+
+	for _, msg := range futureMessages {
+		if msg.Content != "" {
+			break
+		}
+
+		for _, toolCall := range msg.ToolCalls {
+			toolOutput := renderTool(toolCall)
+			allParts = append(allParts, "    "+strings.ReplaceAll(toolOutput, "\n", "\n    "))
+
+			result := findToolResult(toolCall.ID, futureMessages)
+			if result != nil {
+				resultOutput := renderToolResult(*result)
+				allParts = append(allParts, "    "+strings.ReplaceAll(resultOutput, "\n", "\n    "))
+			} else {
+				runningIndicator := runningStyle.Render(fmt.Sprintf("%s Running...", styles.SpinnerIcon))
+				allParts = append(allParts, "    "+runningIndicator)
+			}
+		}
+	}
+
+	return lipgloss.JoinVertical(lipgloss.Left, allParts...)
+}
+
 func (m *messagesCmp) renderView() {
 	stringMessages := make([]string, 0)
 	r, _ := glamour.NewTermRenderer(
 		glamour.WithStyles(styles.CatppuccinMarkdownStyle()),
-		glamour.WithWordWrap(m.width-10),
+		glamour.WithWordWrap(m.width-20),
 		glamour.WithEmoji(),
 	)
 	textStyle := lipgloss.NewStyle().Width(m.width - 4)
 	currentMessage := 1
-	for _, msg := range m.messages {
-		if msg.MessageData.Role == schema.Tool {
-			continue
-		}
-		content := msg.MessageData.Content
-		if content != "" {
-			content, _ = r.Render(msg.MessageData.Content)
-			stringMessages = append(stringMessages, layout.Borderize(
-				textStyle.Render(content),
-				layout.BorderOptions{
-					InactiveBorder: lipgloss.DoubleBorder(),
-					ActiveBorder:   lipgloss.DoubleBorder(),
-					ActiveColor:    borderColor(msg.MessageData.Role),
-					InactiveColor:  borderColor(msg.MessageData.Role),
-					EmbeddedText:   borderText(msg.MessageData.Role, currentMessage),
-				},
-			))
-			currentMessage++
-		}
-		for _, toolCall := range msg.MessageData.ToolCalls {
-			resultInx := slices.IndexFunc(m.messages, func(m message.Message) bool {
-				return m.MessageData.ToolCallID == toolCall.ID
-			})
-			content := fmt.Sprintf("**Arguments**\n```json\n%s\n```\n", toolCall.Function.Arguments)
-			if resultInx == -1 {
-				content += "Running..."
-			} else {
-				result := m.messages[resultInx].MessageData.Content
-				if result != "" {
-					lines := strings.Split(result, "\n")
-					if len(lines) > 15 {
-						result = strings.Join(lines[:15], "\n")
-					}
-					content += fmt.Sprintf("**Result**\n```\n%s\n```\n", result)
-					if len(lines) > 15 {
-						content += fmt.Sprintf("\n\n *...%d lines are truncated* ", len(lines)-15)
-					}
-				}
+	displayedMsgCount := 0 // Track the actual displayed messages count
+
+	prevMessageWasUser := false
+	for inx, msg := range m.messages {
+		content := msg.Content
+		if content != "" || prevMessageWasUser {
+			if msg.Thinking != "" && content == "" {
+				content = msg.Thinking
+			} else if content == "" {
+				content = "..."
 			}
 			content, _ = r.Render(content)
-			stringMessages = append(stringMessages, layout.Borderize(
+
+			isSelected := inx == m.selectedMsgIdx
+
+			border := lipgloss.DoubleBorder()
+			activeColor := borderColor(msg.Role)
+
+			if isSelected {
+				activeColor = styles.Primary // Use primary color for selected message
+			}
+
+			content = layout.Borderize(
 				textStyle.Render(content),
 				layout.BorderOptions{
-					InactiveBorder: lipgloss.DoubleBorder(),
-					ActiveBorder:   lipgloss.DoubleBorder(),
-					ActiveColor:    borderColor(schema.Tool),
-					InactiveColor:  borderColor(schema.Tool),
-					EmbeddedText: map[layout.BorderPosition]string{
-						layout.TopLeftBorder: lipgloss.NewStyle().
-							Padding(0, 1).
-							Bold(true).
-							Foreground(styles.Crust).
-							Background(borderColor(schema.Tool)).
-							Render(
-								fmt.Sprintf("Tool [%s] %s ", toolCall.Function.Name, styles.ToolIcon),
-							),
-						layout.TopRightBorder: lipgloss.NewStyle().
-							Padding(0, 1).
-							Bold(true).
-							Foreground(styles.Crust).
-							Background(borderColor(schema.Tool)).
-							Render(fmt.Sprintf("#%d ", currentMessage)),
-					},
+					InactiveBorder: border,
+					ActiveBorder:   border,
+					ActiveColor:    activeColor,
+					InactiveColor:  borderColor(msg.Role),
+					EmbeddedText:   borderText(msg.Role, currentMessage),
 				},
-			))
+			)
+			if len(msg.ToolCalls) > 0 {
+				content = m.renderMessageWithToolCall(content, msg.ToolCalls, m.messages[inx+1:])
+			}
+			stringMessages = append(stringMessages, content)
 			currentMessage++
+			displayedMsgCount++
+		}
+		if msg.Role == message.User && msg.Content != "" {
+			prevMessageWasUser = true
+		} else {
+			prevMessageWasUser = false
 		}
 	}
 	m.viewport.SetContent(lipgloss.JoinVertical(lipgloss.Top, stringMessages...))
@@ -191,7 +388,9 @@ func (m *messagesCmp) View() string {
 }
 
 func (m *messagesCmp) BindingKeys() []key.Binding {
-	return layout.KeyMapToSlice(m.viewport.KeyMap)
+	keys := layout.KeyMapToSlice(m.viewport.KeyMap)
+
+	return keys
 }
 
 func (m *messagesCmp) Blur() tea.Cmd {
@@ -208,10 +407,17 @@ func (m *messagesCmp) BorderText() map[layout.BorderPosition]string {
 	if m.focused {
 		title = lipgloss.NewStyle().Foreground(styles.Primary).Render(title)
 	}
-	return map[layout.BorderPosition]string{
+	borderTest := map[layout.BorderPosition]string{
 		layout.TopLeftBorder:     title,
 		layout.BottomRightBorder: formatTokensAndCost(m.session.CompletionTokens+m.session.PromptTokens, m.session.Cost),
 	}
+	if hasUnfinishedMessages(m.messages) {
+		borderTest[layout.BottomLeftBorder] = lipgloss.NewStyle().Foreground(styles.Peach).Render("Thinking...")
+	} else {
+		borderTest[layout.BottomLeftBorder] = lipgloss.NewStyle().Foreground(styles.Text).Render("Sleeping " + styles.SleepIcon + " ")
+	}
+
+	return borderTest
 }
 
 func (m *messagesCmp) Focus() tea.Cmd {
@@ -232,6 +438,7 @@ func (m *messagesCmp) SetSize(width int, height int) {
 	m.height = height
 	m.viewport.Width = width - 2   // padding
 	m.viewport.Height = height - 2 // padding
+	m.renderView()
 }
 
 func (m *messagesCmp) Init() tea.Cmd {
diff --git a/internal/tui/components/repl/sessions.go b/internal/tui/components/repl/sessions.go
index 0f208ced9..a1302baae 100644
--- a/internal/tui/components/repl/sessions.go
+++ b/internal/tui/components/repl/sessions.go
@@ -89,7 +89,23 @@ func (i *sessionsCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		}
 		return i, i.list.SetItems(items)
 	case pubsub.Event[session.Session]:
-		if msg.Type == pubsub.UpdatedEvent {
+		if msg.Type == pubsub.CreatedEvent && msg.Payload.ParentSessionID == "" {
+			// Check if the session is already in the list
+			items := i.list.Items()
+			for _, item := range items {
+				s := item.(listItem)
+				if s.id == msg.Payload.ID {
+					return i, nil
+				}
+			}
+			// insert the new session at the top of the list
+			items = append([]list.Item{listItem{
+				id:    msg.Payload.ID,
+				title: msg.Payload.Title,
+				desc:  formatTokensAndCost(msg.Payload.PromptTokens+msg.Payload.CompletionTokens, msg.Payload.Cost),
+			}}, items...)
+			return i, i.list.SetItems(items)
+		} else if msg.Type == pubsub.UpdatedEvent {
 			// update the session in the list
 			items := i.list.Items()
 			for idx, item := range items {
@@ -229,3 +245,4 @@ func NewSessionsCmp(app *app.App) SessionsCmp {
 		focused: false,
 	}
 }
+