feat: user cost #126
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,38 @@ | ||
| package usage | ||
|
|
||
| import ( | ||
| "context" | ||
|
|
||
| "paperdebugger/internal/libs/contextutil" | ||
| usagev1 "paperdebugger/pkg/gen/api/usage/v1" | ||
|
|
||
| "google.golang.org/protobuf/types/known/timestamppb" | ||
| ) | ||
|
|
||
| func (s *UsageServer) GetSessionUsage( | ||
| ctx context.Context, | ||
| req *usagev1.GetSessionUsageRequest, | ||
| ) (*usagev1.GetSessionUsageResponse, error) { | ||
| actor, err := contextutil.GetActor(ctx) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| session, err := s.usageService.GetActiveSession(ctx, actor.ID) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| if session == nil { | ||
| return &usagev1.GetSessionUsageResponse{ | ||
| Session: nil, | ||
| }, nil | ||
| } | ||
|
|
||
| return &usagev1.GetSessionUsageResponse{ | ||
| Session: &usagev1.SessionUsage{ | ||
| SessionExpiry: timestamppb.New(session.SessionExpiry.Time()), | ||
| TotalTokens: session.TotalTokens, | ||
| }, | ||
| }, nil | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,29 @@ | ||
| package usage | ||
|
|
||
| import ( | ||
| "context" | ||
|
|
||
| "paperdebugger/internal/libs/contextutil" | ||
| usagev1 "paperdebugger/pkg/gen/api/usage/v1" | ||
| ) | ||
|
|
||
| func (s *UsageServer) GetWeeklyUsage( | ||
| ctx context.Context, | ||
| req *usagev1.GetWeeklyUsageRequest, | ||
| ) (*usagev1.GetWeeklyUsageResponse, error) { | ||
| actor, err := contextutil.GetActor(ctx) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| stats, err := s.usageService.GetWeeklyUsage(ctx, actor.ID) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| return &usagev1.GetWeeklyUsageResponse{ | ||
| Usage: &usagev1.WeeklyUsage{ | ||
| TotalTokens: stats.TotalTokens, | ||
| }, | ||
| }, nil | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,24 @@ | ||
| package usage | ||
|
|
||
| import ( | ||
| "paperdebugger/internal/libs/logger" | ||
| "paperdebugger/internal/services" | ||
| usagev1 "paperdebugger/pkg/gen/api/usage/v1" | ||
| ) | ||
|
|
||
| type UsageServer struct { | ||
| usagev1.UnimplementedUsageServiceServer | ||
|
|
||
| usageService *services.UsageService | ||
| logger *logger.Logger | ||
| } | ||
|
|
||
| func NewUsageServer( | ||
| usageService *services.UsageService, | ||
| logger *logger.Logger, | ||
| ) usagev1.UsageServiceServer { | ||
| return &UsageServer{ | ||
| usageService: usageService, | ||
| logger: logger, | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,19 @@ | ||
| package models | ||
|
|
||
| import "go.mongodb.org/mongo-driver/v2/bson" | ||
|
|
||
| // LLMSession represents a user's session for tracking LLM usage and token counts. | ||
| type LLMSession struct { | ||
| ID bson.ObjectID `bson:"_id"` | ||
| UserID bson.ObjectID `bson:"user_id"` | ||
| SessionStart bson.DateTime `bson:"session_start"` | ||
| SessionExpiry bson.DateTime `bson:"session_expiry"` | ||
| PromptTokens int64 `bson:"prompt_tokens"` | ||
| CompletionTokens int64 `bson:"completion_tokens"` | ||
| TotalTokens int64 `bson:"total_tokens"` | ||
| RequestCount int64 `bson:"request_count"` | ||
| } | ||
|
|
||
| func (s LLMSession) CollectionName() string { | ||
| return "llm_sessions" | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -4,11 +4,13 @@ import ( | |
| "context" | ||
| "encoding/json" | ||
| "paperdebugger/internal/models" | ||
| "paperdebugger/internal/services" | ||
| "paperdebugger/internal/services/toolkit/handler" | ||
| chatv2 "paperdebugger/pkg/gen/api/chat/v2" | ||
| "strings" | ||
|
|
||
| "github.com/openai/openai-go/v3" | ||
| "go.mongodb.org/mongo-driver/v2/bson" | ||
| ) | ||
|
|
||
| // define []openai.ChatCompletionMessageParamUnion as OpenAIChatHistory | ||
|
|
@@ -25,8 +27,8 @@ import ( | |
| // 1. The full chat history sent to the language model (including any tool call results). | ||
| // 2. The incremental chat history visible to the user (including tool call results and assistant responses). | ||
| // 3. An error, if any occurred during the process. | ||
| func (a *AIClientV2) ChatCompletionV2(ctx context.Context, userID bson.ObjectID, modelSlug string, messages OpenAIChatHistory, llmProvider *models.LLMProviderConfig) (OpenAIChatHistory, AppChatHistory, error) { | ||
| openaiChatHistory, inappChatHistory, err := a.ChatCompletionStreamV2(ctx, nil, userID, "", modelSlug, messages, llmProvider) | ||
| if err != nil { | ||
| return nil, nil, err | ||
| } | ||
|
|
@@ -54,7 +56,7 @@ func (a *AIClientV2) ChatCompletionV2(ctx context.Context, modelSlug string, mes | |
| // - If tool calls are required, it handles them and appends the results to the chat history, then continues the loop. | ||
| // - If no tool calls are needed, it appends the assistant's response and exits the loop. | ||
| // - Finally, it returns the updated chat histories and any error encountered. | ||
| func (a *AIClientV2) ChatCompletionStreamV2(ctx context.Context, callbackStream chatv2.ChatService_CreateConversationMessageStreamServer, userID bson.ObjectID, conversationId string, modelSlug string, messages OpenAIChatHistory, llmProvider *models.LLMProviderConfig) (OpenAIChatHistory, AppChatHistory, error) { | ||
| openaiChatHistory := messages | ||
| inappChatHistory := AppChatHistory{} | ||
|
|
||
|
|
@@ -97,7 +99,22 @@ func (a *AIClientV2) ChatCompletionStreamV2(ctx context.Context, callbackStream | |
|
|
||
| if len(chunk.Choices) == 0 { | ||
| // Handle usage information | ||
| if chunk.Usage.TotalTokens > 0 { | ||
| // Record usage asynchronously to avoid blocking the response | ||
| go func(usage services.UsageRecord) { | ||
| bgCtx := context.Background() | ||
| if err := a.usageService.RecordUsage(bgCtx, usage); err != nil { | ||
| a.logger.Error("Failed to store usage", "error", err) | ||
| return | ||
| } | ||
|
|
||
| }(services.UsageRecord{ | ||
| UserID: userID, | ||
|
Member
There was a problem hiding this comment. Just a quick check, you are using the I noticed there might be an interesting edge case. Ideally we should recognise and combine / avoid re-generating If its convenient, could you also test and verify during integration testing? Can try two different login methods on the same Overleaf project and we should expect two tokens usage tracking.
Member
Author
There was a problem hiding this comment. Yup, it's the cc: @Junyi-99
Member
There was a problem hiding this comment. yeah switching out emails will generate a different Edit: Not sure if Overleaf has any safeguards / cooldowns on switching out emails too frequently. But yeah, we can keep this in mind since its a separate problem. Ideally |
||
| PromptTokens: chunk.Usage.PromptTokens, | ||
| CompletionTokens: chunk.Usage.CompletionTokens, | ||
| TotalTokens: chunk.Usage.TotalTokens, | ||
| }) | ||
| } | ||
wjiayis marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| continue | ||
| } | ||
|
|
||
|
|
@@ -185,7 +202,6 @@ func (a *AIClientV2) ChatCompletionStreamV2(ctx context.Context, callbackStream | |
| // answer_content += chunk.Choices[0].Delta.Content | ||
| // fmt.Printf("answer_content: %s\n", answer_content) | ||
| streamHandler.HandleTextDoneItem(chunk, answer_content, reasoning_content) | ||
| } | ||
| } | ||
|
|
||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.