Add postgres support and anthropic ai support

This commit is contained in:
ferreo 2024-08-21 22:52:39 +01:00
parent 1b3bbed1c7
commit 9add1599a4
7 changed files with 98 additions and 18 deletions

View File

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

View File

@ -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
}

14
go.mod
View File

@ -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
)

12
go.sum
View File

@ -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=

BIN
mango Executable file

Binary file not shown.

View File

@ -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,6 +78,22 @@ func generateSummary(b *Benchmark, bds []*BenchmarkData) {
userPrompt := writeAIPrompt(bds, b.Title, b.Description)
// Retrieve AI response
resp, err := getAIResponse(systemMessage, userPrompt)
if err != nil {
log.Println("Failed to generate AI summary:", err)
return
}
db.Model(&Benchmark{}).Where("id = ?", b.ID).Update("AiSummary", resp)
// Update status
inProgressSummariesMux.Lock()
delete(inProgressSummaries, b.ID)
inProgressSummariesMux.Unlock()
}
func getAIResponse(systemMessage string, userPrompt string) (string, error) {
if openaiClient != nil {
resp, err := openaiClient.CreateChatCompletion(
context.Background(),
openai.ChatCompletionRequest{
@ -88,16 +106,24 @@ func generateSummary(b *Benchmark, bds []*BenchmarkData) {
},
)
if err != nil {
log.Println("Failed to generate AI summary:", err)
return
return "", err
}
db.Model(&Benchmark{}).Where("id = ?", b.ID).Update("AiSummary", resp.Choices[0].Message.Content)
// Update status
inProgressSummariesMux.Lock()
delete(inProgressSummaries, b.ID)
inProgressSummariesMux.Unlock()
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 {

View File

@ -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)
}