Add postgres support and anthropic ai support

benchmarks.go

@@ -225,7 +225,7 @@ func postBenchmarkCreate(c *gin.Context) {
 		return
 	}
 
-	if openaiClient != nil {
+	if openaiClient != nil || anthroClient != nil {
 		go generateSummary(&benchmark, csvFiles)
 	}
 

config.go

@@ -20,6 +20,11 @@ type Config struct {
 	OpenAIApiKey string
 	OpenAIModel  string
 
+	AnthroAPIKey string
+	AnthroModel  string
+
+	DBurl string
+
 	Version bool
 }
 
@@ -38,6 +43,11 @@ func NewConfig() (*Config, error) {
 	flag.StringVar(&config.OpenAIModel, "openai-model", "gpt-4o", "OpenAI model ID")
 	flag.StringVar(&config.OpenAIApiKey, "openai-api-key", "", "OpenAI API Key (leave empty to disable OpenAI integration)")
 
+	flag.StringVar(&config.AnthroAPIKey, "anthro-api-key", "", "Anthropic API Key (leave empty to disable Anthropic integration)")
+	flag.StringVar(&config.AnthroModel, "anthro-model", "claude-3-haiku-20240307", "Anthropic model ID")
+
+	flag.StringVar(&config.DBurl, "db-url", "", "Database URL")
+
 	flag.BoolVar(&config.Version, "version", false, "prints version of the application")
 
 	envflag.Parse(envflag.WithPrefix("FS_"))
@@ -69,6 +79,11 @@ func NewConfig() (*Config, error) {
 			return nil, errors.New("missing openai-url argument")
 		}
 	}
+	if config.AnthroAPIKey != "" {
+		if config.AnthroModel == "" {
+			return nil, errors.New("missing anthro-model argument")
+		}
+	}
 
 	return config, nil
 }

go.mod

@@ -2,12 +2,24 @@ module flightlesssomething
 
 go 1.22
 
-require github.com/gin-gonic/gin v1.10.0
+require (
+	github.com/gin-gonic/gin v1.10.0
+	github.com/liushuangls/go-anthropic v1.6.0
+	gorm.io/driver/postgres v1.4.1
+)
 
 require (
 	github.com/gorilla/context v1.1.2 // indirect
 	github.com/gorilla/securecookie v1.1.2 // indirect
 	github.com/gorilla/sessions v1.2.2 // indirect
+	github.com/jackc/chunkreader/v2 v2.0.1 // indirect
+	github.com/jackc/pgconn v1.13.0 // indirect
+	github.com/jackc/pgio v1.0.0 // indirect
+	github.com/jackc/pgpassfile v1.0.0 // indirect
+	github.com/jackc/pgproto3/v2 v2.3.1 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b // indirect
+	github.com/jackc/pgtype v1.12.0 // indirect
+	github.com/jackc/pgx/v4 v4.17.2 // indirect
 	github.com/wader/gormstore/v2 v2.0.3 // indirect
 )
 

go.sum

@@ -1,4 +1,5 @@
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
+github.com/Masterminds/semver/v3 v3.1.1 h1:hLg3sBzpNErnxhQtUy/mmLR2I9foDujNK030IGemrRc=
 github.com/Masterminds/semver/v3 v3.1.1/go.mod h1:VPu/7SZ7ePZ3QOrcuXROw5FAcLl4a0cBrbBpGY/8hQs=
 github.com/bytedance/sonic v1.11.6 h1:oUp34TzMlL+OY1OUWxHqsdkgC/Zfc85zGqw9siXjrc0=
 github.com/bytedance/sonic v1.11.6/go.mod h1:LysEHSvpvDySVdC2f87zGWf6CIKJcAvqab1ZaiQtds4=
@@ -8,6 +9,7 @@ github.com/cloudwego/base64x v0.1.4 h1:jwCgWpFanWmN8xoIUHa2rtzmkd5J2plF/dnLS6Xd/
 github.com/cloudwego/base64x v0.1.4/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
 github.com/cloudwego/iasm v0.2.0 h1:1KNIy1I1H9hNNFEEH3DVnI4UujN+1zjpuk6gwHLTssg=
 github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY=
+github.com/cockroachdb/apd v1.1.0 h1:3LFP3629v+1aKXU5Q37mxmRxX/pIu1nijXydLShEq5I=
 github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ=
 github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
 github.com/coreos/go-systemd v0.0.0-20190719114852-fd7a80b32e1f/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
@@ -46,6 +48,7 @@ github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LB
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
 github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
+github.com/gofrs/uuid v4.0.0+incompatible h1:1SD/1F5pU8p29ybwgQSwpQk+mwdRrXCYuPhW6m+TnJw=
 github.com/gofrs/uuid v4.0.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
 github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
@@ -65,7 +68,6 @@ github.com/gorilla/securecookie v1.1.2/go.mod h1:NfCASbcHqRSY+3a8tlWJwsQap2VX5pw
 github.com/gorilla/sessions v1.2.1/go.mod h1:dk2InVEVJ0sfLlnXv9EAgkf6ecYs/i80K/zI+bUmuGM=
 github.com/gorilla/sessions v1.2.2 h1:lqzMYz6bOfvn2WriPUjNByzeXIlVzURcPmgMczkmTjY=
 github.com/gorilla/sessions v1.2.2/go.mod h1:ePLdVu+jbEgHH+KWw8I1z2wqd0BAdAQh/8LRvBeoNcQ=
-github.com/jackc/chunkreader v1.0.0 h1:4s39bBR8ByfqH+DKm8rQA3E1LHZWB9XWcrz8fqaZbe0=
 github.com/jackc/chunkreader v1.0.0/go.mod h1:RT6O25fNZIuasFJRyZ4R/Y2BbhasbmZXF9QQ7T3kePo=
 github.com/jackc/chunkreader/v2 v2.0.0/go.mod h1:odVSm741yZoC3dpHEUXIqA9tQRhFrgOHwnPIn9lDKlk=
 github.com/jackc/chunkreader/v2 v2.0.1 h1:i+RDz65UE+mmpjTfyz0MoVTnzeYxroil2G82ki7MGG8=
@@ -82,10 +84,10 @@ github.com/jackc/pgio v1.0.0 h1:g12B9UwVnzGhueNavwioyEEpAmqMe1E/BN9ES+8ovkE=
 github.com/jackc/pgio v1.0.0/go.mod h1:oP+2QK2wFfUWgr+gxjoBH9KGBb31Eio69xUb0w5bYf8=
 github.com/jackc/pgmock v0.0.0-20190831213851-13a1b77aafa2/go.mod h1:fGZlG77KXmcq05nJLRkk0+p82V8B8Dw8KN2/V9c/OAE=
 github.com/jackc/pgmock v0.0.0-20201204152224-4fe30f7445fd/go.mod h1:hrBW0Enj2AZTNpt/7Y5rr2xe/9Mn757Wtb2xeBzPv2c=
+github.com/jackc/pgmock v0.0.0-20210724152146-4ad1a8207f65 h1:DadwsjnMwFjfWc9y5Wi/+Zz7xoE5ALHsRQlOctkOiHc=
 github.com/jackc/pgmock v0.0.0-20210724152146-4ad1a8207f65/go.mod h1:5R2h2EEX+qri8jOWMbJCtaPWkrrNc7OHwsp2TCqp7ak=
 github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
 github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
-github.com/jackc/pgproto3 v1.1.0 h1:FYYE4yRw+AgI8wXIinMlNjBbp/UitDJwfj5LqqewP1A=
 github.com/jackc/pgproto3 v1.1.0/go.mod h1:eR5FA3leWg7p9aeAqi37XOTgTIbkABlvcPB3E5rlc78=
 github.com/jackc/pgproto3/v2 v2.0.0-alpha1.0.20190420180111-c116219b62db/go.mod h1:bhq50y+xrl9n5mRYyCBFKkpRVTLYJVWeCc+mEAI3yXA=
 github.com/jackc/pgproto3/v2 v2.0.0-alpha1.0.20190609003834-432c2951c711/go.mod h1:uH0AWtUmuShn0bcesswc4aBTWGvw0cAxIJp+6OB//Wg=
@@ -141,6 +143,10 @@ github.com/lib/pq v1.0.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
 github.com/lib/pq v1.1.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
 github.com/lib/pq v1.2.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
 github.com/lib/pq v1.10.2/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
+github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw=
+github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
+github.com/liushuangls/go-anthropic v1.6.0 h1:8hDEn/EJkeerOFwnQ10efTlRIU6VO+IxE6u6IinphBg=
+github.com/liushuangls/go-anthropic v1.6.0/go.mod h1:sUg9f/ZHoia6Nc8zoNvT7+KavHwMq2eL3VY1Mgf6I7Y=
 github.com/mattn/go-colorable v0.1.1/go.mod h1:FuOcm+DKB9mbwrcAfNl7/TZVBZ6rcnceauSikq3lYCQ=
 github.com/mattn/go-colorable v0.1.6/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
 github.com/mattn/go-isatty v0.0.5/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s=
@@ -159,6 +165,7 @@ github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjY
 github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM=
 github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
@@ -175,6 +182,7 @@ github.com/sashabaranov/go-openai v1.28.0 h1:WS9F9BriSvtHvknPQy2Oi3b+8zkmJdEXcyc
 github.com/sashabaranov/go-openai v1.28.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/satori/go.uuid v1.2.0/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0=
 github.com/shopspring/decimal v0.0.0-20180709203117-cd690d0c9e24/go.mod h1:M+9NzErvs504Cn4c5DxATwIqPbtswREoFCre64PpcG4=
+github.com/shopspring/decimal v1.2.0 h1:abSATXmQEYyShuxI4/vyW3tV1MrKAJzCZ/0zLUXYbsQ=
 github.com/shopspring/decimal v1.2.0/go.mod h1:DKyhrW/HYNuLGql+MJL6WCR6knT2jwCFRcu2hWCYk4o=
 github.com/sirupsen/logrus v1.4.1/go.mod h1:ni0Sbl8bgC9z8RoU9G6nDWqqs/fq4eDPysMBDgk/93Q=
 github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=

openai.go

@@ -2,6 +2,7 @@ package flightlesssomething
 
 import (
 	"context"
+	"errors"
 	"log"
 	"math"
 	"sort"
@@ -9,6 +10,7 @@ import (
 	"strings"
 	"sync"
 
+	anthro "github.com/liushuangls/go-anthropic"
 	openai "github.com/sashabaranov/go-openai"
 )
 
@@ -57,7 +59,7 @@ var (
 
 func generateSummary(b *Benchmark, bds []*BenchmarkData) {
 	// Check if OpenAI integration is not enabled
-	if openaiClient == nil {
+	if openaiClient == nil && anthroClient == nil {
 		return
 	}
 
@@ -76,23 +78,13 @@ func generateSummary(b *Benchmark, bds []*BenchmarkData) {
 	userPrompt := writeAIPrompt(bds, b.Title, b.Description)
 
 	// Retrieve AI response
-	resp, err := openaiClient.CreateChatCompletion(
+	resp, err := getAIResponse(systemMessage, userPrompt)
-		context.Background(),
-		openai.ChatCompletionRequest{
-			Model:       openaiModel,
-			Temperature: 0.0,
-			Messages: []openai.ChatCompletionMessage{
-				{Role: openai.ChatMessageRoleSystem, Content: systemMessage},
-				{Role: openai.ChatMessageRoleUser, Content: userPrompt},
-			},
-		},
-	)
 	if err != nil {
 		log.Println("Failed to generate AI summary:", err)
 		return
 	}
 
-	db.Model(&Benchmark{}).Where("id = ?", b.ID).Update("AiSummary", resp.Choices[0].Message.Content)
+	db.Model(&Benchmark{}).Where("id = ?", b.ID).Update("AiSummary", resp)
 
 	// Update status
 	inProgressSummariesMux.Lock()
@@ -100,6 +92,40 @@ func generateSummary(b *Benchmark, bds []*BenchmarkData) {
 	inProgressSummariesMux.Unlock()
 }
 
+func getAIResponse(systemMessage string, userPrompt string) (string, error) {
+	if openaiClient != nil {
+		resp, err := openaiClient.CreateChatCompletion(
+			context.Background(),
+			openai.ChatCompletionRequest{
+				Model:       openaiModel,
+				Temperature: 0.0,
+				Messages: []openai.ChatCompletionMessage{
+					{Role: openai.ChatMessageRoleSystem, Content: systemMessage},
+					{Role: openai.ChatMessageRoleUser, Content: userPrompt},
+				},
+			},
+		)
+		if err != nil {
+			return "", err
+		}
+		return resp.Choices[0].Message.Content, nil
+	} else if anthroClient != nil {
+		resp, err := anthroClient.CreateMessages(context.Background(), anthro.MessagesRequest{
+			Model:     anthroModel,
+			System:    systemMessage,
+			MaxTokens: 1024,
+			Messages: []anthro.Message{
+				anthro.NewUserTextMessage(userPrompt),
+			},
+		})
+		if err != nil {
+			return "", err
+		}
+		return resp.Content[0].Text, nil
+	}
+	return "", errors.New("no AI client configured")
+}
+
 func writeAIPrompt(bds []*BenchmarkData, bdTitle, bdDescription string) string {
 	sb := strings.Builder{}
 	sb.WriteString("Benchmark title: ")

server.go

@@ -14,9 +14,11 @@ import (
 	gormsessions "github.com/gin-contrib/sessions/gorm"
 	"github.com/gin-gonic/gin"
 	"github.com/glebarez/sqlite"
+	anthro "github.com/liushuangls/go-anthropic"
 	"github.com/ravener/discord-oauth2"
 	openai "github.com/sashabaranov/go-openai"
 	"golang.org/x/oauth2"
+	"gorm.io/driver/postgres"
 	"gorm.io/gorm"
 )
 
@@ -33,6 +35,9 @@ var (
 	// OpenAI
 	openaiClient *openai.Client
 	openaiModel  string
+
+	anthroClient *anthro.Client
+	anthroModel  string
 )
 
 func Start(c *Config, version string) {
@@ -45,6 +50,13 @@ func Start(c *Config, version string) {
 		openaiModel = c.OpenAIModel
 	}
 
+	// Setup Anthro client //
+
+	if c.AnthroAPIKey != "" {
+		anthroClient = anthro.NewClient(c.AnthroAPIKey)
+		anthroModel = c.AnthroModel
+	}
+
 	// Setup data dir //
 
 	_, err := os.Stat(c.DataDir)
@@ -80,7 +92,14 @@ func Start(c *Config, version string) {
 
 	// Setup gorm (database) //
 
-	db, err = gorm.Open(sqlite.Open(filepath.Join(c.DataDir, "database.db")), &gorm.Config{})
+	var gormDialector gorm.Dialector
+	if c.DBurl != "" {
+		gormDialector = postgres.Open(c.DBurl)
+	} else {
+		gormDialector = sqlite.Open(filepath.Join(c.DataDir, "database.db"))
+	}
+
+	db, err = gorm.Open(gormDialector, &gorm.Config{})
 	if err != nil {
 		panic(err)
 	}