Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ import (
"context"
"database/sql"
"fmt"
_ "github.com/alexbrainman/odbc"
)

type RedshiftIntegrationConnection struct {
Expand Down
8 changes: 8 additions & 0 deletions processor/taskrunner/.env
Original file line number Diff line number Diff line change
Expand Up @@ -11,3 +11,11 @@ TASK_RUNNER_TABLEAU_HTTP_WRAPPER_BASE_URL=http://host.docker.internal:7890
# More information: https://docs.ylem.co/pipelines/tasks-ip/gpt
TASK_RUNNER_OPENAI_GPT_KEY=
TASK_RUNNER_OPENAI_MODEL=gpt-4o-mini

# To enable Gemini integration, create its API secret key
# And place it here
TASK_RUNNER_GEMINI_KEY=
TASK_RUNNER_GEMINI_MODEL=gemini-2.0-flash

# To tell Ylem which AI provider it should use, write it here. Possible values: openai | gemini
TASK_RUNNER_AI_PROVIDER=openai
7 changes: 7 additions & 0 deletions processor/taskrunner/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ type config struct {
Gopyk gopyk
YlemStatistics ylemStatistics
Openai openai
Gemini gemini
AIProvider string `default:"openai"`
}

type taskRunner struct {
Expand Down Expand Up @@ -82,6 +84,11 @@ type openai struct {
Model string `split_words:"true" default:"gpt-4o-mini"`
}

type gemini struct {
Key string `split_words:"true"`
Model string `split_words:"true" default:"gemini-2.0-flash"`
}

var c config

func init() {
Expand Down
18 changes: 13 additions & 5 deletions processor/taskrunner/domain/runner/gpt.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,9 @@ package runner

import (
"fmt"
"ylem_taskrunner/config"
"ylem_taskrunner/services/openai"
"ylem_taskrunner/services/gemini"

messaging "github.com/ylem-co/shared-messaging"
)
Expand Down Expand Up @@ -39,11 +41,17 @@ func GptTaskRunner(t *messaging.CallOpenapiGptTask) *messaging.TaskRunResult {
return tr
}

inst := openai.Instance()
resp, err := inst.CompleteText(openai.Completion{
JSON: string(t.Input),
UserPrompt: t.Prompt,
})
var err error
var resp string
if config.Cfg().AIProvider == "openai" {
inst := openai.Instance()
resp, err = inst.CompleteText(openai.Completion{
JSON: string(t.Input),
UserPrompt: t.Prompt,
})
} else {
resp, err = gemini.Process(string(t.Input), t.Prompt)
}

if err != nil {
tr.IsSuccessful = false
Expand Down
86 changes: 54 additions & 32 deletions processor/taskrunner/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ module ylem_taskrunner
go 1.23

require (
cloud.google.com/go/bigquery v1.27.0
cloud.google.com/go/bigquery v1.62.0
github.com/IBM/sarama v1.42.1
github.com/PaesslerAG/gval v1.1.2
github.com/andygrunwald/go-jira v1.15.1
Expand All @@ -15,7 +15,7 @@ require (
github.com/go-resty/resty/v2 v2.7.0
github.com/go-sql-driver/mysql v1.6.0
github.com/golang/gddo v0.0.0-20210115222349-20d68f94ee1f
github.com/google/uuid v1.3.0
github.com/google/uuid v1.6.0
github.com/gorilla/mux v1.8.0
github.com/hashicorp/golang-lru v0.5.1
github.com/joho/godotenv v1.4.0
Expand All @@ -28,31 +28,53 @@ require (
github.com/slack-go/slack v0.10.0
github.com/snowflakedb/gosnowflake v1.6.5
github.com/steinfletcher/apitest v1.5.11
github.com/tidwall/gjson v1.12.1
github.com/tidwall/gjson v1.14.2
github.com/twilio/twilio-go v1.23.12
github.com/urfave/cli/v2 v2.3.0
golang.org/x/crypto v0.14.0
golang.org/x/oauth2 v0.2.0
golang.org/x/time v0.5.0
google.golang.org/api v0.97.0
github.com/ylem-co/es-sql-client v0.0.0-20240902163915-13ef8638db5d
github.com/ylem-co/hubspot-client v0.0.0-20240902164341-00f6a99cfdd1
github.com/ylem-co/opsgenie-client v0.0.0-20240902173345-eb25eabd0214
github.com/ylem-co/salesforce-client v0.0.0-20240902165912-6a570283e8d5
github.com/ylem-co/shared-messaging v0.0.0-20250214095426-8c2750adeae1
golang.org/x/crypto v0.27.0
golang.org/x/oauth2 v0.23.0
golang.org/x/time v0.6.0
google.golang.org/api v0.197.0
google.golang.org/genai v1.3.0
gopkg.in/gomail.v2 v2.0.0-20160411212932-81ebce5c23df
)

require (
cloud.google.com/go/auth v0.9.3 // indirect
cloud.google.com/go/auth/oauth2adapt v0.2.4 // indirect
cloud.google.com/go/compute/metadata v0.5.0 // indirect
github.com/alexbrainman/odbc v0.0.0-20240810052813-bcbcb6842ce9 // indirect
github.com/apache/arrow/go/v15 v15.0.2 // indirect
github.com/felixge/httpsnoop v1.0.4 // indirect
github.com/go-logr/logr v1.4.2 // indirect
github.com/go-logr/stdr v1.2.2 // indirect
github.com/goccy/go-json v0.10.2 // indirect
github.com/google/s2a-go v0.1.8 // indirect
github.com/itchyny/timefmt-go v0.1.5 // indirect
github.com/ylem-co/es-sql-client v0.0.0-20240902163915-13ef8638db5d // indirect
github.com/ylem-co/hubspot-client v0.0.0-20240902164341-00f6a99cfdd1 // indirect
github.com/ylem-co/opsgenie-client v0.0.0-20240902173345-eb25eabd0214 // indirect
github.com/ylem-co/salesforce-client v0.0.0-20240902165912-6a570283e8d5 // indirect
github.com/ylem-co/shared-messaging v0.0.0-20250214095426-8c2750adeae1 // indirect
github.com/klauspost/cpuid/v2 v2.2.5 // indirect
github.com/zeebo/xxh3 v1.0.2 // indirect
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.54.0 // indirect
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.54.0 // indirect
go.opentelemetry.io/otel v1.29.0 // indirect
go.opentelemetry.io/otel/metric v1.29.0 // indirect
go.opentelemetry.io/otel/trace v1.29.0 // indirect
golang.org/x/exp v0.0.0-20231006140011-7918f672742d // indirect
golang.org/x/mod v0.18.0 // indirect
golang.org/x/sys v0.30.0 // indirect
golang.org/x/tools v0.22.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20240903143218-8af14fe29dc1 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1 // indirect
)

require (
cloud.google.com/go v0.102.0 // indirect
cloud.google.com/go/compute v1.7.0 // indirect
cloud.google.com/go/iam v0.3.0 // indirect
cloud.google.com/go v0.116.0 // indirect
cloud.google.com/go/compute v1.28.0 // indirect
cloud.google.com/go/iam v1.2.0 // indirect
github.com/Azure/azure-pipeline-go v0.2.3 // indirect
github.com/Azure/azure-storage-blob-go v0.14.0 // indirect
github.com/apache/arrow/go/arrow v0.0.0-20211112161151-bc219186db40 // indirect
Expand All @@ -71,7 +93,7 @@ require (
github.com/aws/aws-sdk-go-v2/service/sso v1.6.0 // indirect
github.com/aws/aws-sdk-go-v2/service/sts v1.10.0 // indirect
github.com/aws/smithy-go v1.11.2 // indirect
github.com/cespare/xxhash/v2 v2.1.2 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/cpuguy83/go-md2man/v2 v2.0.0 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
Expand All @@ -83,16 +105,16 @@ require (
github.com/gabriel-vasile/mimetype v1.4.0 // indirect
github.com/go-stack/stack v1.8.1 // indirect
github.com/golang-jwt/jwt/v4 v4.3.0 // indirect
github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e // indirect
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
github.com/golang/mock v1.6.0 // indirect
github.com/golang/protobuf v1.5.2 // indirect
github.com/golang/protobuf v1.5.4 // indirect
github.com/golang/snappy v0.0.4 // indirect
github.com/google/flatbuffers v2.0.0+incompatible // indirect
github.com/google/go-cmp v0.5.8 // indirect
github.com/google/flatbuffers v23.5.26+incompatible // indirect
github.com/google/go-cmp v0.6.0 // indirect
github.com/google/go-querystring v1.1.0 // indirect
github.com/googleapis/enterprise-certificate-proxy v0.1.0 // indirect
github.com/googleapis/gax-go/v2 v2.4.0 // indirect
github.com/gorilla/websocket v1.4.2 // indirect
github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
github.com/googleapis/gax-go/v2 v2.13.0 // indirect
github.com/gorilla/websocket v1.5.3 // indirect
github.com/hashicorp/errwrap v1.1.0 // indirect
github.com/hashicorp/go-multierror v1.1.1 // indirect
github.com/hashicorp/go-uuid v1.0.3 // indirect
Expand All @@ -115,14 +137,14 @@ require (
github.com/tidwall/match v1.1.1 // indirect
github.com/tidwall/pretty v1.2.0 // indirect
github.com/trivago/tgo v1.0.7 // indirect
go.opencensus.io v0.23.0 // indirect
golang.org/x/net v0.17.0 // indirect
golang.org/x/sync v0.4.0 // indirect
golang.org/x/text v0.13.0 // indirect
golang.org/x/xerrors v0.0.0-20220609144429-65e65417b02f // indirect
google.golang.org/appengine v1.6.7 // indirect
google.golang.org/genproto v0.0.0-20220624142145-8cd45d7dbd1f // indirect
google.golang.org/grpc v1.47.0 // indirect
google.golang.org/protobuf v1.28.1 // indirect
go.opencensus.io v0.24.0 // indirect
golang.org/x/net v0.29.0 // indirect
golang.org/x/sync v0.8.0 // indirect
golang.org/x/text v0.18.0 // indirect
golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028 // indirect
google.golang.org/appengine v1.6.8 // indirect
google.golang.org/genproto v0.0.0-20240903143218-8af14fe29dc1 // indirect
google.golang.org/grpc v1.66.2 // indirect
google.golang.org/protobuf v1.34.2 // indirect
gopkg.in/alexcesaro/quotedprintable.v3 v3.0.0-20150716171945-2caba252f4dc // indirect
)
Loading
Loading