Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: vertexai support proxy url(example: cloudflare ai gateway) and fix some vertexai bug (#1642)

Browse Source 
* feat: vertexai support proxy url(example: cloudflare ai gateway)

* fix: do resp model mapping

* fix: missing system

* fix: stream need query alt=sse
This commit is contained in:
zijiren 2024-07-16 01:02:06 +08:00 committed by GitHub
parent 1c44d7e1cd
commit 6209ff9ea9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 30 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
relay/adaptor/vertexai/adaptor.go
View File

@ -19,7 +19,7 @@ var _ adaptor.Adaptor = new(Adaptor)
const channelName = "vertexai" const channelName = "vertexai"
type Adaptor struct {} type Adaptor struct{}
func (a *Adaptor) Init(meta *meta.Meta) { func (a *Adaptor) Init(meta *meta.Meta) {
} }
@ -38,7 +38,7 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G
} }
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) {
adaptor := GetAdaptor(meta.OriginModelName) adaptor := GetAdaptor(meta.ActualModelName)
if adaptor == nil { if adaptor == nil {
return nil, &relaymodel.ErrorWithStatusCode{ return nil, &relaymodel.ErrorWithStatusCode{
StatusCode: http.StatusInternalServerError, StatusCode: http.StatusInternalServerError,
@ -63,20 +63,36 @@ func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) {
suffix := "" suffix := ""
if strings.HasPrefix(meta.ActualModelName, "gemini") { if strings.HasPrefix(meta.ActualModelName, "gemini") {
if meta.IsStream { if meta.IsStream {
suffix = "streamGenerateContent" suffix = "streamGenerateContent?alt=sse"
} else { } else {
suffix = "generateContent" suffix = "generateContent"
} }
} else { } else {
if meta.IsStream { if meta.IsStream {
suffix = "streamRawPredict" suffix = "streamRawPredict?alt=sse"
} else { } else {
suffix = "rawPredict" suffix = "rawPredict"
} }
} }
baseUrl := fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:%s", meta.Config.Region, meta.Config.VertexAIProjectID, meta.Config.Region, meta.ActualModelName, suffix) if meta.BaseURL != "" {
return baseUrl, nil return fmt.Sprintf(
"%s/v1/projects/%s/locations/%s/publishers/google/models/%s:%s",
meta.BaseURL,
meta.Config.VertexAIProjectID,
meta.Config.Region,
meta.ActualModelName,
suffix,
), nil
}
return fmt.Sprintf(
"https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:%s",
meta.Config.Region,
meta.Config.VertexAIProjectID,
meta.Config.Region,
meta.ActualModelName,
suffix,
), nil
} }
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error {

1
relay/adaptor/vertexai/claude/adapter.go
View File

@ -31,6 +31,7 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G
AnthropicVersion: anthropicVersion, AnthropicVersion: anthropicVersion,
// Model: claudeReq.Model, // Model: claudeReq.Model,
Messages: claudeReq.Messages, Messages: claudeReq.Messages,
System: claudeReq.System,
MaxTokens: claudeReq.MaxTokens, MaxTokens: claudeReq.MaxTokens,
Temperature: claudeReq.Temperature, Temperature: claudeReq.Temperature,
TopP: claudeReq.TopP, TopP: claudeReq.TopP,