session.go 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  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(title string) (Session, error)
  23. CreateTaskSession(toolCallID, parentSessionID, title string) (Session, error)
  24. Get(id string) (Session, error)
  25. List() ([]Session, error)
  26. Save(session Session) (Session, error)
  27. Delete(id string) error
  28. }
  29. type service struct {
  30. *pubsub.Broker[Session]
  31. q db.Querier
  32. ctx context.Context
  33. }
  34. func (s *service) Create(title string) (Session, error) {
  35. dbSession, err := s.q.CreateSession(s.ctx, db.CreateSessionParams{
  36. ID: uuid.New().String(),
  37. Title: title,
  38. })
  39. if err != nil {
  40. return Session{}, err
  41. }
  42. session := s.fromDBItem(dbSession)
  43. s.Publish(pubsub.CreatedEvent, session)
  44. return session, nil
  45. }
  46. func (s *service) CreateTaskSession(toolCallID, parentSessionID, title string) (Session, error) {
  47. dbSession, err := s.q.CreateSession(s.ctx, db.CreateSessionParams{
  48. ID: toolCallID,
  49. ParentSessionID: sql.NullString{String: parentSessionID, Valid: true},
  50. Title: title,
  51. })
  52. if err != nil {
  53. return Session{}, err
  54. }
  55. session := s.fromDBItem(dbSession)
  56. s.Publish(pubsub.CreatedEvent, session)
  57. return session, nil
  58. }
  59. func (s *service) Delete(id string) error {
  60. session, err := s.Get(id)
  61. if err != nil {
  62. return err
  63. }
  64. err = s.q.DeleteSession(s.ctx, session.ID)
  65. if err != nil {
  66. return err
  67. }
  68. s.Publish(pubsub.DeletedEvent, session)
  69. return nil
  70. }
  71. func (s *service) Get(id string) (Session, error) {
  72. dbSession, err := s.q.GetSessionByID(s.ctx, id)
  73. if err != nil {
  74. return Session{}, err
  75. }
  76. return s.fromDBItem(dbSession), nil
  77. }
  78. func (s *service) Save(session Session) (Session, error) {
  79. dbSession, err := s.q.UpdateSession(s.ctx, db.UpdateSessionParams{
  80. ID: session.ID,
  81. Title: session.Title,
  82. PromptTokens: session.PromptTokens,
  83. CompletionTokens: session.CompletionTokens,
  84. Cost: session.Cost,
  85. })
  86. if err != nil {
  87. return Session{}, err
  88. }
  89. session = s.fromDBItem(dbSession)
  90. s.Publish(pubsub.UpdatedEvent, session)
  91. return session, nil
  92. }
  93. func (s *service) List() ([]Session, error) {
  94. dbSessions, err := s.q.ListSessions(s.ctx)
  95. if err != nil {
  96. return nil, err
  97. }
  98. sessions := make([]Session, len(dbSessions))
  99. for i, dbSession := range dbSessions {
  100. sessions[i] = s.fromDBItem(dbSession)
  101. }
  102. return sessions, nil
  103. }
  104. func (s service) fromDBItem(item db.Session) Session {
  105. return Session{
  106. ID: item.ID,
  107. ParentSessionID: item.ParentSessionID.String,
  108. Title: item.Title,
  109. MessageCount: item.MessageCount,
  110. PromptTokens: item.PromptTokens,
  111. CompletionTokens: item.CompletionTokens,
  112. Cost: item.Cost,
  113. CreatedAt: item.CreatedAt,
  114. UpdatedAt: item.UpdatedAt,
  115. }
  116. }
  117. func NewService(ctx context.Context, q db.Querier) Service {
  118. broker := pubsub.NewBroker[Session]()
  119. return &service{
  120. broker,
  121. q,
  122. ctx,
  123. }
  124. }