diff --git a/internal/entity/models/vllm.go b/internal/entity/models/vllm.go index 1d5055a950..def9297ccf 100644 --- a/internal/entity/models/vllm.go +++ b/internal/entity/models/vllm.go @@ -17,11 +17,13 @@ package models import ( + "bufio" "bytes" "encoding/json" "fmt" "io" "net/http" + "ragflow/internal/logger" "strings" "time" ) @@ -227,7 +229,7 @@ func (z *VllmModel) ChatWithMessages(modelName string, apiConfig *APIConfig, mes // Build request body reqBody := map[string]interface{}{ "model": modelName, - "messages": apiMessages, + "messages": apiMessages, "stream": false, "temperature": 1, } @@ -332,7 +334,154 @@ func (z *VllmModel) ChatWithMessages(modelName string, apiConfig *APIConfig, mes // ChatStreamlyWithSender sends a message and streams response via sender function (best performance, no channel) func (z *VllmModel) ChatStreamlyWithSender(modelName, message *string, apiConfig *APIConfig, modelConfig *ChatConfig, sender func(*string, *string) error) error { - return fmt.Errorf("not implemented") + var region = "default" + if apiConfig.Region != nil { + region = *apiConfig.Region + } + + url := fmt.Sprintf("%s/%s", z.BaseURL[region], z.URLSuffix.Chat) + // TODO figure out what's the point of these codes + modelType := strings.Split(*modelName, "-")[0] + if modelType == "qwen" || modelType == "glm" { + url = fmt.Sprintf("%s/%s", z.BaseURL[region], z.URLSuffix.AsyncChat) + } + + // Build request body with streaming enabled + reqBody := map[string]interface{}{ + "model": modelName, + "messages": []map[string]string{ + {"role": "user", "content": *message}, + }, + "stream": true, + } + + if modelConfig.Stream != nil { + reqBody["stream"] = *modelConfig.Stream + } + + if modelConfig.MaxTokens != nil { + reqBody["max_tokens"] = *modelConfig.MaxTokens + } + + if modelConfig.Temperature != nil { + reqBody["temperature"] = *modelConfig.Temperature + } + + if modelConfig.DoSample != nil { + reqBody["do_sample"] = *modelConfig.DoSample + } + + if modelConfig.TopP != nil { + reqBody["top_p"] = *modelConfig.TopP + } + + if modelConfig.Stop != nil { + reqBody["stop"] = *modelConfig.Stop + } + + if modelConfig.Thinking != nil { + if *modelConfig.Thinking { + reqBody["thinking"] = map[string]interface{}{ + "type": "enabled", + } + } else { + reqBody["thinking"] = map[string]interface{}{ + "type": "disabled", + } + } + } + + 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", *apiConfig.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 + } + + reasoningContent, ok := delta["reasoning_content"].(string) + if ok && reasoningContent != "" { + if err := sender(nil, &reasoningContent); err != nil { + return err + } + } + + content, ok := delta["content"].(string) + if ok && content != "" { + if err := sender(&content, nil); err != nil { + return err + } + } + + finishReason, ok := firstChoice["finish_reason"].(string) + if ok && finishReason != "" { + break + } + } + + // Send [DONE] marker for OpenAI compatibility + endOfStream := "[DONE]" + if err = sender(&endOfStream, nil); err != nil { + return err + } + + return scanner.Err() } // Encode encodes a list of texts into embeddings @@ -341,7 +490,59 @@ func (z *VllmModel) Encode(modelName *string, texts []string, apiConfig *APIConf } func (z *VllmModel) ListModels(apiConfig *APIConfig) ([]string, error) { - return nil, fmt.Errorf("not implemented") + var region = "default" + + if apiConfig.Region != nil { + region = *apiConfig.Region + } + + url := fmt.Sprintf("%s/%s", z.BaseURL[region], z.URLSuffix.Models) + + reqBody := map[string]interface{}{} + + jsonData, err := json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + req, err := http.NewRequest("GET", 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", *apiConfig.ApiKey)) + + resp, err := z.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) + } + + // Parse response + var result map[string]interface{} + if err = json.Unmarshal(body, &result); err != nil { + return nil, fmt.Errorf("failed to parse response: %w", err) + } + + // convert result["data"] to []map[string]interface{} + models := make([]string, 0) + for _, model := range result["data"].([]interface{}) { + modelMap := model.(map[string]interface{}) + modelName := modelMap["id"].(string) + models = append(models, modelName) + } + + return models, nil } func (z *VllmModel) Balance(apiConfig *APIConfig) (map[string]interface{}, error) { diff --git a/internal/service/model_service.go b/internal/service/model_service.go index f2a4f2f750..8dd6c211b2 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -199,6 +199,14 @@ func (m *ModelProviderService) ListSupportedModels(providerName, instanceName, u apiConfig.Region = ®ion apiConfig.ApiKey = &instance.APIKey + // For local deployed models + if baseURL, ok := extra["base_url"]; ok && baseURL != "" { + newURL := map[string]string{ + region: baseURL, + } + providerInfo.ModelDriver = providerInfo.ModelDriver.NewInstance(newURL) + } + return providerInfo.ModelDriver.ListModels(apiConfig) } @@ -831,7 +839,7 @@ func (m *ModelProviderService) ChatToModelStreamWithSender(providerName, instanc return common.CodeServerError, err } - _, err = m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(provider.ID, instance.ID, modelName) + modelInfo, err := m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(provider.ID, instance.ID, modelName) if err != nil { providerInfo := dao.GetModelProviderManager().FindProvider(providerName) if providerInfo == nil { @@ -853,7 +861,6 @@ func (m *ModelProviderService) ChatToModelStreamWithSender(providerName, instanc apiConfig.Region = ®ion apiConfig.ApiKey = &instance.APIKey - // Direct call with sender function err = providerInfo.ModelDriver.ChatStreamlyWithSender(&modelName, &message, apiConfig, modelConfig, sender) if err != nil { return common.CodeServerError, err @@ -862,6 +869,37 @@ func (m *ModelProviderService) ChatToModelStreamWithSender(providerName, instanc return common.CodeSuccess, nil } + if modelInfo.Status == "active" { + // For local deployed models + providerInfo := dao.GetModelProviderManager().FindProvider(providerName) + if providerInfo == nil { + return common.CodeNotFound, errors.New("provider not found") + } + + var extra map[string]string + err = json.Unmarshal([]byte(instance.Extra), &extra) + if err != nil { + return common.CodeServerError, err + } + + region := extra["region"] + apiConfig.Region = ®ion + apiConfig.ApiKey = &instance.APIKey + + modelConfig.ModelClass = &providerInfo.Class + + newURL := map[string]string{ + region: extra["base_url"], + } + newProviderInfo := providerInfo.ModelDriver.NewInstance(newURL) + + err = newProviderInfo.ChatStreamlyWithSender(&modelName, &message, apiConfig, modelConfig, sender) + if err != nil { + return common.CodeServerError, err + } + return common.CodeSuccess, nil + } + return common.CodeServerError, errors.New("model is disabled") }