Skip to content

Commit dba9fa1

Browse files
Dot-LiuDot-Liu
authored
Merge pull request #201 from eolinker/feature/local-model
support ollama local model
2 parents a9e3011 + d53b051 commit dba9fa1

File tree

43 files changed

+211
-79
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

43 files changed

+211
-79
lines changed

‎convert/entity.go

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,14 @@ const (
1616

1717
type ClientRequest struct {
1818
Messages []*Message `json:"messages"`
19+
Stream bool `json:"stream"`
1920
}
2021

2122
type ClientResponse struct {
22-
Message Message `json:"message,omitempty"`
23-
FinishReason string `json:"finish_reason,omitempty"`
24-
Code int `json:"code"`
25-
Error string `json:"error"`
23+
Message *Message `json:"message,omitempty"`
24+
FinishReason string `json:"finish_reason"`
25+
Code int `json:"code"`
26+
Error string `json:"error"`
2627
}
2728

2829
type Message struct {

‎convert/manager.go

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,10 @@ func NewBalanceManager() *BalanceManager {
155155
}
156156

157157
func (m *BalanceManager) SetProvider(id string, p IProvider) {
158-
m.providers.Set(p.Provider(), p)
158+
if p.Priority() < 1 {
159+
m.providers.Set(p.Provider(), p)
160+
}
161+
159162
m.balances.Set(id, p)
160163
tmp, has := m.ids.Get(p.Provider())
161164
if !has {
@@ -190,17 +193,24 @@ func (m *BalanceManager) Del(id string) string {
190193
if !ok {
191194
return ""
192195
}
196+
if p.Priority() == 0 {
197+
// 供应商本身
198+
m.providers.Del(p.Provider())
199+
m.ids.Del(p.Provider())
200+
m.sortBalances()
201+
return p.Provider()
202+
}
193203
tmp, has := m.ids.Get(p.Provider())
194204
if !has {
195205
return ""
196206
}
197207
tmp.Del(id)
208+
m.sortBalances()
198209
if tmp.Count() < 1 {
199210
m.providers.Del(p.Provider())
200211
return p.Provider()
201212
}
202213

203-
m.sortBalances()
204214
return ""
205215
}
206216

‎convert/status.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ var (
1010
StatusExceeded = "exceeded"
1111
StatusInvalid = "invalid"
1212
StatusTimeout = "timeout"
13+
14+
ParamAIModel = "param_ai_model"
1315
)
1416

1517
func SetAIStatusTimeout(ctx eocontext.EoContext) {

‎drivers/ai-provider/authropic/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
104104
responseBody := &convert.ClientResponse{}
105105
if len(data.Config.Contents) > 0 {
106106
msg := data.Config.Contents[0]
107-
responseBody.Message = convert.Message{
107+
responseBody.Message = &convert.Message{
108108
Role: data.Config.Role,
109109
Content: msg.Text,
110110
}

‎drivers/ai-provider/baichuan/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
107107
responseBody := &convert.ClientResponse{}
108108
if len(data.Config.Choices) > 0 {
109109
msg := data.Config.Choices[0]
110-
responseBody.Message = convert.Message{
110+
responseBody.Message = &convert.Message{
111111
Role: msg.Message.Role,
112112
Content: msg.Message.Content,
113113
}

‎drivers/ai-provider/bedrock/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
100100
responseBody := &convert.ClientResponse{}
101101
if data.Config.Output.Message != nil && len(data.Config.Output.Message.Content) > 0 {
102102
msg := data.Config.Output.Message
103-
responseBody.Message = convert.Message{
103+
responseBody.Message = &convert.Message{
104104
Role: msg.Role,
105105
Content: msg.Content[0].Text,
106106
}

‎drivers/ai-provider/chatglm/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
101101
responseBody := &convert.ClientResponse{}
102102
if len(data.Config.Choices) > 0 {
103103
msg := data.Config.Choices[0]
104-
responseBody.Message = convert.Message{
104+
responseBody.Message = &convert.Message{
105105
Role: msg.Message.Role,
106106
Content: msg.Message.Content,
107107
}

‎drivers/ai-provider/cohere/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
110110
case map[string]interface{}:
111111
{
112112
responseMessage := convert.MapToStruct[ResponseMessage](tmp)
113-
responseBody.Message = convert.Message{
113+
responseBody.Message = &convert.Message{
114114
Role: responseMessage.Role,
115115
Content: responseMessage.Content[0].Text,
116116
}

‎drivers/ai-provider/deepseek/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
106106
responseBody := &convert.ClientResponse{}
107107
if len(data.Config.Choices) > 0 {
108108
msg := data.Config.Choices[0]
109-
responseBody.Message = convert.Message{
109+
responseBody.Message = &convert.Message{
110110
Role: msg.Message.Role,
111111
Content: msg.Message.Content,
112112
}

‎drivers/ai-provider/fakegpt/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
6565
responseBody := &convert.ClientResponse{}
6666
if len(data.Config.Choices) > 0 {
6767
msg := data.Config.Choices[0]
68-
responseBody.Message = convert.Message{
68+
responseBody.Message = &convert.Message{
6969
Role: msg.Message.Role,
7070
Content: msg.Message.Content,
7171
}

‎drivers/ai-provider/fireworks/mode.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ package fireworks
22

33
import (
44
"encoding/json"
5-
5+
66
"github.com/eolinker/apinto/convert"
77
"github.com/eolinker/eosc"
88
"github.com/eolinker/eosc/eocontext"
@@ -124,7 +124,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
124124
responseBody := &convert.ClientResponse{}
125125
if len(data.Config.Choices) > 0 {
126126
msg := data.Config.Choices[0]
127-
responseBody.Message = convert.Message{
127+
responseBody.Message = &convert.Message{
128128
Role: msg.Message.Role,
129129
Content: msg.Message.Content,
130130
}

‎drivers/ai-provider/google/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
131131
}
132132
}
133133

134-
responseBody.Message = convert.Message{
134+
responseBody.Message = &convert.Message{
135135
Role: role,
136136
Content: text,
137137
}

‎drivers/ai-provider/groq/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
104104
responseBody := &convert.ClientResponse{}
105105
if len(data.Config.Choices) > 0 {
106106
msg := data.Config.Choices[0]
107-
responseBody.Message = convert.Message{
107+
responseBody.Message = &convert.Message{
108108
Role: msg.Message.Role,
109109
Content: msg.Message.Content,
110110
}

‎drivers/ai-provider/hunyuan/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
221221
if len(data.Config.Response.Choices) > 0 {
222222
//if data.Config.Object == "chat.completion" {
223223
msg := data.Config.Response.Choices[0]
224-
responseBody.Message = convert.Message{
224+
responseBody.Message = &convert.Message{
225225
Role: msg.Message.Role,
226226
Content: msg.Message.Content,
227227
}

‎drivers/ai-provider/minimax/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
105105
responseBody := &convert.ClientResponse{}
106106
if len(data.Config.Choices) > 0 {
107107
msg := data.Config.Choices[0]
108-
responseBody.Message = convert.Message{
108+
responseBody.Message = &convert.Message{
109109
Role: msg.Message.Role,
110110
Content: msg.Message.Content,
111111
}

‎drivers/ai-provider/mistralai/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
118118
responseBody := &convert.ClientResponse{}
119119
if len(data.Config.Choices) > 0 {
120120
msg := data.Config.Choices[0]
121-
responseBody.Message = convert.Message{
121+
responseBody.Message = &convert.Message{
122122
Role: msg.Message.Role,
123123
Content: msg.Message.Content,
124124
}

‎drivers/ai-provider/moonshot/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
122122
responseBody := &convert.ClientResponse{}
123123
if len(data.Config.Choices) > 0 {
124124
msg := data.Config.Choices[0]
125-
responseBody.Message = convert.Message{
125+
responseBody.Message = &convert.Message{
126126
Role: msg.Message.Role,
127127
Content: msg.Message.Content,
128128
}

‎drivers/ai-provider/novita/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
114114
responseBody := &convert.ClientResponse{}
115115
if len(data.Config.Choices) > 0 {
116116
msg := data.Config.Choices[0]
117-
responseBody.Message = convert.Message{
117+
responseBody.Message = &convert.Message{
118118
Role: msg.Message.Role,
119119
Content: msg.Message.Content,
120120
}

‎drivers/ai-provider/nvidia/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
112112
responseBody := &convert.ClientResponse{}
113113
if data != nil && len(data.Config.Choices) > 0 {
114114
msg := data.Config.Choices[0]
115-
responseBody.Message = convert.Message{
115+
responseBody.Message = &convert.Message{
116116
Role: msg.Message.Role,
117117
Content: msg.Message.Content,
118118
}

‎drivers/ai-provider/ollama/converter.go

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -45,20 +45,14 @@ func newConverterDriver(cfg *Config) (convert.IConverterDriver, error) {
4545
}
4646

4747
func (e *converterDriver) GetConverter(model string) (convert.IConverter, bool) {
48-
//converter, ok := modelConvert[model]
49-
//if !ok {
50-
// return nil, false
51-
//}
52-
5348
return &Converter{balanceHandler: e.BalanceHandler, converter: NewChat()}, true
5449
}
5550

5651
func (e *converterDriver) GetModel(model string) (convert.FGenerateConfig, bool) {
5752
return func(cfg string) (map[string]interface{}, error) {
5853

5954
result := map[string]interface{}{
60-
"model": model,
61-
"stream": false,
55+
"model": model,
6256
}
6357
if cfg != "" {
6458
tmp := make(map[string]interface{})

‎drivers/ai-provider/ollama/message.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,4 +27,5 @@ type Response struct {
2727
PromptEvalDuration int `json:"prompt_eval_duration"`
2828
EvalCount int `json:"eval_count"`
2929
EvalDuration int64 `json:"eval_duration"`
30+
Error string `json:"error"`
3031
}

‎drivers/ai-provider/ollama/mode.go

Lines changed: 65 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -73,11 +73,17 @@ func (c *Chat) RequestConvert(ctx eocontext.EoContext, extender map[string]inter
7373
for k, v := range extender {
7474
baseCfg.SetAppend(k, v)
7575
}
76+
//if paramModel := ctx.GetLabel(convert.ParamAIModel); paramModel != "" {
77+
// baseCfg.SetAppend("model", paramModel)
78+
//}
7679
// Marshal the updated configuration back into JSON.
7780
body, err = json.Marshal(baseCfg)
7881
if err != nil {
7982
return err
8083
}
84+
if baseCfg.Config.Stream {
85+
httpContext.SetLabel("stream", "true")
86+
}
8187

8288
// SetProvider the modified body in the HTTP context.
8389
httpContext.Proxy().Body().SetRaw("application/json", body)
@@ -93,9 +99,64 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
9399
if err != nil {
94100
return err
95101
}
96-
102+
status := httpContext.Response().StatusCode()
103+
switch status {
104+
case 200:
105+
convert.SetAIStatusNormal(ctx)
106+
}
107+
if httpContext.Response().IsBodyStream() {
108+
bodyStream := httpContext.Response().GetBodyStream()
109+
if bodyStream != nil {
110+
bodyStream.AppendWriterFunc(func(p []byte) ([]byte, error) {
111+
// Parse the response body into a base configuration.
112+
data := eosc.NewBase[Response]()
113+
err = json.Unmarshal(p, data)
114+
if err != nil {
115+
return nil, err
116+
}
117+
switch status {
118+
case 200:
119+
// Calculate the token consumption for a successful request.
120+
usage := data.Config
121+
if usage.Done {
122+
convert.SetAIStatusNormal(ctx)
123+
convert.SetAIModelInputToken(ctx, usage.PromptEvalCount)
124+
convert.SetAIModelOutputToken(ctx, usage.EvalCount)
125+
convert.SetAIModelTotalToken(ctx, usage.PromptEvalCount+usage.EvalCount)
126+
}
127+
}
128+
129+
// Prepare the response body for the client.
130+
responseBody := &convert.ClientResponse{}
131+
resp := data.Config
132+
if resp.Message != nil {
133+
responseBody.Message = &convert.Message{
134+
Role: resp.Message.Role,
135+
Content: resp.Message.Content,
136+
}
137+
if resp.Done {
138+
responseBody.FinishReason = convert.FinishStop
139+
}
140+
} else {
141+
responseBody.Code = -1
142+
responseBody.Error = "response message is nil"
143+
}
144+
145+
// Marshal the modified response body back into JSON.
146+
body, err := json.Marshal(responseBody)
147+
if err != nil {
148+
return nil, err
149+
}
150+
return body, nil
151+
})
152+
}
153+
return nil
154+
}
97155
// Retrieve the response body.
98156
body := httpContext.Response().GetBody()
157+
if body == nil {
158+
return nil
159+
}
99160

100161
// Parse the response body into a base configuration.
101162
data := eosc.NewBase[Response]()
@@ -111,21 +172,20 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
111172
convert.SetAIModelInputToken(ctx, usage.PromptEvalCount)
112173
convert.SetAIModelOutputToken(ctx, usage.EvalCount)
113174
convert.SetAIModelTotalToken(ctx, usage.PromptEvalCount+usage.EvalCount)
114-
115175
}
116176

117177
// Prepare the response body for the client.
118178
responseBody := &convert.ClientResponse{}
119179
resp := data.Config
120180
if resp.Message != nil {
121-
responseBody.Message = convert.Message{
181+
responseBody.Message = &convert.Message{
122182
Role: resp.Message.Role,
123183
Content: resp.Message.Content,
124184
}
125185
responseBody.FinishReason = convert.FinishStop
126186
} else {
127187
responseBody.Code = -1
128-
responseBody.Error = "response message is nil"
188+
responseBody.Error = resp.Error
129189
}
130190

131191
// Marshal the modified response body back into JSON.
@@ -136,5 +196,6 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
136196

137197
httpContext.Response().SetBody(body)
138198

199+
// SetProvider the modified response in the HTTP context.
139200
return nil
140201
}

‎drivers/ai-provider/openAI/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
150150
responseBody := &convert.ClientResponse{}
151151
if len(data.Config.Choices) > 0 {
152152
msg := data.Config.Choices[0]
153-
responseBody.Message = convert.Message{
153+
responseBody.Message = &convert.Message{
154154
Role: msg.Message.Role,
155155
Content: msg.Message.Content,
156156
}

‎drivers/ai-provider/openrouter/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
132132
responseBody := &convert.ClientResponse{}
133133
if len(data.Config.Choices) > 0 {
134134
msg := data.Config.Choices[0]
135-
responseBody.Message = convert.Message{
135+
responseBody.Message = &convert.Message{
136136
Role: msg.Message.Role,
137137
Content: msg.Message.Content,
138138
}

‎drivers/ai-provider/perfxcloud/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
8686
responseBody := &convert.ClientResponse{}
8787
if len(data.Config.Choices) > 0 {
8888
msg := data.Config.Choices[0]
89-
responseBody.Message = convert.Message{
89+
responseBody.Message = &convert.Message{
9090
Role: msg.Message.Role,
9191
Content: msg.Message.Content,
9292
}

‎drivers/ai-provider/spark/mode.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ func (c *Chat) ResponseConvert(ctx eocontext.EoContext) error {
106106
responseBody := &convert.ClientResponse{}
107107
if len(data.Config.Choices) > 0 {
108108
msg := data.Config.Choices[0]
109-
responseBody.Message = convert.Message{
109+
responseBody.Message = &convert.Message{
110110
Role: msg.Message.Role,
111111
Content: msg.Message.Content,
112112
}

0 commit comments

Comments
 (0)