Merge pull request #25 from bat-bs/reporting

Add cost collector and caching

Author: hofq

api/apiserver.go

@@ -18,11 +18,12 @@ func ApiInit(mux *http.ServeMux, a *auth.Auth, db *db.Database) {
 	}
 
 	api := ApiHandler{db, a, timeZone}
+	graph := NewGraphHandler(&api)
 	mux.HandleFunc("/api2/user/widget", api.GetUserWidget)
 	mux.HandleFunc("/api2/user/logout", api.LogoutUser)
 	mux.HandleFunc("/api2/admin/table/get", api.GetAdminTable)
-	mux.HandleFunc("/api2/admin/table/graph/get/", api.GetAdminTableGraph)
-	mux.HandleFunc("/api2/table/graph/get/", api.GetTableGraph)
+	mux.HandleFunc("/api2/admin/table/graph/get/", graph.GetAdminTableGraph)
+	mux.HandleFunc("/api2/table/graph/get/", graph.GetTableGraph)
 	mux.HandleFunc("/api2/table/get", api.GetTable)
 	mux.HandleFunc("/api2/table/entry/save", api.CreateEntry)
 	mux.HandleFunc("/api2/table/entry/delete/", api.DeleteEntry)

api/graph.go

@@ -15,21 +15,41 @@ import (
 	"github.com/go-echarts/go-echarts/v2/opts"
 )
 
+type GraphHandler struct {
+	a     *ApiHandler
+	cache []Cache
+}
+
+// if basecount, filter and DB Rows missmatch, Cache will be updated
+type Cache struct {
+	ID        string
+	BaseCount int //
+	Filter    string
+	Data      []db.RequestSummary
+}
+
+func NewGraphHandler(a *ApiHandler) *GraphHandler {
+	var cache []Cache
+	return &GraphHandler{a, cache}
+}
+
+// Generate Token Graphs for API Key and Admin overview and allow to filter by timeframes
+
 type Graph struct {
 	key  string
 	kind string // can be "user" or "apiKey"
 	w    http.ResponseWriter
 	r    *http.Request
 }
 
-func (a *ApiHandler) GetTableGraph(w http.ResponseWriter, r *http.Request) {
-	ok := a.auth.ValidateSessionToken(w, r)
+func (g *GraphHandler) GetTableGraph(w http.ResponseWriter, r *http.Request) {
+	ok := g.a.auth.ValidateSessionToken(w, r)
 	if !ok {
 		http.Error(w, "Not Authorized", http.StatusForbidden)
 		return
 	}
 	key := strings.TrimPrefix(r.URL.Path, "/api2/table/graph/get/")
-	a.RenderGraph(&Graph{
+	g.RenderGraph(&Graph{
 		key:  key,
 		kind: "apiKey",
 		w:    w,
@@ -38,16 +58,16 @@ func (a *ApiHandler) GetTableGraph(w http.ResponseWriter, r *http.Request) {
 
 }
 
-func (a *ApiHandler) GetAdminTableGraph(w http.ResponseWriter, r *http.Request) {
+func (g *GraphHandler) GetAdminTableGraph(w http.ResponseWriter, r *http.Request) {
 
-	ok, err := a.auth.ValidateAdminSession(w, r)
+	ok, err := g.a.auth.ValidateAdminSession(w, r)
 	if err != nil || !ok {
 		http.Error(w, "Not Authorized", http.StatusForbidden)
 		return
 	}
 
 	key := strings.TrimPrefix(r.URL.Path, "/api2/admin/table/graph/get/")
-	a.RenderGraph(&Graph{
+	g.RenderGraph(&Graph{
 		key:  key,
 		kind: "user",
 		w:    w,
@@ -56,10 +76,10 @@ func (a *ApiHandler) GetAdminTableGraph(w http.ResponseWriter, r *http.Request)
 
 }
 
-func (a *ApiHandler) RenderGraph(g *Graph) {
+func (g *GraphHandler) RenderGraph(gr *Graph) {
 
 	selectedfilter := "24h" // default Value
-	header := g.r.Header.Get("HX-Current-URL")
+	header := gr.r.Header.Get("HX-Current-URL")
 	currentURL, err := url.Parse(header)
 	if err != nil {
 		log.Println("Error Parsing HX-Current-URL for Graph Table")
@@ -90,13 +110,6 @@ func (a *ApiHandler) RenderGraph(g *Graph) {
 		filter = "This Year"
 	}
 
-	data, err := a.db.LookupApiKeyUserStats(g.key, g.kind, filter)
-	if err != nil {
-		log.Println(err)
-		http.Error(g.w, "Could not get Data from DB for User "+string(g.key), 500)
-		return
-	}
-
 	// create a new line instance
 	line := charts.NewLine()
 	// set some global options like Title/Legend/ToolTip or anything else
@@ -108,9 +121,10 @@ func (a *ApiHandler) RenderGraph(g *Graph) {
 		charts.WithYAxisOpts(opts.YAxis{SplitNumber: 2}),
 		// charts.WithVisualMapOpts(opts.VisualMap{Show: opts.Bool(false)})
 	)
+	data := g.GetTableGraphData(gr, filter)
 
 	// Put data into instance
-	td, err := a.GetAdminTableGraphData(g.w, data, filter)
+	td, err := g.SetTableGraphData(gr.w, data, filter)
 	if err != nil {
 		return
 	}
@@ -141,7 +155,7 @@ func (a *ApiHandler) RenderGraph(g *Graph) {
 		Filter:     filter,
 	}
 	// var buf bytes.Buffer
-	if err := t.Execute(g.w, snippetData); err != nil {
+	if err := t.Execute(gr.w, snippetData); err != nil {
 		log.Println("Error Templating Chart", err)
 		return
 	}
@@ -157,7 +171,46 @@ type TableData struct {
 	totalCount int
 }
 
-func (a *ApiHandler) GetAdminTableGraphData(w http.ResponseWriter, d []db.RequestSummary, filter string) (*TableData, error) {
+func (g *GraphHandler) GetTableGraphData(gr *Graph, filter string) []db.RequestSummary {
+	for _, row := range g.cache {
+		if row.ID == gr.key && row.Filter == filter {
+			count, err := g.a.db.LookupApiKeyUserStatsRows(gr.key, gr.kind)
+			if err == nil && count == row.BaseCount {
+				return row.Data
+			}
+			if err != nil {
+				log.Println(err)
+			}
+			row.Data = g.LookupTableGraphData(gr, filter)
+			row.BaseCount = count
+		}
+	}
+	rowCount, err := g.a.db.LookupApiKeyUserStatsRows(gr.key, gr.kind)
+	if err != nil {
+		log.Println("could not count rows for caching: ", err)
+	}
+
+	row := Cache{
+		Data:      g.LookupTableGraphData(gr, filter),
+		BaseCount: rowCount,
+		Filter:    filter,
+		ID:        gr.key,
+	}
+	g.cache = append(g.cache, row)
+	return row.Data
+}
+
+func (g *GraphHandler) LookupTableGraphData(gr *Graph, filter string) []db.RequestSummary {
+	data, err := g.a.db.LookupApiKeyUserStats(gr.key, gr.kind, filter)
+	if err != nil && data != nil {
+		log.Println(err)
+		http.Error(gr.w, "Could not get Data from DB for User "+string(gr.key), 500)
+		return nil
+	}
+	return data
+}
+
+func (g *GraphHandler) SetTableGraphData(w http.ResponseWriter, d []db.RequestSummary, filter string) (*TableData, error) {
 
 	td := &TableData{
 		data:     make([]opts.LineData, 0),
@@ -188,7 +241,7 @@ func (a *ApiHandler) GetAdminTableGraphData(w http.ResponseWriter, d []db.Reques
 		totalTokens = item.TokenCountComplete + item.TokenCountPrompt
 		td.data = append(td.data, opts.LineData{Value: totalTokens})
 		td.totalCount = td.totalCount + totalTokens
-		loc, err := time.LoadLocation(a.timeZone)
+		loc, err := time.LoadLocation(g.a.timeZone)
 		if err != nil {
 			log.Println("Error Displaying Timezone, maybe the TIMEZONE env is wrongly set")
 			td.timeAxis = append(td.timeAxis, item.RequestTime.Format(format))

apiproxy/response.go

@@ -11,6 +11,8 @@ import (
 	"strings"
 )
 
+// Process Response of Azure and Write Key Parameters to DB
+
 type ResponseConf struct {
 	db *db.Database
 }

cmd/main.go

@@ -6,11 +6,13 @@ import (
 	api "openai-api-proxy/api"
 	proxy "openai-api-proxy/apiproxy"
 	auth "openai-api-proxy/auth"
+	costs "openai-api-proxy/costs"
 	db "openai-api-proxy/db"
 	web "openai-api-proxy/webui"
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/joho/godotenv"
 )
@@ -23,6 +25,11 @@ func main() {
 	}
 	db := db.DatabaseInit()
 	defer db.Close()
+
+	ticker := time.NewTicker(24 * time.Hour)
+
+	go costs.GetAllCosts(db, ticker.C)
+
 	osExit(db)
 	defer log.Println("Closing DB Clients :)")
 

costs/azureCostCollector.go

@@ -0,0 +1,117 @@
+package costs
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	db "openai-api-proxy/db"
+	"regexp"
+	"time"
+)
+
+// Collect Cost of Azure DB and Write them to DB
+
+type BaseJson struct {
+	Costs []Costs `json:"Items"`
+}
+type Costs struct {
+	ModelName     string
+	RetailPrice   float32 `json:"retailPrice"`
+	SKUName       string  `json:"skuName"`
+	UnitOfMeasure string  `json:"unitOfMeasure"`
+	TokenType     string
+	Currency      string
+	IsRegional    bool
+}
+
+const MoneyUnit = 10000000
+
+func GetAllCosts(d *db.Database, g <-chan time.Time) {
+	for ; true; <-g {
+		log.Println("Azure: Collecting Prices")
+		totalCosts := []*Costs{}
+		tokenTypes := []string{"Input", "Output"}
+		models := d.LookupModels()
+		cleanModel := regexp.MustCompile(`-[0-9]{4}-[0-9]{2}-[0-9]{2}`)
+		if len(models) == 0 {
+			log.Println("No Models in DB yet, no costs can be calculated")
+			return
+		}
+
+		for _, model := range models {
+			modelName := cleanModel.ReplaceAllString(model, "")
+			for _, tokenType := range tokenTypes {
+				c := GetCosts(modelName, tokenType)
+				c.TokenType = tokenType
+				c.ModelName = modelName
+				totalCosts = append(totalCosts, c)
+			}
+		}
+		dbcosts := []*db.Costs{}
+		for _, costs := range totalCosts {
+			dbc := &db.Costs{
+				ModelName:     costs.ModelName,
+				RetailPrice:   int64(costs.RetailPrice * float32(MoneyUnit)),
+				TokenType:     costs.TokenType,
+				UnitOfMeasure: costs.UnitOfMeasure,
+				IsRegional:    costs.IsRegional,
+				BackendName:   "azure",
+			}
+			dbcosts = append(dbcosts, dbc)
+		}
+		err := d.WriteCosts(dbcosts)
+		if err != nil {
+			log.Println(err)
+		}
+	}
+}
+
+func GetCosts(modelName string, tokenType string) *Costs {
+	currency := "EUR"
+	regionName := "swedencentral"
+	productName := "Azure OpenAI"
+	regional := true
+	errorContext := fmt.Sprintf("Check ENV - regional: %s, regionName: %s, currency: %s", regional, regionName, currency)
+
+	var region string
+
+	if regional {
+		region = "regional"
+	} else {
+		region = "global"
+	}
+
+	rq, err := http.NewRequest("GET", "https://prices.azure.com/api/retail/prices", nil)
+	if err != nil {
+		log.Println("Could not get latest Prices from API", errorContext)
+	}
+
+	filter := fmt.Sprintf("productName eq '%s' and armRegionName eq '%s' and skuName eq '%s-%s-%s'", productName, regionName, modelName, tokenType, region)
+	q := rq.URL.Query()
+	q.Add("currencyCode", currency)
+	q.Add("$filter", filter)
+	rq.URL.RawQuery = q.Encode()
+
+	client := &http.Client{}
+	resp, err := client.Do(rq)
+	if err != nil {
+		log.Println("Error Requesting Data from Azure API.", errorContext)
+	}
+	defer resp.Body.Close()
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		log.Println("Error Parsing Body")
+		log.Println(err)
+	}
+	var base *BaseJson
+
+	resp.Body = io.NopCloser(bytes.NewReader(body))
+	json.Unmarshal(body, &base)
+	c := &base.Costs[0]
+	c.IsRegional = regional
+	c.Currency = currency
+	return c
+}