zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-18 05:37:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add extra field to model instance (#14203)

Browse Source 
### What problem does this PR solve?

Now each model support region with different URL

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

---------

Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-04-20 15:31:12 +08:00
committed by GitHub
parent 939933649a
commit af2ed416a7
22 changed files with 398 additions and 550 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
internal/entity/model.go
View File

@@ -145,11 +145,10 @@ type Model struct {
// Provider represents an LLM provider
type Provider struct {
Name string `json:"name"`
Tags string `json:"tags"`
URL string `json:"url"`
URLSuffix models.URLSuffix `json:"url_suffix"`
Models []*Model `json:"models"`
Name string `json:"name"`
URL map[string]string `json:"url"`
URLSuffix models.URLSuffix `json:"url_suffix"`
Models []*Model `json:"models"`
ModelDriver models.ModelDriver
}
@@ -236,11 +235,24 @@ func (pm *ProviderManager) ListProviders() ([]map[string]interface{}, error) {
var providers []map[string]interface{}
for _, provider := range pm.Providers {
modelTypeSet := make(map[string]struct{})
for _, model := range provider.Models {
for _, modelType := range model.ModelTypes {
modelTypeSet[modelType] = struct{}{}
}
}
var modelTypes []string
for modelType := range modelTypeSet {
modelTypes = append(modelTypes, modelType)
}
providerData := map[string]interface{}{
"name": provider.Name,
"tags": provider.Tags,
"url": provider.URL,
"url_suffix": provider.URLSuffix,
"name": provider.Name,
"url": provider.URL,
"model_types": modelTypes,
"url_suffix": provider.URLSuffix,
}
providers = append(providers, providerData)
}
@@ -262,7 +274,6 @@ func (pm *ProviderManager) GetProviderByName(providerName string) (map[string]in
providerInfo := map[string]interface{}{
"name": provider.Name,
"tags": provider.Tags,
"base_url": provider.URL,
"total_models": len(provider.Models),
}

20
internal/entity/models/dummy.go
View File

@@ -20,14 +20,14 @@ import (
"fmt"
)
// DummyModel implements ModelDriver for Zhipu AI (智谱 AI)
// DummyModel implements ModelDriver for Zhipu AI
type DummyModel struct {
BaseURL string
BaseURL map[string]string
URLSuffix URLSuffix
}
// NewDummyModel creates a new Zhipu AI model instance
func NewDummyModel(baseURL string, urlSuffix URLSuffix) *DummyModel {
func NewDummyModel(baseURL map[string]string, urlSuffix URLSuffix) *DummyModel {
return &DummyModel{
BaseURL: baseURL,
URLSuffix: urlSuffix,
@@ -35,26 +35,16 @@ func NewDummyModel(baseURL string, urlSuffix URLSuffix) *DummyModel {
}
// Chat sends a message and returns response
func (z *DummyModel) Chat(modelName, apiKey, message *string, genConf map[string]interface{}) (string, error) {
func (z *DummyModel) Chat(modelName, apiKey, message *string, modelConfig *ChatConfig) (string, error) {
return "", fmt.Errorf("not implemented")
}
// ChatStreamly sends a message and streams response
func (z *DummyModel) ChatStreamly(modelName, apiKey, message *string, genConf map[string]interface{}) (<-chan string, error) {
return nil, fmt.Errorf("not implemented")
}
// ChatStreamlyWithChannel sends a message and streams response to channel (better performance)
func (z *DummyModel) ChatStreamlyWithChannel(modelName, apiKey, message *string, genConf map[string]interface{}, resultChan chan<- string) error {
return fmt.Errorf("not implemented")
}
// ChatStreamlyWithSender sends a message and streams response via sender function (best performance, no channel)
func (z *DummyModel) ChatStreamlyWithSender(modelName, apiKey, message *string, modelConfig *ChatConfig, sender func(*string, *string) error) error {
return fmt.Errorf("not implemented")
}
// EncodeToEmbedding encodes a list of texts into embeddings
func (z *DummyModel) EncodeToEmbedding(modelName, apiKey *string, texts []string) ([][]float64, error) {
func (z *DummyModel) EncodeToEmbedding(modelName, apiKey *string, texts []string, embeddingConfig *EmbeddingConfig) ([][]float64, error) {
return nil, fmt.Errorf("not implemented")
}

2
internal/entity/models/factory.go
View File

@@ -30,7 +30,7 @@ func NewModelFactory() *ModelFactory {
}
// CreateModelDriver creates a ModelDriver for the given provider and model
func (f *ModelFactory) CreateModelDriver(providerName string, baseURL string, urlSuffix URLSuffix) (ModelDriver, error) {
func (f *ModelFactory) CreateModelDriver(providerName string, baseURL map[string]string, urlSuffix URLSuffix) (ModelDriver, error) {
providerLower := strings.ToLower(providerName)
switch providerLower {
case "zhipu-ai":

13
internal/entity/models/types.go
View File

@@ -3,15 +3,11 @@ package models
// EmbeddingModel interface for embedding models
type ModelDriver interface {
// Chat sends a message and returns response
Chat(modelName, apiKey, message *string, genConf map[string]interface{}) (string, error)
// ChatStreamly sends a message and streams response
ChatStreamly(modelName, apiKey, message *string, genConf map[string]interface{}) (<-chan string, error)
// ChatStreamlyWithChannel sends a message and streams response to channel (better performance)
ChatStreamlyWithChannel(modelName, apiKey, message *string, genConf map[string]interface{}, resultChan chan<- string) error
Chat(modelName, apiKey, message *string, modelConfig *ChatConfig) (string, error)
// ChatStreamlyWithSender sends a message and streams response via sender function (best performance, no channel)
ChatStreamlyWithSender(modelName, apiKey, message *string, modelConfig *ChatConfig, sender func(*string, *string) error) error
// Encode encodes a list of texts into embeddings
EncodeToEmbedding(modelName, apiKey *string, texts []string) ([][]float64, error)
EncodeToEmbedding(modelName, apiKey *string, texts []string, embeddingConfig *EmbeddingConfig) ([][]float64, error)
}
// URLSuffix represents the URL suffixes for different API endpoints
@@ -31,4 +27,9 @@ type ChatConfig struct {
TopP *float64
DoSample *bool
Stop *[]string
Region *string
}
type EmbeddingConfig struct {
Region *string
}

291
internal/entity/models/zhipu-ai.go
View File

@@ -30,13 +30,13 @@ import (
// ZhipuAIModel implements ModelDriver for Zhipu AI
type ZhipuAIModel struct {
BaseURL string
BaseURL map[string]string
URLSuffix URLSuffix
httpClient *http.Client // Reusable HTTP client with connection pool
}
// NewZhipuAIModel creates a new Zhipu AI model instance
func NewZhipuAIModel(baseURL string, urlSuffix URLSuffix) *ZhipuAIModel {
func NewZhipuAIModel(baseURL map[string]string, urlSuffix URLSuffix) *ZhipuAIModel {
return &ZhipuAIModel{
BaseURL: baseURL,
URLSuffix: urlSuffix,
@@ -53,7 +53,7 @@ func NewZhipuAIModel(baseURL string, urlSuffix URLSuffix) *ZhipuAIModel {
}
// Chat sends a message and returns response
func (z *ZhipuAIModel) Chat(modelName, apiKey, message *string, genConf map[string]interface{}) (string, error) {
func (z *ZhipuAIModel) Chat(modelName, apiKey, message *string, chatModelConfig *ChatConfig) (string, error) {
if message == nil {
return "", fmt.Errorf("message is nil")
}
@@ -70,16 +70,17 @@ func (z *ZhipuAIModel) Chat(modelName, apiKey, message *string, genConf map[stri
"temperature": 1,
}
// Add generation config if provided
if genConf != nil {
if maxTokens, ok := genConf["max_tokens"]; ok {
reqBody["max_tokens"] = maxTokens
if chatModelConfig != nil {
if chatModelConfig.MaxTokens != nil {
reqBody["max_tokens"] = *chatModelConfig.MaxTokens
}
if temperature, ok := genConf["temperature"]; ok {
reqBody["temperature"] = temperature
if chatModelConfig.Temperature != nil {
reqBody["temperature"] = *chatModelConfig.Temperature
}
if topP, ok := genConf["top_p"]; ok {
reqBody["top_p"] = topP
if chatModelConfig.TopP != nil {
reqBody["top_p"] = *chatModelConfig.TopP
}
}
@@ -140,229 +141,14 @@ func (z *ZhipuAIModel) Chat(modelName, apiKey, message *string, genConf map[stri
return content, nil
}
// ChatStreamly sends a message and streams response
func (z *ZhipuAIModel) ChatStreamly(modelName, apiKey, message *string, genConf map[string]interface{}) (<-chan string, error) {
url := fmt.Sprintf("%s/chat/completions", z.BaseURL)
// Build request body with streaming enabled
reqBody := map[string]interface{}{
"model": modelName,
"messages": []map[string]string{
{"role": "user", "content": *message},
},
"stream": true,
"temperature": 1,
}
// Add generation config if provided
if genConf != nil {
if maxTokens, ok := genConf["max_tokens"]; ok {
reqBody["max_tokens"] = maxTokens
}
if temperature, ok := genConf["temperature"]; ok {
reqBody["temperature"] = temperature
}
if topP, ok := genConf["top_p"]; ok {
reqBody["top_p"] = topP
}
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiKey))
resp, err := z.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
resp.Body.Close()
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// Create channel for streaming
resultChan := make(chan string)
go func() {
defer close(resultChan)
defer resp.Body.Close()
// SSE parsing: read line by line
scanner := bufio.NewScanner(resp.Body)
for scanner.Scan() {
line := scanner.Text()
// SSE data line starts with "data:"
if !strings.HasPrefix(line, "data:") {
continue
}
// Extract JSON after "data:"
data := strings.TrimSpace(line[5:])
// [DONE] marks the end of stream
if data == "[DONE]" {
break
}
// Parse the JSON event
var event map[string]interface{}
if err := json.Unmarshal([]byte(data), &event); err != nil {
continue
}
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
continue
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
continue
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
continue
}
content, ok := delta["content"].(string)
if ok && content != "" {
resultChan <- content
}
finishReason, ok := firstChoice["finish_reason"].(string)
if ok && finishReason != "" {
break
}
}
}()
return resultChan, nil
}
// ChatStreamlyWithChannel sends a message and streams response to channel (better performance)
func (z *ZhipuAIModel) ChatStreamlyWithChannel(modelName, apiKey, message *string, genConf map[string]interface{}, resultChan chan<- string) error {
url := fmt.Sprintf("%s/chat/completions", z.BaseURL)
// Build request body with streaming enabled
reqBody := map[string]interface{}{
"model": modelName,
"messages": []map[string]string{
{"role": "user", "content": *message},
},
"stream": true,
"temperature": 1,
}
// Add generation config if provided
if genConf != nil {
if maxTokens, ok := genConf["max_tokens"]; ok {
reqBody["max_tokens"] = maxTokens
}
if temperature, ok := genConf["temperature"]; ok {
reqBody["temperature"] = temperature
}
if topP, ok := genConf["top_p"]; ok {
reqBody["top_p"] = topP
}
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
}
req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiKey))
resp, err := z.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// SSE parsing: read line by line
scanner := bufio.NewScanner(resp.Body)
for scanner.Scan() {
line := scanner.Text()
logger.Info(line)
// SSE data line starts with "data:"
if !strings.HasPrefix(line, "data:") {
continue
}
// Extract JSON after "data:"
data := strings.TrimSpace(line[5:])
// [DONE] marks the end of stream
if data == "[DONE]" {
break
}
// Parse the JSON event
var event map[string]interface{}
if err := json.Unmarshal([]byte(data), &event); err != nil {
continue
}
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
continue
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
continue
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
continue
}
content, ok := delta["content"].(string)
if ok && content != "" {
resultChan <- content
}
finishReason, ok := firstChoice["finish_reason"].(string)
if ok && finishReason != "" {
break
}
}
// Send [DONE] marker for OpenAI compatibility
resultChan <- "[DONE]"
return scanner.Err()
}
// ChatStreamlyWithSender sends a message and streams response via sender function (best performance, no channel)
func (z *ZhipuAIModel) ChatStreamlyWithSender(modelName, apiKey, message *string, modelConfig *ChatConfig, sender func(*string, *string) error) error {
url := fmt.Sprintf("%s/chat/completions", z.BaseURL)
func (z *ZhipuAIModel) ChatStreamlyWithSender(modelName, apiKey, message *string, chatModelConfig *ChatConfig, sender func(*string, *string) error) error {
var region = "default"
if chatModelConfig.Region != nil {
region = *chatModelConfig.Region
}
url := fmt.Sprintf("%s/chat/completions", z.BaseURL[region])
// Build request body with streaming enabled
reqBody := map[string]interface{}{
@@ -374,33 +160,33 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(modelName, apiKey, message *string
"temperature": 1,
}
if modelConfig != nil {
if modelConfig.Stream != nil {
reqBody["stream"] = *modelConfig.Stream
if chatModelConfig != nil {
if chatModelConfig.Stream != nil {
reqBody["stream"] = *chatModelConfig.Stream
}
if modelConfig.MaxTokens != nil {
reqBody["max_tokens"] = *modelConfig.MaxTokens
if chatModelConfig.MaxTokens != nil {
reqBody["max_tokens"] = *chatModelConfig.MaxTokens
}
if modelConfig.Temperature != nil {
reqBody["temperature"] = *modelConfig.Temperature
if chatModelConfig.Temperature != nil {
reqBody["temperature"] = *chatModelConfig.Temperature
}
if modelConfig.DoSample != nil {
reqBody["do_sample"] = *modelConfig.DoSample
if chatModelConfig.DoSample != nil {
reqBody["do_sample"] = *chatModelConfig.DoSample
}
if modelConfig.TopP != nil {
reqBody["top_p"] = *modelConfig.TopP
if chatModelConfig.TopP != nil {
reqBody["top_p"] = *chatModelConfig.TopP
}
if modelConfig.Stop != nil {
reqBody["stop"] = *modelConfig.Stop
if chatModelConfig.Stop != nil {
reqBody["stop"] = *chatModelConfig.Stop
}
if modelConfig.Reasoning != nil {
if *modelConfig.Reasoning {
if chatModelConfig.Reasoning != nil {
if *chatModelConfig.Reasoning {
reqBody["thinking"] = map[string]interface{}{
"type": "enabled",
}
@@ -506,8 +292,13 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(modelName, apiKey, message *string
}
// EncodeToEmbedding encodes a list of texts into embeddings
func (z *ZhipuAIModel) EncodeToEmbedding(modelName, apiKey *string, texts []string) ([][]float64, error) {
url := fmt.Sprintf("%s/embedding", z.BaseURL)
func (z *ZhipuAIModel) EncodeToEmbedding(modelName, apiKey *string, texts []string, embeddingConfig *EmbeddingConfig) ([][]float64, error) {
var region = "default"
if embeddingConfig.Region != nil {
region = *embeddingConfig.Region
}
url := fmt.Sprintf("%s/embedding", z.BaseURL[region])
embeddings := make([][]float64, len(texts))

2
internal/entity/tenant.go
View File

@@ -34,6 +34,8 @@ type Tenant struct {
TTSID *string `gorm:"column:tts_id;size:256;index" json:"tts_id,omitempty"`
TenantTTSID *int64 `gorm:"column:tenant_tts_id;index" json:"tenant_tts_id,omitempty"`
ParserIDs string `gorm:"column:parser_ids;size:256;not null;index" json:"parser_ids"`
OCRID string `gorm:"column:ocr_id;size:256;not null" json:"ocr_id"`
TenantOCRID *int64 `gorm:"column:tenant_ocr_id" json:"tenant_ocr_id,omitempty"`
Credit int64 `gorm:"column:credit;default:512;index" json:"credit"`
Status *string `gorm:"column:status;size:1;index" json:"status,omitempty"`
BaseModel

1
internal/entity/tenant_model_instance.go
View File

@@ -23,6 +23,7 @@ type TenantModelInstance struct {
ProviderID string `gorm:"column:provider_id;size:32;not null;uniqueIndex:idx_api_key_provider_id" json:"provider_id"`
APIKey string `gorm:"column:api_key;size:512;not null;uniqueIndex:idx_api_key_provider_id" json:"api_key"`
Status string `gorm:"column:status;size:32;default:'active'" json:"status"`
Extra string `gorm:"column:extra;size:512;default:'active'" json:"extra"`
BaseModel
}