session.go 3.4 KB

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