@@ -11,6 +11,23 @@ import (
1111 "github.com/leporo/sqlf"
1212)
1313
14+ type AiUsageReport struct {
15+ TeamId string `json:"team_id"`
16+ Model string `json:"model"`
17+ InputTokens uint64 `json:"input_tokens"`
18+ OutputTokens uint64 `json:"output_tokens"`
19+ }
20+
21+ type AiUsage struct {
22+ TeamId string `json:"team_id"`
23+ MonthlyAiUsage []MonthlyAiUsage `json:"monthly_ai_usage"`
24+ }
25+
26+ type MonthlyAiUsage struct {
27+ MonthName string `json:"month_year"`
28+ TotalTokenCount uint64 `json:"total_token_count"`
29+ }
30+
1431type AppUsage struct {
1532 AppId string `json:"app_id"`
1633 AppName string `json:"app_name"`
@@ -159,6 +176,61 @@ func GetUsage(c *gin.Context) {
159176 return
160177 }
161178
179+ // Query ai metrics for team
180+ aiMetricsStmt := sqlf .
181+ From (`ai_metrics` ).
182+ Select ("team_id" ).
183+ Select ("formatDateTime(toStartOfMonth(timestamp), '%b %Y') AS month_year" ).
184+ Select ("sumMerge(input_token_count) + sumMerge(output_token_count) AS total_token_count" ).
185+ Where ("timestamp >= addMonths(toStartOfMonth(?), -2) AND timestamp < toStartOfMonth(addMonths(?, 1))" , now , now ).
186+ GroupBy ("team_id, toStartOfMonth(timestamp)" ).
187+ OrderBy ("team_id, toStartOfMonth(timestamp) DESC" )
188+
189+ defer aiMetricsStmt .Close ()
190+
191+ aiMetricsRows , err := server .Server .ChPool .Query (ctx , aiMetricsStmt .String (), aiMetricsStmt .Args ()... )
192+ if err != nil {
193+ msg := fmt .Sprintf ("error occurred while querying AI metrics for team: %s" , teamId )
194+ fmt .Println (msg , err )
195+ c .JSON (http .StatusInternalServerError , gin.H {"error" : msg })
196+ return
197+ }
198+
199+ aiUsageMap := make (map [string ]* AiUsage )
200+
201+ // Initialize aiUsageMap
202+ aiUsageMap [teamId .String ()] = & AiUsage {
203+ TeamId : teamId .String (),
204+ MonthlyAiUsage : make ([]MonthlyAiUsage , 0 , 3 ),
205+ }
206+
207+ // Populate aiUsageMap with metrics rows from DB
208+ for aiMetricsRows .Next () {
209+ var teamId , monthYear string
210+ var totalTokenCount uint64
211+
212+ if err := aiMetricsRows .Scan (& teamId , & monthYear , & totalTokenCount ); err != nil {
213+ msg := fmt .Sprintf ("error occurred while scanning AI metrics row for team: %s" , teamId )
214+ fmt .Println (msg , err )
215+ c .JSON (http .StatusInternalServerError , gin.H {"error" : msg })
216+ return
217+ }
218+
219+ if aiUsage , exists := aiUsageMap [teamId ]; exists {
220+ aiUsage .MonthlyAiUsage = append (aiUsage .MonthlyAiUsage , MonthlyAiUsage {
221+ MonthName : monthYear ,
222+ TotalTokenCount : totalTokenCount ,
223+ })
224+ }
225+ }
226+
227+ if err := aiMetricsRows .Err (); err != nil {
228+ msg := fmt .Sprintf ("error occurred while iterating AI metrics rows for team: %s" , teamId )
229+ fmt .Println (msg , err )
230+ c .JSON (http .StatusInternalServerError , gin.H {"error" : msg })
231+ return
232+ }
233+
162234 // Ensure all apps have entries for all three months by adding 0 values for missing months
163235 for _ , appUsage := range appUsageMap {
164236 monthDataMap := make (map [string ]MonthlyAppUsage )
@@ -183,11 +255,115 @@ func GetUsage(c *gin.Context) {
183255 appUsage .MonthlyAppUsage = newMonthlyAppUsage
184256 }
185257
258+ // Ensure all AI usages have entries for all three months by adding 0 values for missing months
259+ for _ , aiUsage := range aiUsageMap {
260+ monthDataMap := make (map [string ]MonthlyAiUsage )
261+ for _ , usage := range aiUsage .MonthlyAiUsage {
262+ monthDataMap [usage .MonthName ] = usage
263+ }
264+
265+ newMonthlyAiUsage := make ([]MonthlyAiUsage , 0 , 3 )
266+ for _ , monthName := range monthNames {
267+ if usage , exists := monthDataMap [monthName ]; exists {
268+ newMonthlyAiUsage = append (newMonthlyAiUsage , usage )
269+ } else {
270+ newMonthlyAiUsage = append (newMonthlyAiUsage , MonthlyAiUsage {
271+ MonthName : monthName ,
272+ TotalTokenCount : 0 ,
273+ })
274+ }
275+ }
276+ aiUsage .MonthlyAiUsage = newMonthlyAiUsage
277+ }
278+
186279 // Convert map to slice for JSON response
187- var result []AppUsage
280+ var appUsageResult []AppUsage
188281 for _ , appUsage := range appUsageMap {
189- result = append (result , * appUsage )
282+ appUsageResult = append (appUsageResult , * appUsage )
283+ }
284+
285+ var aiUsageResult []AiUsage
286+ for _ , aiUsage := range aiUsageMap {
287+ aiUsageResult = append (aiUsageResult , * aiUsage )
288+ }
289+
290+ result := gin.H {
291+ "app_usage" : appUsageResult ,
292+ "ai_usage" : aiUsageResult ,
190293 }
191294
192295 c .JSON (http .StatusOK , result )
193296}
297+
298+ func ReportAiUsage (c * gin.Context ) {
299+ ctx := c .Request .Context ()
300+ userId := c .GetString ("userId" )
301+ teamId , err := uuid .Parse (c .Param ("id" ))
302+ if err != nil {
303+ msg := `team id invalid or missing`
304+ fmt .Println (msg , err )
305+ c .JSON (http .StatusBadRequest , gin.H {"error" : msg })
306+ return
307+ }
308+
309+ var aiUsageReport AiUsageReport
310+
311+ if err := c .ShouldBindJSON (& aiUsageReport ); err != nil {
312+ msg := "failed to parse AI usage payload"
313+ fmt .Println (msg , err )
314+ c .JSON (http .StatusBadRequest , gin.H {
315+ "error" : msg ,
316+ })
317+ return
318+ }
319+
320+ if ok , err := PerformAuthz (userId , teamId .String (), * ScopeTeamRead ); err != nil {
321+ msg := `couldn't perform authorization checks`
322+ fmt .Println (msg , err )
323+ c .JSON (http .StatusInternalServerError , gin.H {"error" : msg })
324+ return
325+ } else if ! ok {
326+ msg := fmt .Sprintf (`you don't have permissions for team [%s]` , teamId )
327+ c .JSON (http .StatusForbidden , gin.H {"error" : msg })
328+ return
329+ }
330+
331+ if ok , err := PerformAuthz (userId , teamId .String (), * ScopeAppRead ); err != nil {
332+ msg := `couldn't perform authorization checks`
333+ fmt .Println (msg , err )
334+ c .JSON (http .StatusInternalServerError , gin.H {"error" : msg })
335+ return
336+ } else if ! ok {
337+ msg := fmt .Sprintf (`you don't have permissions to read apps in team [%s]` , teamId )
338+ c .JSON (http .StatusForbidden , gin.H {"error" : msg })
339+ return
340+ }
341+
342+ var team = new (Team )
343+ team .ID = & teamId
344+
345+ // insert metrics into clickhouse table
346+ metricsSelectStmt := sqlf .
347+ Select ("? AS team_id" , team .ID ).
348+ Select ("? AS timestamp" , time .Now ()).
349+ Select ("? AS model" , aiUsageReport .Model ).
350+ Select ("sumState(CAST(? AS UInt32)) AS input_token_count" , aiUsageReport .InputTokens ).
351+ Select ("sumState(CAST(? AS UInt32)) AS output_token_count" , aiUsageReport .OutputTokens )
352+ selectSQL := metricsSelectStmt .String ()
353+ args := metricsSelectStmt .Args ()
354+ defer metricsSelectStmt .Close ()
355+ metricsInsertStmt := "INSERT INTO ai_metrics " + selectSQL
356+
357+ if err := server .Server .ChPool .Exec (ctx , metricsInsertStmt , args ... ); err != nil {
358+ msg := `failed to insert ai metrics into clickhouse`
359+ fmt .Println (msg , err )
360+ c .JSON (http .StatusInternalServerError , gin.H {
361+ "error" : msg ,
362+ })
363+ return
364+ }
365+
366+ c .JSON (http .StatusOK , gin.H {
367+ "ok" : "done" ,
368+ })
369+ }
0 commit comments