Skip to content

Commit 10c474a

Browse files
ArdemArdem
authored
Added API integration with Gemini (#3)
1 parent 3ae77b1 commit 10c474a

9 files changed

Lines changed: 232 additions & 41 deletions

File tree

‎backend/integrations/services/sql/RedshiftIntegrationConnection.go‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"context"
55
"database/sql"
66
"fmt"
7-
_ "github.com/alexbrainman/odbc"
87
)
98

109
type RedshiftIntegrationConnection struct {

‎processor/taskrunner/.env‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,3 +11,11 @@ TASK_RUNNER_TABLEAU_HTTP_WRAPPER_BASE_URL=http://host.docker.internal:7890
1111
# More information: https://docs.ylem.co/pipelines/tasks-ip/gpt
1212
TASK_RUNNER_OPENAI_GPT_KEY=
1313
TASK_RUNNER_OPENAI_MODEL=gpt-4o-mini
14+
15+
# To enable Gemini integration, create its API secret key
16+
# And place it here
17+
TASK_RUNNER_GEMINI_KEY=
18+
TASK_RUNNER_GEMINI_MODEL=gemini-2.0-flash
19+
20+
# To tell Ylem which AI provider it should use, write it here. Possible values: openai | gemini
21+
TASK_RUNNER_AI_PROVIDER=openai

‎processor/taskrunner/config/config.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ type config struct {
2626
Gopyk gopyk
2727
YlemStatistics ylemStatistics
2828
Openai openai
29+
Gemini gemini
30+
AIProvider string `default:"openai"`
2931
}
3032

3133
type taskRunner struct {
@@ -82,6 +84,11 @@ type openai struct {
8284
Model string `split_words:"true" default:"gpt-4o-mini"`
8385
}
8486

87+
type gemini struct {
88+
Key string `split_words:"true"`
89+
Model string `split_words:"true" default:"gemini-2.0-flash"`
90+
}
91+
8592
var c config
8693

8794
func init() {

‎processor/taskrunner/domain/runner/gpt.go‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@ package runner
22

33
import (
44
"fmt"
5+
"ylem_taskrunner/config"
56
"ylem_taskrunner/services/openai"
7+
"ylem_taskrunner/services/gemini"
68

79
messaging "github.com/ylem-co/shared-messaging"
810
)
@@ -39,11 +41,17 @@ func GptTaskRunner(t *messaging.CallOpenapiGptTask) *messaging.TaskRunResult {
3941
return tr
4042
}
4143

42-
inst := openai.Instance()
43-
resp, err := inst.CompleteText(openai.Completion{
44-
JSON: string(t.Input),
45-
UserPrompt: t.Prompt,
46-
})
44+
var err error
45+
var resp string
46+
if config.Cfg().AIProvider == "openai" {
47+
inst := openai.Instance()
48+
resp, err = inst.CompleteText(openai.Completion{
49+
JSON: string(t.Input),
50+
UserPrompt: t.Prompt,
51+
})
52+
} else {
53+
resp, err = gemini.Process(string(t.Input), t.Prompt)
54+
}
4755

4856
if err != nil {
4957
tr.IsSuccessful = false

‎processor/taskrunner/go.mod‎

Lines changed: 54 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ module ylem_taskrunner
33
go 1.23
44

55
require (
6-
cloud.google.com/go/bigquery v1.27.0
6+
cloud.google.com/go/bigquery v1.62.0
77
github.com/IBM/sarama v1.42.1
88
github.com/PaesslerAG/gval v1.1.2
99
github.com/andygrunwald/go-jira v1.15.1
@@ -15,7 +15,7 @@ require (
1515
github.com/go-resty/resty/v2 v2.7.0
1616
github.com/go-sql-driver/mysql v1.6.0
1717
github.com/golang/gddo v0.0.0-20210115222349-20d68f94ee1f
18-
github.com/google/uuid v1.3.0
18+
github.com/google/uuid v1.6.0
1919
github.com/gorilla/mux v1.8.0
2020
github.com/hashicorp/golang-lru v0.5.1
2121
github.com/joho/godotenv v1.4.0
@@ -28,31 +28,53 @@ require (
2828
github.com/slack-go/slack v0.10.0
2929
github.com/snowflakedb/gosnowflake v1.6.5
3030
github.com/steinfletcher/apitest v1.5.11
31-
github.com/tidwall/gjson v1.12.1
31+
github.com/tidwall/gjson v1.14.2
3232
github.com/twilio/twilio-go v1.23.12
3333
github.com/urfave/cli/v2 v2.3.0
34-
golang.org/x/crypto v0.14.0
35-
golang.org/x/oauth2 v0.2.0
36-
golang.org/x/time v0.5.0
37-
google.golang.org/api v0.97.0
34+
github.com/ylem-co/es-sql-client v0.0.0-20240902163915-13ef8638db5d
35+
github.com/ylem-co/hubspot-client v0.0.0-20240902164341-00f6a99cfdd1
36+
github.com/ylem-co/opsgenie-client v0.0.0-20240902173345-eb25eabd0214
37+
github.com/ylem-co/salesforce-client v0.0.0-20240902165912-6a570283e8d5
38+
github.com/ylem-co/shared-messaging v0.0.0-20250214095426-8c2750adeae1
39+
golang.org/x/crypto v0.27.0
40+
golang.org/x/oauth2 v0.23.0
41+
golang.org/x/time v0.6.0
42+
google.golang.org/api v0.197.0
43+
google.golang.org/genai v1.3.0
3844
gopkg.in/gomail.v2 v2.0.0-20160411212932-81ebce5c23df
3945
)
4046

4147
require (
48+
cloud.google.com/go/auth v0.9.3 // indirect
49+
cloud.google.com/go/auth/oauth2adapt v0.2.4 // indirect
50+
cloud.google.com/go/compute/metadata v0.5.0 // indirect
4251
github.com/alexbrainman/odbc v0.0.0-20240810052813-bcbcb6842ce9 // indirect
52+
github.com/apache/arrow/go/v15 v15.0.2 // indirect
53+
github.com/felixge/httpsnoop v1.0.4 // indirect
54+
github.com/go-logr/logr v1.4.2 // indirect
55+
github.com/go-logr/stdr v1.2.2 // indirect
56+
github.com/goccy/go-json v0.10.2 // indirect
57+
github.com/google/s2a-go v0.1.8 // indirect
4358
github.com/itchyny/timefmt-go v0.1.5 // indirect
44-
github.com/ylem-co/es-sql-client v0.0.0-20240902163915-13ef8638db5d // indirect
45-
github.com/ylem-co/hubspot-client v0.0.0-20240902164341-00f6a99cfdd1 // indirect
46-
github.com/ylem-co/opsgenie-client v0.0.0-20240902173345-eb25eabd0214 // indirect
47-
github.com/ylem-co/salesforce-client v0.0.0-20240902165912-6a570283e8d5 // indirect
48-
github.com/ylem-co/shared-messaging v0.0.0-20250214095426-8c2750adeae1 // indirect
59+
github.com/klauspost/cpuid/v2 v2.2.5 // indirect
60+
github.com/zeebo/xxh3 v1.0.2 // indirect
61+
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.54.0 // indirect
62+
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.54.0 // indirect
63+
go.opentelemetry.io/otel v1.29.0 // indirect
64+
go.opentelemetry.io/otel/metric v1.29.0 // indirect
65+
go.opentelemetry.io/otel/trace v1.29.0 // indirect
66+
golang.org/x/exp v0.0.0-20231006140011-7918f672742d // indirect
67+
golang.org/x/mod v0.18.0 // indirect
4968
golang.org/x/sys v0.30.0 // indirect
69+
golang.org/x/tools v0.22.0 // indirect
70+
google.golang.org/genproto/googleapis/api v0.0.0-20240903143218-8af14fe29dc1 // indirect
71+
google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1 // indirect
5072
)
5173

5274
require (
53-
cloud.google.com/go v0.102.0 // indirect
54-
cloud.google.com/go/compute v1.7.0 // indirect
55-
cloud.google.com/go/iam v0.3.0 // indirect
75+
cloud.google.com/go v0.116.0 // indirect
76+
cloud.google.com/go/compute v1.28.0 // indirect
77+
cloud.google.com/go/iam v1.2.0 // indirect
5678
github.com/Azure/azure-pipeline-go v0.2.3 // indirect
5779
github.com/Azure/azure-storage-blob-go v0.14.0 // indirect
5880
github.com/apache/arrow/go/arrow v0.0.0-20211112161151-bc219186db40 // indirect
@@ -71,7 +93,7 @@ require (
7193
github.com/aws/aws-sdk-go-v2/service/sso v1.6.0 // indirect
7294
github.com/aws/aws-sdk-go-v2/service/sts v1.10.0 // indirect
7395
github.com/aws/smithy-go v1.11.2 // indirect
74-
github.com/cespare/xxhash/v2 v2.1.2 // indirect
96+
github.com/cespare/xxhash/v2 v2.3.0 // indirect
7597
github.com/cpuguy83/go-md2man/v2 v2.0.0 // indirect
7698
github.com/davecgh/go-spew v1.1.1 // indirect
7799
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
@@ -83,16 +105,16 @@ require (
83105
github.com/gabriel-vasile/mimetype v1.4.0 // indirect
84106
github.com/go-stack/stack v1.8.1 // indirect
85107
github.com/golang-jwt/jwt/v4 v4.3.0 // indirect
86-
github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e // indirect
108+
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
87109
github.com/golang/mock v1.6.0 // indirect
88-
github.com/golang/protobuf v1.5.2 // indirect
110+
github.com/golang/protobuf v1.5.4 // indirect
89111
github.com/golang/snappy v0.0.4 // indirect
90-
github.com/google/flatbuffers v2.0.0+incompatible // indirect
91-
github.com/google/go-cmp v0.5.8 // indirect
112+
github.com/google/flatbuffers v23.5.26+incompatible // indirect
113+
github.com/google/go-cmp v0.6.0 // indirect
92114
github.com/google/go-querystring v1.1.0 // indirect
93-
github.com/googleapis/enterprise-certificate-proxy v0.1.0 // indirect
94-
github.com/googleapis/gax-go/v2 v2.4.0 // indirect
95-
github.com/gorilla/websocket v1.4.2 // indirect
115+
github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
116+
github.com/googleapis/gax-go/v2 v2.13.0 // indirect
117+
github.com/gorilla/websocket v1.5.3 // indirect
96118
github.com/hashicorp/errwrap v1.1.0 // indirect
97119
github.com/hashicorp/go-multierror v1.1.1 // indirect
98120
github.com/hashicorp/go-uuid v1.0.3 // indirect
@@ -115,14 +137,14 @@ require (
115137
github.com/tidwall/match v1.1.1 // indirect
116138
github.com/tidwall/pretty v1.2.0 // indirect
117139
github.com/trivago/tgo v1.0.7 // indirect
118-
go.opencensus.io v0.23.0 // indirect
119-
golang.org/x/net v0.17.0 // indirect
120-
golang.org/x/sync v0.4.0 // indirect
121-
golang.org/x/text v0.13.0 // indirect
122-
golang.org/x/xerrors v0.0.0-20220609144429-65e65417b02f // indirect
123-
google.golang.org/appengine v1.6.7 // indirect
124-
google.golang.org/genproto v0.0.0-20220624142145-8cd45d7dbd1f // indirect
125-
google.golang.org/grpc v1.47.0 // indirect
126-
google.golang.org/protobuf v1.28.1 // indirect
140+
go.opencensus.io v0.24.0 // indirect
141+
golang.org/x/net v0.29.0 // indirect
142+
golang.org/x/sync v0.8.0 // indirect
143+
golang.org/x/text v0.18.0 // indirect
144+
golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028 // indirect
145+
google.golang.org/appengine v1.6.8 // indirect
146+
google.golang.org/genproto v0.0.0-20240903143218-8af14fe29dc1 // indirect
147+
google.golang.org/grpc v1.66.2 // indirect
148+
google.golang.org/protobuf v1.34.2 // indirect
127149
gopkg.in/alexcesaro/quotedprintable.v3 v3.0.0-20150716171945-2caba252f4dc // indirect
128150
)

0 commit comments

Comments
 (0)