Skip to content

Commit e54ea71

Browse files
fanweixiaofanweixiao
authored
feat(ai-bridge): support Github Models (#885)
1 parent 8eedd7d commit e54ea71

6 files changed

Lines changed: 173 additions & 5 deletions

File tree

‎cli/serve.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import (
3333
"github.com/yomorun/yomo/pkg/bridge/ai/provider/cfazure"
3434
"github.com/yomorun/yomo/pkg/bridge/ai/provider/cfopenai"
3535
"github.com/yomorun/yomo/pkg/bridge/ai/provider/gemini"
36+
"github.com/yomorun/yomo/pkg/bridge/ai/provider/githubmodels"
3637
"github.com/yomorun/yomo/pkg/bridge/ai/provider/ollama"
3738
"github.com/yomorun/yomo/pkg/bridge/ai/provider/openai"
3839
)
@@ -49,7 +50,7 @@ var serveCmd = &cobra.Command{
4950
}
5051

5152
// log.InfoStatusEvent(os.Stdout, "")
52-
ylog.Info("Starting Zipper...")
53+
ylog.Info("Starting YoMo Zipper...")
5354
// config
5455
conf, err := pkgconfig.ParseConfigFile(config)
5556
if err != nil {
@@ -150,12 +151,14 @@ func registerAIProvider(aiConfig *ai.Config) error {
150151
providerpkg.RegisterProvider(ollama.NewProvider(provider["api_endpoint"]))
151152
case "gemini":
152153
providerpkg.RegisterProvider(gemini.NewProvider(provider["api_key"]))
154+
case "githubmodels":
155+
providerpkg.RegisterProvider(githubmodels.NewProvider(provider["api_key"], provider["model"]))
153156
default:
154157
log.WarningStatusEvent(os.Stdout, "unknown provider: %s", name)
155158
}
156159
}
157160

158-
ylog.Info("registered AI providers", "len", len(providerpkg.ListProviders()))
161+
ylog.Info("register LLM providers", "num", len(providerpkg.ListProviders()))
159162
return nil
160163
}
161164

‎pkg/bridge/ai/ai.go‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,6 @@ func ParseConfig(conf map[string]any) (config *Config, err error) {
144144
if config.Server.Addr == "" {
145145
config.Server.Addr = ":8000"
146146
}
147-
ylog.Info("parse AI config success")
148147
return
149148
}
150149

‎pkg/bridge/ai/api_server.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ func Serve(config *Config, zipperListenAddr string, credential string, logger *s
4949
return err
5050
}
5151

52-
logger.Info("start bridge server", "addr", config.Server.Addr, "provider", provider.Name())
52+
logger.Info("start AI Bridge service", "addr", config.Server.Addr, "provider", provider.Name())
5353
return srv.ServeAddr(config.Server.Addr)
5454
}
5555

‎pkg/bridge/ai/provider/githubmodels/provider.go‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
// Package githubmodels is the Github Models llm provider, see https://github.com/marketplace/models
2+
package githubmodels
3+
4+
import (
5+
"context"
6+
"os"
7+
8+
// automatically load .env file
9+
_ "github.com/joho/godotenv/autoload"
10+
"github.com/sashabaranov/go-openai"
11+
"github.com/yomorun/yomo/core/metadata"
12+
13+
provider "github.com/yomorun/yomo/pkg/bridge/ai/provider"
14+
)
15+
16+
// Provider is the provider for Github Models
17+
type Provider struct {
18+
// APIKey is the API key for Github Models
19+
APIKey string
20+
// Model is the model for Github Models, see https://github.com/marketplace/models
21+
// e.g. "Meta-Llama-3.1-405B-Instruct", "Mistral-large-2407", "gpt-4o"
22+
Model string
23+
client *openai.Client
24+
}
25+
26+
// check if implements ai.Provider
27+
var _ provider.LLMProvider = &Provider{}
28+
29+
// NewProvider creates a new OpenAIProvider
30+
func NewProvider(apiKey string, model string) *Provider {
31+
if apiKey == "" {
32+
apiKey = os.Getenv("GITHUB_TOKEN")
33+
}
34+
35+
config := openai.DefaultConfig(apiKey)
36+
config.BaseURL = "https://models.inference.ai.azure.com"
37+
38+
return &Provider{
39+
APIKey: apiKey,
40+
Model: model,
41+
client: openai.NewClientWithConfig(config),
42+
}
43+
}
44+
45+
// Name returns the name of the provider
46+
func (p *Provider) Name() string {
47+
return "githubmodels"
48+
}
49+
50+
// GetChatCompletions implements ai.LLMProvider.
51+
func (p *Provider) GetChatCompletions(ctx context.Context, req openai.ChatCompletionRequest, _ metadata.M) (openai.ChatCompletionResponse, error) {
52+
if p.Model != "" {
53+
req.Model = p.Model
54+
}
55+
56+
return p.client.CreateChatCompletion(ctx, req)
57+
}
58+
59+
// GetChatCompletionsStream implements ai.LLMProvider.
60+
func (p *Provider) GetChatCompletionsStream(ctx context.Context, req openai.ChatCompletionRequest, _ metadata.M) (provider.ResponseRecver, error) {
61+
if p.Model != "" {
62+
req.Model = p.Model
63+
}
64+
65+
return p.client.CreateChatCompletionStream(ctx, req)
66+
}

‎pkg/bridge/ai/provider/githubmodels/provider_test.go‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
package githubmodels
2+
3+
import (
4+
"context"
5+
"os"
6+
"testing"
7+
8+
openai "github.com/sashabaranov/go-openai"
9+
"github.com/stretchr/testify/assert"
10+
)
11+
12+
func TestGithubModelsProvider_Name(t *testing.T) {
13+
provider := &Provider{}
14+
15+
name := provider.Name()
16+
17+
assert.Equal(t, "githubmodels", name)
18+
}
19+
20+
func TestNewProvider(t *testing.T) {
21+
t.Run("with parameters", func(t *testing.T) {
22+
provider := NewProvider("test_api_key", "test_model")
23+
24+
assert.Equal(t, "test_api_key", provider.APIKey)
25+
assert.Equal(t, "test_model", provider.Model)
26+
})
27+
28+
t.Run("with environment variables", func(t *testing.T) {
29+
os.Setenv("GITHUB_TOKEN", "env_api_key")
30+
31+
provider := NewProvider("", "test_model")
32+
33+
assert.Equal(t, "env_api_key", provider.APIKey)
34+
assert.Equal(t, "test_model", provider.Model)
35+
36+
os.Unsetenv("GITHUB_TOKEN")
37+
})
38+
}
39+
40+
func TestGithubModelsProvider_GetChatCompletions(t *testing.T) {
41+
config := openai.DefaultConfig("test_api_key")
42+
config.BaseURL = "https://models.inference.ai.azure.com"
43+
client := openai.NewClientWithConfig(config)
44+
45+
provider := &Provider{
46+
APIKey: "test_api_key",
47+
Model: "test_model",
48+
client: client,
49+
}
50+
51+
msgs := []openai.ChatCompletionMessage{
52+
{
53+
Role: "user",
54+
Content: "hello",
55+
},
56+
{
57+
Role: "system",
58+
Content: "I'm a bot",
59+
},
60+
}
61+
req := openai.ChatCompletionRequest{
62+
Messages: msgs,
63+
}
64+
65+
_, err := provider.GetChatCompletions(context.TODO(), req, nil)
66+
assert.Error(t, err)
67+
assert.Contains(t, err.Error(), "401")
68+
assert.Contains(t, err.Error(), "Bad credentials")
69+
}
70+
71+
func TestGithubModelsProvider_GetChatCompletionsStream(t *testing.T) {
72+
config := openai.DefaultConfig("test_api_key")
73+
config.BaseURL = "https://models.inference.ai.azure.com"
74+
client := openai.NewClientWithConfig(config)
75+
76+
provider := &Provider{
77+
APIKey: "test_api_key",
78+
Model: "test_model",
79+
client: client,
80+
}
81+
82+
msgs := []openai.ChatCompletionMessage{
83+
{
84+
Role: "user",
85+
Content: "hello",
86+
},
87+
{
88+
Role: "system",
89+
Content: "I'm a bot",
90+
},
91+
}
92+
req := openai.ChatCompletionRequest{
93+
Messages: msgs,
94+
}
95+
96+
_, err := provider.GetChatCompletionsStream(context.TODO(), req, nil)
97+
assert.Error(t, err)
98+
assert.Contains(t, err.Error(), "401")
99+
assert.Contains(t, err.Error(), "Bad credentials")
100+
}

‎zipper_notwindows.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import (
2121
func waitSignalForShutdownServer(server *core.Server) {
2222
c := make(chan os.Signal, 1)
2323
signal.Notify(c, syscall.SIGTERM, syscall.SIGUSR2, syscall.SIGUSR1, syscall.SIGINT)
24-
ylog.Info("Listening SIGUSR1, SIGUSR2, SIGTERM/SIGINT...")
24+
ylog.Info("listening SIGUSR1, SIGUSR2, SIGTERM/SIGINT...")
2525
for p1 := range c {
2626
ylog.Debug("Received signal", "signal", p1)
2727
if p1 == syscall.SIGTERM || p1 == syscall.SIGINT {

0 commit comments

Comments
 (0)