initial agent setup

2025-12-21 17:54:23 +01:00 · 2025-03-23 22:25:31 +01:00
parent 8daa6e774a
commit e7258e38ae
29 changed files with 2207 additions and 109 deletions
--- a/internal/llm/llm.go
+++ b/internal/llm/llm.go
@@ -0,0 +1,206 @@
+package llm
+
+import (
+	"context"
+	"log"
+	"sync"
+	"time"
+
+	"github.com/cloudwego/eino/callbacks"
+	"github.com/cloudwego/eino/compose"
+	"github.com/cloudwego/eino/schema"
+	"github.com/google/uuid"
+	"github.com/kujtimiihoxha/termai/internal/llm/agent"
+	"github.com/kujtimiihoxha/termai/internal/logging"
+	"github.com/kujtimiihoxha/termai/internal/message"
+	"github.com/kujtimiihoxha/termai/internal/pubsub"
+	"github.com/kujtimiihoxha/termai/internal/session"
+
+	eModel "github.com/cloudwego/eino/components/model"
+	enioAgent "github.com/cloudwego/eino/flow/agent"
+	"github.com/spf13/viper"
+)
+
+const (
+	AgentRequestoEvent pubsub.EventType = "agent_request"
+	AgentErrorEvent    pubsub.EventType = "agent_error"
+	AgentResponseEvent pubsub.EventType = "agent_response"
+)
+
+type AgentMessageType int
+
+const (
+	AgentMessageTypeNewUserMessage AgentMessageType = iota
+	AgentMessageTypeAgentResponse
+	AgentMessageTypeError
+)
+
+type agentID string
+
+const (
+	RootAgent agentID = "root"
+	TaskAgent agentID = "task"
+)
+
+type AgentEvent struct {
+	ID        string           `json:"id"`
+	Type      AgentMessageType `json:"type"`
+	AgentID   agentID          `json:"agent_id"`
+	MessageID string           `json:"message_id"`
+	SessionID string           `json:"session_id"`
+	Content   string           `json:"content"`
+}
+
+type Service interface {
+	pubsub.Suscriber[AgentEvent]
+
+	SendRequest(sessionID string, content string)
+}
+type service struct {
+	*pubsub.Broker[AgentEvent]
+	Requests       sync.Map
+	ctx            context.Context
+	activeRequests sync.Map
+	messages       message.Service
+	sessions       session.Service
+	logger         logging.Interface
+}
+
+func (s *service) handleRequest(id string, sessionID string, content string) {
+	cancel, ok := s.activeRequests.Load(id)
+	if !ok {
+		return
+	}
+	defer cancel.(context.CancelFunc)()
+	defer s.activeRequests.Delete(id)
+
+	history, err := s.messages.List(sessionID)
+	if err != nil {
+		s.Publish(AgentErrorEvent, AgentEvent{
+			ID:        id,
+			Type:      AgentMessageTypeError,
+			AgentID:   RootAgent,
+			MessageID: "",
+			SessionID: sessionID,
+			Content:   err.Error(),
+		})
+		return
+	}
+
+	log.Printf("Request: %s", content)
+	agent, systemMessage, err := agent.GetAgent(s.ctx, viper.GetString("agents.default"))
+	if err != nil {
+		s.Publish(AgentErrorEvent, AgentEvent{
+			ID:        id,
+			Type:      AgentMessageTypeError,
+			AgentID:   RootAgent,
+			MessageID: "",
+			SessionID: sessionID,
+			Content:   err.Error(),
+		})
+		return
+	}
+
+	messages := []*schema.Message{
+		{
+			Role:    schema.System,
+			Content: systemMessage,
+		},
+	}
+	for _, m := range history {
+		messages = append(messages, &m.MessageData)
+	}
+	builder := callbacks.NewHandlerBuilder()
+	builder.OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
+		i, ok := input.(*eModel.CallbackInput)
+		if info.Component == "ChatModel" && ok {
+			if len(messages) < len(i.Messages) {
+				// find new messages
+				newMessages := i.Messages[len(messages):]
+				for _, m := range newMessages {
+					_, err = s.messages.Create(sessionID, *m)
+					if err != nil {
+						s.Publish(AgentErrorEvent, AgentEvent{
+							ID:        id,
+							Type:      AgentMessageTypeError,
+							AgentID:   RootAgent,
+							MessageID: "",
+							SessionID: sessionID,
+							Content:   err.Error(),
+						})
+					}
+					messages = append(messages, m)
+				}
+			}
+		}
+
+		return ctx
+	})
+	builder.OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
+		return ctx
+	})
+
+	out, err := agent.Generate(s.ctx, messages, enioAgent.WithComposeOptions(compose.WithCallbacks(builder.Build())))
+	if err != nil {
+		s.Publish(AgentErrorEvent, AgentEvent{
+			ID:        id,
+			Type:      AgentMessageTypeError,
+			AgentID:   RootAgent,
+			MessageID: "",
+			SessionID: sessionID,
+			Content:   err.Error(),
+		})
+		return
+	}
+	usage := out.ResponseMeta.Usage
+	if usage != nil {
+		log.Printf("Prompt Tokens: %d, Completion Tokens: %d, Total Tokens: %d", usage.PromptTokens, usage.CompletionTokens, usage.TotalTokens)
+		session, err := s.sessions.Get(sessionID)
+		if err != nil {
+			s.Publish(AgentErrorEvent, AgentEvent{
+				ID:        id,
+				Type:      AgentMessageTypeError,
+				AgentID:   RootAgent,
+				MessageID: "",
+				SessionID: sessionID,
+				Content:   err.Error(),
+			})
+			return
+		}
+		session.PromptTokens += int64(usage.PromptTokens)
+		session.CompletionTokens += int64(usage.CompletionTokens)
+		// TODO: calculate cost
+		_, err = s.sessions.Save(session)
+		if err != nil {
+			s.Publish(AgentErrorEvent, AgentEvent{
+				ID:        id,
+				Type:      AgentMessageTypeError,
+				AgentID:   RootAgent,
+				MessageID: "",
+				SessionID: sessionID,
+				Content:   err.Error(),
+			})
+			return
+		}
+	}
+	s.messages.Create(sessionID, *out)
+}
+
+func (s *service) SendRequest(sessionID string, content string) {
+	id := uuid.New().String()
+
+	_, cancel := context.WithTimeout(s.ctx, 5*time.Minute)
+	s.activeRequests.Store(id, cancel)
+	log.Printf("Request: %s", content)
+	go s.handleRequest(id, sessionID, content)
+}
+
+func NewService(ctx context.Context, logger logging.Interface, sessions session.Service, messages message.Service) Service {
+	return &service{
+		Broker:   pubsub.NewBroker[AgentEvent](),
+		ctx:      ctx,
+		sessions: sessions,
+		messages: messages,
+		logger:   logger,
+	}
+}