Skip to content

Commit 4206592

Browse files
wjiayiswjiayis
committed
feat: LLM usage tracking, backend
1 parent 2e593f7 commit 4206592

24 files changed

Lines changed: 1464 additions & 19 deletions

‎go.sum‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,8 @@ github.com/golang/snappy v1.0.0/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEW
5555
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
5656
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
5757
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
58+
github.com/google/subcommands v1.2.0 h1:vWQspBTo2nEqTUFita5/KeEWlUL8kQObDFbub/EN9oE=
59+
github.com/google/subcommands v1.2.0/go.mod h1:ZjhPrFU+Olkh9WazFPsl27BQ4UPiG37m3yTrtFlrHVk=
5860
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
5961
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
6062
github.com/google/wire v0.7.0 h1:JxUKI6+CVBgCO2WToKy/nQk0sS+amI9z9EjVmdaocj4=
@@ -160,6 +162,8 @@ golang.org/x/crypto v0.41.0/go.mod h1:pO5AFd7FA68rFak7rOAGVuygIISepHftHnr8dr6+sU
160162
golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b h1:DXr+pvt3nC887026GRP39Ej11UATqWDmWuS99x26cD0=
161163
golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b/go.mod h1:4QTo5u+SEIbbKW1RacMZq1YEfOBqeXa19JeshGi+zc4=
162164
golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
165+
golang.org/x/mod v0.27.0 h1:kb+q2PyFnEADO2IEF935ehFUXlWiNjJWtRNgBLSfbxQ=
166+
golang.org/x/mod v0.27.0/go.mod h1:rWI627Fq0DEoudcK+MBkNkCe0EetEaDSwJJkCcjpazc=
163167
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
164168
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
165169
golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
@@ -188,6 +192,8 @@ golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU=
188192
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
189193
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
190194
golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
195+
golang.org/x/tools v0.36.0 h1:kWS0uv/zsvHEle1LbV5LE8QujrxB3wfQyxHfhOk0Qkg=
196+
golang.org/x/tools v0.36.0/go.mod h1:WBDiHKJK8YgLHlcQPYQzNCkUxUypCaa5ZegCVutKm+s=
191197
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
192198
gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
193199
gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=

‎internal/api/chat/create_conversation_message_stream_v2.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -281,7 +281,7 @@ func (s *ChatServerV2) CreateConversationMessageStream(
281281
APIKey: settings.OpenAIAPIKey,
282282
}
283283

284-
openaiChatHistory, inappChatHistory, err := s.aiClientV2.ChatCompletionStreamV2(ctx, stream, conversation.ID.Hex(), modelSlug, conversation.OpenaiChatHistoryCompletion, llmProvider)
284+
openaiChatHistory, inappChatHistory, err := s.aiClientV2.ChatCompletionStreamV2(ctx, stream, conversation.UserID, conversation.ID.Hex(), modelSlug, conversation.OpenaiChatHistoryCompletion, llmProvider)
285285
if err != nil {
286286
return s.sendStreamError(stream, err)
287287
}
@@ -307,7 +307,7 @@ func (s *ChatServerV2) CreateConversationMessageStream(
307307
for i, bsonMsg := range conversation.InappChatHistory {
308308
protoMessages[i] = mapper.BSONToChatMessageV2(bsonMsg)
309309
}
310-
title, err := s.aiClientV2.GetConversationTitleV2(ctx, protoMessages, llmProvider)
310+
title, err := s.aiClientV2.GetConversationTitleV2(ctx, conversation.UserID, protoMessages, llmProvider)
311311
if err != nil {
312312
s.logger.Error("Failed to get conversation title", "error", err, "conversationID", conversation.ID.Hex())
313313
return

‎internal/api/grpc.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
chatv2 "paperdebugger/pkg/gen/api/chat/v2"
1616
commentv1 "paperdebugger/pkg/gen/api/comment/v1"
1717
projectv1 "paperdebugger/pkg/gen/api/project/v1"
18+
usagev1 "paperdebugger/pkg/gen/api/usage/v1"
1819
userv1 "paperdebugger/pkg/gen/api/user/v1"
1920

2021
// "github.com/grpc-ecosystem/go-grpc-middleware"
@@ -106,6 +107,7 @@ func NewGrpcServer(
106107
userServer userv1.UserServiceServer,
107108
projectServer projectv1.ProjectServiceServer,
108109
commentServer commentv1.CommentServiceServer,
110+
usageServer usagev1.UsageServiceServer,
109111
) *GrpcServer {
110112
grpcServer := &GrpcServer{}
111113
grpcServer.userService = userService
@@ -121,5 +123,6 @@ func NewGrpcServer(
121123
userv1.RegisterUserServiceServer(grpcServer.Server, userServer)
122124
projectv1.RegisterProjectServiceServer(grpcServer.Server, projectServer)
123125
commentv1.RegisterCommentServiceServer(grpcServer.Server, commentServer)
126+
usagev1.RegisterUsageServiceServer(grpcServer.Server, usageServer)
124127
return grpcServer
125128
}

‎internal/api/server.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717
commentv1 "paperdebugger/pkg/gen/api/comment/v1"
1818
projectv1 "paperdebugger/pkg/gen/api/project/v1"
1919
sharedv1 "paperdebugger/pkg/gen/api/shared/v1"
20+
usagev1 "paperdebugger/pkg/gen/api/usage/v1"
2021
userv1 "paperdebugger/pkg/gen/api/user/v1"
2122

2223
"github.com/gin-gonic/gin"
@@ -105,6 +106,11 @@ func (s *Server) Run(addr string) {
105106
s.logger.Fatalf("failed to register comment service grpc gateway: %v", err)
106107
return
107108
}
109+
err = usagev1.RegisterUsageServiceHandler(context.Background(), mux, client)
110+
if err != nil {
111+
s.logger.Fatalf("failed to register usage service grpc gateway: %v", err)
112+
return
113+
}
108114

109115
s.logger.Infof("[PAPERDEBUGGER] http server listening on %s", addr)
110116
s.ginServer.Any("/_pd/api/*path", func(c *gin.Context) { mux.ServeHTTP(c.Writer, c.Request) })

‎internal/api/usage/get_session_usage.go‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
package usage
2+
3+
import (
4+
"context"
5+
6+
"paperdebugger/internal/libs/contextutil"
7+
usagev1 "paperdebugger/pkg/gen/api/usage/v1"
8+
9+
"google.golang.org/protobuf/types/known/timestamppb"
10+
)
11+
12+
func (s *UsageServer) GetSessionUsage(
13+
ctx context.Context,
14+
req *usagev1.GetSessionUsageRequest,
15+
) (*usagev1.GetSessionUsageResponse, error) {
16+
actor, err := contextutil.GetActor(ctx)
17+
if err != nil {
18+
return nil, err
19+
}
20+
21+
session, err := s.usageService.GetActiveSession(ctx, actor.ID)
22+
if err != nil {
23+
return nil, err
24+
}
25+
26+
if session == nil {
27+
return &usagev1.GetSessionUsageResponse{
28+
Session: nil,
29+
}, nil
30+
}
31+
32+
return &usagev1.GetSessionUsageResponse{
33+
Session: &usagev1.SessionUsage{
34+
Id: session.ID.Hex(),
35+
SessionStart: timestamppb.New(session.SessionStart.Time()),
36+
SessionExpiry: timestamppb.New(session.SessionExpiry.Time()),
37+
PromptTokens: session.PromptTokens,
38+
CompletionTokens: session.CompletionTokens,
39+
TotalTokens: session.TotalTokens,
40+
RequestCount: session.RequestCount,
41+
},
42+
}, nil
43+
}

‎internal/api/usage/get_weekly_usage.go‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package usage
2+
3+
import (
4+
"context"
5+
6+
"paperdebugger/internal/libs/contextutil"
7+
usagev1 "paperdebugger/pkg/gen/api/usage/v1"
8+
)
9+
10+
func (s *UsageServer) GetWeeklyUsage(
11+
ctx context.Context,
12+
req *usagev1.GetWeeklyUsageRequest,
13+
) (*usagev1.GetWeeklyUsageResponse, error) {
14+
actor, err := contextutil.GetActor(ctx)
15+
if err != nil {
16+
return nil, err
17+
}
18+
19+
stats, err := s.usageService.GetWeeklyUsage(ctx, actor.ID)
20+
if err != nil {
21+
return nil, err
22+
}
23+
24+
return &usagev1.GetWeeklyUsageResponse{
25+
Usage: &usagev1.WeeklyUsage{
26+
PromptTokens: stats.PromptTokens,
27+
CompletionTokens: stats.CompletionTokens,
28+
TotalTokens: stats.TotalTokens,
29+
RequestCount: stats.RequestCount,
30+
SessionCount: stats.SessionCount,
31+
},
32+
}, nil
33+
}

‎internal/api/usage/server.go‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
package usage
2+
3+
import (
4+
"paperdebugger/internal/libs/logger"
5+
"paperdebugger/internal/services"
6+
usagev1 "paperdebugger/pkg/gen/api/usage/v1"
7+
)
8+
9+
type UsageServer struct {
10+
usagev1.UnimplementedUsageServiceServer
11+
12+
usageService *services.UsageService
13+
logger *logger.Logger
14+
}
15+
16+
func NewUsageServer(
17+
usageService *services.UsageService,
18+
logger *logger.Logger,
19+
) usagev1.UsageServiceServer {
20+
return &UsageServer{
21+
usageService: usageService,
22+
logger: logger,
23+
}
24+
}

‎internal/libs/db/db.go‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66

77
"paperdebugger/internal/libs/cfg"
88
"paperdebugger/internal/libs/logger"
9+
"paperdebugger/internal/models"
910

1011
"go.mongodb.org/mongo-driver/v2/bson"
1112
"go.mongodb.org/mongo-driver/v2/mongo"
@@ -43,5 +44,33 @@ func NewDB(cfg *cfg.Cfg, logger *logger.Logger) (*DB, error) {
4344
}
4445

4546
logger.Info("[MONGO] initialized")
46-
return &DB{Client: client, cfg: cfg, logger: logger}, nil
47+
48+
db := &DB{Client: client, cfg: cfg, logger: logger}
49+
db.ensureIndexes()
50+
return db, nil
51+
}
52+
53+
// ensureIndexes creates necessary indexes for the database collections.
54+
func (db *DB) ensureIndexes() {
55+
sessions := db.Database("paperdebugger").Collection((models.LLMSession{}).CollectionName())
56+
57+
// TTL index: auto-delete sessions after 30 days
58+
_, err := sessions.Indexes().CreateOne(context.Background(), mongo.IndexModel{
59+
Keys: bson.D{{Key: "session_expiry", Value: 1}},
60+
Options: options.Index().SetExpireAfterSeconds(30 * 24 * 60 * 60),
61+
})
62+
if err != nil {
63+
db.logger.Error("Failed to create TTL index on llm_sessions", "error", err)
64+
}
65+
66+
// Compound index for efficient active session lookups
67+
_, err = sessions.Indexes().CreateOne(context.Background(), mongo.IndexModel{
68+
Keys: bson.D{
69+
{Key: "user_id", Value: 1},
70+
{Key: "session_expiry", Value: -1},
71+
},
72+
})
73+
if err != nil {
74+
db.logger.Error("Failed to create compound index on llm_sessions", "error", err)
75+
}
4776
}

‎internal/models/usage.go‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
package models
2+
3+
import "go.mongodb.org/mongo-driver/v2/bson"
4+
5+
// LLMSession represents a user's session for tracking LLM usage and token counts.
6+
type LLMSession struct {
7+
ID bson.ObjectID `bson:"_id"`
8+
UserID bson.ObjectID `bson:"user_id"`
9+
SessionStart bson.DateTime `bson:"session_start"`
10+
SessionExpiry bson.DateTime `bson:"session_expiry"`
11+
PromptTokens int64 `bson:"prompt_tokens"`
12+
CompletionTokens int64 `bson:"completion_tokens"`
13+
TotalTokens int64 `bson:"total_tokens"`
14+
RequestCount int64 `bson:"request_count"`
15+
}
16+
17+
func (s LLMSession) CollectionName() string {
18+
return "llm_sessions"
19+
}

‎internal/services/toolkit/client/client_v2.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ type AIClientV2 struct {
2020

2121
reverseCommentService *services.ReverseCommentService
2222
projectService *services.ProjectService
23+
usageService *services.UsageService
2324
cfg *cfg.Cfg
2425
logger *logger.Logger
2526
}
@@ -60,6 +61,7 @@ func NewAIClientV2(
6061

6162
reverseCommentService *services.ReverseCommentService,
6263
projectService *services.ProjectService,
64+
usageService *services.UsageService,
6365
cfg *cfg.Cfg,
6466
logger *logger.Logger,
6567
) *AIClientV2 {
@@ -107,6 +109,7 @@ func NewAIClientV2(
107109

108110
reverseCommentService: reverseCommentService,
109111
projectService: projectService,
112+
usageService: usageService,
110113
cfg: cfg,
111114
logger: logger,
112115
}

0 commit comments

Comments
 (0)