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(ctx context.Context, title string) (Session, error)
27 CreateTitleSession(ctx context.Context, parentSessionID string) (Session, error)
28 CreateTaskSession(ctx context.Context, toolCallID, parentSessionID, title string) (Session, error)
29 Get(ctx context.Context, id string) (Session, error)
30 List(ctx context.Context) ([]Session, error)
31 Save(ctx context.Context, session Session) (Session, error)
32 Delete(ctx context.Context, id string) error
33}
34
35type service struct {
36 *pubsub.Broker[Session]
37 q db.Querier
38}
39
40func (s *service) Create(ctx context.Context, title string) (Session, error) {
41 dbSession, err := s.q.CreateSession(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(ctx context.Context, toolCallID, parentSessionID, title string) (Session, error) {
54 dbSession, err := s.q.CreateSession(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) CreateTitleSession(ctx context.Context, parentSessionID string) (Session, error) {
68 dbSession, err := s.q.CreateSession(ctx, db.CreateSessionParams{
69 ID: "title-" + parentSessionID,
70 ParentSessionID: sql.NullString{String: parentSessionID, Valid: true},
71 Title: "Generate a title",
72 })
73 if err != nil {
74 return Session{}, err
75 }
76 session := s.fromDBItem(dbSession)
77 s.Publish(pubsub.CreatedEvent, session)
78 return session, nil
79}
80
81func (s *service) Delete(ctx context.Context, id string) error {
82 session, err := s.Get(ctx, id)
83 if err != nil {
84 return err
85 }
86 err = s.q.DeleteSession(ctx, session.ID)
87 if err != nil {
88 return err
89 }
90 s.Publish(pubsub.DeletedEvent, session)
91 return nil
92}
93
94func (s *service) Get(ctx context.Context, id string) (Session, error) {
95 dbSession, err := s.q.GetSessionByID(ctx, id)
96 if err != nil {
97 return Session{}, err
98 }
99 return s.fromDBItem(dbSession), nil
100}
101
102func (s *service) Save(ctx context.Context, session Session) (Session, error) {
103 dbSession, err := s.q.UpdateSession(ctx, db.UpdateSessionParams{
104 ID: session.ID,
105 Title: session.Title,
106 PromptTokens: session.PromptTokens,
107 CompletionTokens: session.CompletionTokens,
108 Cost: session.Cost,
109 })
110 if err != nil {
111 return Session{}, err
112 }
113 session = s.fromDBItem(dbSession)
114 s.Publish(pubsub.UpdatedEvent, session)
115 return session, nil
116}
117
118func (s *service) List(ctx context.Context) ([]Session, error) {
119 dbSessions, err := s.q.ListSessions(ctx)
120 if err != nil {
121 return nil, err
122 }
123 sessions := make([]Session, len(dbSessions))
124 for i, dbSession := range dbSessions {
125 sessions[i] = s.fromDBItem(dbSession)
126 }
127 return sessions, nil
128}
129
130func (s service) fromDBItem(item db.Session) Session {
131 return Session{
132 ID: item.ID,
133 ParentSessionID: item.ParentSessionID.String,
134 Title: item.Title,
135 MessageCount: item.MessageCount,
136 PromptTokens: item.PromptTokens,
137 CompletionTokens: item.CompletionTokens,
138 Cost: item.Cost,
139 CreatedAt: item.CreatedAt,
140 UpdatedAt: item.UpdatedAt,
141 }
142}
143
144func NewService(q db.Querier) Service {
145 broker := pubsub.NewBroker[Session]()
146 return &service{
147 broker,
148 q,
149 }
150}