From cfdd687216799cb5b47f099f1e7cd5dd16b3bdd0 Mon Sep 17 00:00:00 2001
From: Kujtim Hoxha <kujtimii.h@gmail.com>
Date: Thu, 3 Apr 2025 15:20:15 +0200
Subject: add initial lsp support

---
 internal/llm/agent/agent.go | 49 +++++++++++++++++++++++++++++----------------
 1 file changed, 32 insertions(+), 17 deletions(-)

(limited to 'internal/llm/agent/agent.go')

diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go
index 06dbca4e8..cb123e78c 100644
--- a/internal/llm/agent/agent.go
+++ b/internal/llm/agent/agent.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"strings"
 	"sync"
 
 	"github.com/kujtimiihoxha/termai/internal/app"
@@ -33,8 +34,12 @@ func (c *agent) handleTitleGeneration(sessionID, content string) {
 		c.Context,
 		[]message.Message{
 			{
-				Role:    message.User,
-				Content: content,
+				Role: message.User,
+				Parts: []message.ContentPart{
+					message.TextContent{
+						Text: content,
+					},
+				},
 			},
 		},
 		nil,
@@ -49,6 +54,8 @@ func (c *agent) handleTitleGeneration(sessionID, content string) {
 	}
 	if response.Content != "" {
 		session.Title = response.Content
+		session.Title = strings.TrimSpace(session.Title)
+		session.Title = strings.ReplaceAll(session.Title, "\n", " ")
 		c.Sessions.Save(session)
 	}
 }
@@ -79,17 +86,18 @@ func (c *agent) processEvent(
 ) error {
 	switch event.Type {
 	case provider.EventThinkingDelta:
-		assistantMsg.Thinking += event.Thinking
+		assistantMsg.AppendReasoningContent(event.Content)
 		return c.Messages.Update(*assistantMsg)
 	case provider.EventContentDelta:
-		assistantMsg.Content += event.Content
+		assistantMsg.AppendContent(event.Content)
 		return c.Messages.Update(*assistantMsg)
 	case provider.EventError:
 		log.Println("error", event.Error)
 		return event.Error
 
 	case provider.EventComplete:
-		assistantMsg.ToolCalls = event.Response.ToolCalls
+		assistantMsg.SetToolCalls(event.Response.ToolCalls)
+		assistantMsg.AddFinish(event.Response.FinishReason)
 		err := c.Messages.Update(*assistantMsg)
 		if err != nil {
 			return err
@@ -157,18 +165,21 @@ func (c *agent) handleToolExecution(
 	ctx context.Context,
 	assistantMsg message.Message,
 ) (*message.Message, error) {
-	if len(assistantMsg.ToolCalls) == 0 {
+	if len(assistantMsg.ToolCalls()) == 0 {
 		return nil, nil
 	}
 
-	toolResults, err := c.ExecuteTools(ctx, assistantMsg.ToolCalls, c.tools)
+	toolResults, err := c.ExecuteTools(ctx, assistantMsg.ToolCalls(), c.tools)
 	if err != nil {
 		return nil, err
 	}
-
+	parts := make([]message.ContentPart, 0)
+	for _, toolResult := range toolResults {
+		parts = append(parts, toolResult)
+	}
 	msg, err := c.Messages.Create(assistantMsg.SessionID, message.CreateMessageParams{
-		Role:        message.Tool,
-		ToolResults: toolResults,
+		Role:  message.Tool,
+		Parts: parts,
 	})
 
 	return &msg, err
@@ -185,8 +196,12 @@ func (c *agent) generate(sessionID string, content string) error {
 	}
 
 	userMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
-		Role:    message.User,
-		Content: content,
+		Role: message.User,
+		Parts: []message.ContentPart{
+			message.TextContent{
+				Text: content,
+			},
+		},
 	})
 	if err != nil {
 		return err
@@ -201,8 +216,8 @@ func (c *agent) generate(sessionID string, content string) error {
 		}
 
 		assistantMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
-			Role:    message.Assistant,
-			Content: "",
+			Role:  message.Assistant,
+			Parts: []message.ContentPart{},
 		})
 		if err != nil {
 			return err
@@ -210,20 +225,20 @@ func (c *agent) generate(sessionID string, content string) error {
 		for event := range eventChan {
 			err = c.processEvent(sessionID, &assistantMsg, event)
 			if err != nil {
-				assistantMsg.Finished = true
+				assistantMsg.AddFinish("error:" + err.Error())
 				c.Messages.Update(assistantMsg)
 				return err
 			}
 		}
 
 		msg, err := c.handleToolExecution(c.Context, assistantMsg)
-		assistantMsg.Finished = true
+
 		c.Messages.Update(assistantMsg)
 		if err != nil {
 			return err
 		}
 
-		if len(assistantMsg.ToolCalls) == 0 {
+		if len(assistantMsg.ToolCalls()) == 0 {
 			break
 		}
 
-- 
cgit v1.2.3