session.go

  1package session
  2
  3import (
  4	"context"
  5	"database/sql"
  6
  7	"github.com/google/uuid"
  8	"github.com/kujtimiihoxha/termai/internal/db"
  9	"github.com/kujtimiihoxha/termai/internal/pubsub"
 10)
 11
 12type Session struct {
 13	ID               string
 14	ParentSessionID  string
 15	Title            string
 16	MessageCount     int64
 17	PromptTokens     int64
 18	CompletionTokens int64
 19	Cost             float64
 20	CreatedAt        int64
 21	UpdatedAt        int64
 22}
 23
 24type Service interface {
 25	pubsub.Suscriber[Session]
 26	Create(title string) (Session, error)
 27	CreateTaskSession(toolCallID, parentSessionID, title string) (Session, error)
 28	Get(id string) (Session, error)
 29	List() ([]Session, error)
 30	Save(session Session) (Session, error)
 31	Delete(id string) error
 32}
 33
 34type service struct {
 35	*pubsub.Broker[Session]
 36	q   db.Querier
 37	ctx context.Context
 38}
 39
 40func (s *service) Create(title string) (Session, error) {
 41	dbSession, err := s.q.CreateSession(s.ctx, db.CreateSessionParams{
 42		ID:    uuid.New().String(),
 43		Title: title,
 44	})
 45	if err != nil {
 46		return Session{}, err
 47	}
 48	session := s.fromDBItem(dbSession)
 49	s.Publish(pubsub.CreatedEvent, session)
 50	return session, nil
 51}
 52
 53func (s *service) CreateTaskSession(toolCallID, parentSessionID, title string) (Session, error) {
 54	dbSession, err := s.q.CreateSession(s.ctx, db.CreateSessionParams{
 55		ID:              toolCallID,
 56		ParentSessionID: sql.NullString{String: parentSessionID, Valid: true},
 57		Title:           title,
 58	})
 59	if err != nil {
 60		return Session{}, err
 61	}
 62	session := s.fromDBItem(dbSession)
 63	s.Publish(pubsub.CreatedEvent, session)
 64	return session, nil
 65}
 66
 67func (s *service) Delete(id string) error {
 68	session, err := s.Get(id)
 69	if err != nil {
 70		return err
 71	}
 72	err = s.q.DeleteSession(s.ctx, session.ID)
 73	if err != nil {
 74		return err
 75	}
 76	s.Publish(pubsub.DeletedEvent, session)
 77	return nil
 78}
 79
 80func (s *service) Get(id string) (Session, error) {
 81	dbSession, err := s.q.GetSessionByID(s.ctx, id)
 82	if err != nil {
 83		return Session{}, err
 84	}
 85	return s.fromDBItem(dbSession), nil
 86}
 87
 88func (s *service) Save(session Session) (Session, error) {
 89	dbSession, err := s.q.UpdateSession(s.ctx, db.UpdateSessionParams{
 90		ID:               session.ID,
 91		Title:            session.Title,
 92		PromptTokens:     session.PromptTokens,
 93		CompletionTokens: session.CompletionTokens,
 94		Cost:             session.Cost,
 95	})
 96	if err != nil {
 97		return Session{}, err
 98	}
 99	session = s.fromDBItem(dbSession)
100	s.Publish(pubsub.UpdatedEvent, session)
101	return session, nil
102}
103
104func (s *service) List() ([]Session, error) {
105	dbSessions, err := s.q.ListSessions(s.ctx)
106	if err != nil {
107		return nil, err
108	}
109	sessions := make([]Session, len(dbSessions))
110	for i, dbSession := range dbSessions {
111		sessions[i] = s.fromDBItem(dbSession)
112	}
113	return sessions, nil
114}
115
116func (s service) fromDBItem(item db.Session) Session {
117	return Session{
118		ID:               item.ID,
119		ParentSessionID:  item.ParentSessionID.String,
120		Title:            item.Title,
121		MessageCount:     item.MessageCount,
122		PromptTokens:     item.PromptTokens,
123		CompletionTokens: item.CompletionTokens,
124		Cost:             item.Cost,
125		CreatedAt:        item.CreatedAt,
126		UpdatedAt:        item.UpdatedAt,
127	}
128}
129
130func NewService(ctx context.Context, q db.Querier) Service {
131	broker := pubsub.NewBroker[Session]()
132	return &service{
133		broker,
134		q,
135		ctx,
136	}
137}