Go: implement provider: lm-studio (#14586)

### What problem does this PR solve?

implement `lm-studio` provider

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
- [x] Refactoring
This commit is contained in:
Haruko386
2026-05-06 19:23:11 +08:00
committed by GitHub
parent 33d8320ce8
commit dd7a0ce1d3
4 changed files with 444 additions and 9 deletions

View File

@@ -0,0 +1,8 @@
{
"name": "lmstudio",
"url_suffix": {
"chat": "chat/completions",
"models": "models"
},
"class": "local"
}

View File

@@ -55,6 +55,8 @@ func (f *ModelFactory) CreateModelDriver(providerName string, baseURL map[string
return NewVllmModel(baseURL, urlSuffix), nil
case "xai":
return NewXAIModel(baseURL, urlSuffix), nil
case "lmstudio":
return NewLmStudioModel(baseURL, urlSuffix), nil
default:
return NewDummyModel(baseURL, urlSuffix), nil
}

View File

@@ -0,0 +1,425 @@
package models
import (
"bufio"
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
"ragflow/internal/common"
"strings"
"time"
)
// LmStudioModel implements ModelDriver for lm-studio
type LmStudioModel struct {
BaseURL map[string]string
URLSuffix URLSuffix
httpClient *http.Client
}
// NewLmStudioModel
func NewLmStudioModel(baseURL map[string]string, urlSuffix URLSuffix) *LmStudioModel {
return &LmStudioModel{
BaseURL: baseURL,
URLSuffix: urlSuffix,
httpClient: &http.Client{
Timeout: 120 * time.Second,
Transport: &http.Transport{
MaxIdleConns: 100,
MaxIdleConnsPerHost: 10,
IdleConnTimeout: 90 * time.Second,
DisableCompression: false,
},
},
}
}
func (l *LmStudioModel) NewInstance(baseURL map[string]string) ModelDriver {
return &LmStudioModel{
BaseURL: baseURL,
URLSuffix: l.URLSuffix,
httpClient: &http.Client{
Timeout: 120 * time.Second,
Transport: &http.Transport{
MaxIdleConns: 100,
MaxIdleConnsPerHost: 10,
IdleConnTimeout: 90 * time.Second,
DisableCompression: false,
},
},
}
}
func (l *LmStudioModel) Name() string {
return "lmstudio"
}
// ChatWithMessages sends multiple messages with roles and returns response
func (l *LmStudioModel) ChatWithMessages(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig) (*ChatResponse, error) {
if len(messages) == 0 {
return nil, fmt.Errorf("messages is empty")
}
var region = "default"
if apiConfig.Region != nil {
region = *apiConfig.Region
}
url := fmt.Sprintf("%s/%s", l.BaseURL[region], l.URLSuffix.Chat)
// For qwen/glm models, use async chat endpoint
modelType := strings.Split(modelName, "-")[0]
if modelType == "qwen" || modelType == "glm" {
url = fmt.Sprintf("%s/%s", l.BaseURL[region], l.URLSuffix.AsyncChat)
}
// Convert messages to API format
apiMessages := make([]map[string]interface{}, len(messages))
for i, msg := range messages {
apiMessages[i] = map[string]interface{}{
"role": msg.Role,
"content": msg.Content,
}
}
// Build request body
reqBody := map[string]interface{}{
"model": modelName,
"messages": apiMessages,
"stream": false,
"temperature": 1,
}
if chatModelConfig != nil {
if chatModelConfig.Stream != nil {
reqBody["stream"] = *chatModelConfig.Stream
}
if chatModelConfig.MaxTokens != nil {
reqBody["max_tokens"] = *chatModelConfig.MaxTokens
}
if chatModelConfig.Temperature != nil {
reqBody["temperature"] = *chatModelConfig.Temperature
}
if chatModelConfig.TopP != nil {
reqBody["top_p"] = *chatModelConfig.TopP
}
if chatModelConfig.Stop != nil {
reqBody["stop"] = *chatModelConfig.Stop
}
if chatModelConfig.Thinking != nil {
if *chatModelConfig.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 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", *apiConfig.ApiKey))
resp, err := l.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 :%s", resp.StatusCode, string(body), messages[0].Content)
}
// 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)
}
choices, ok := result["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil, fmt.Errorf("no choices in response")
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("invalid choice format")
}
messageMap, ok := firstChoice["message"].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("invalid message format")
}
content, ok := messageMap["content"].(string)
if !ok {
return nil, fmt.Errorf("invalid content format")
}
thinking, answer := GetThinkingAndAnswer(chatModelConfig.ModelClass, &content)
chatResponse := &ChatResponse{
Answer: answer,
ReasonContent: thinking,
}
return chatResponse, nil
}
// ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel)
func (l *LmStudioModel) ChatStreamlyWithSender(modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, sender func(*string, *string) error) error {
if len(messages) == 0 {
return fmt.Errorf("messages is empty")
}
var region = "default"
if apiConfig.Region != nil {
region = *apiConfig.Region
}
url := fmt.Sprintf("%s/%s", l.BaseURL[region], l.URLSuffix.Chat)
modelType := strings.Split(modelName, "-")[0]
if modelType == "qwen" || modelType == "glm" {
url = fmt.Sprintf("%s/%s", l.BaseURL[region], l.URLSuffix.AsyncChat)
}
// Convert messages to API format (supporting multimodal content)
apiMessages := make([]map[string]interface{}, len(messages))
for i, msg := range messages {
apiMessages[i] = map[string]interface{}{
"role": msg.Role,
"content": msg.Content,
}
}
// Build request body with streaming enabled
reqBody := map[string]interface{}{
"model": modelName,
"messages": apiMessages,
"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 := l.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()
common.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()
}
func (l *LmStudioModel) Encode(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([][]float64, error) {
return nil, fmt.Errorf("no such method")
}
func (l *LmStudioModel) Rerank(modelName *string, query string, texts []string, apiConfig *APIConfig) ([]float64, error) {
return nil, fmt.Errorf("no such method")
}
// ListModels list supported models
func (l *LmStudioModel) ListModels(apiConfig *APIConfig) ([]string, error) {
var region = "default"
if apiConfig.Region != nil {
region = *apiConfig.Region
}
url := fmt.Sprintf("%s/%s", l.BaseURL[region], l.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 := l.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 body: %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"] 2 []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 (l *LmStudioModel) Balance(apiConfig *APIConfig) (map[string]interface{}, error) {
return nil, fmt.Errorf("no such method")
}
func (l *LmStudioModel) CheckConnection(apiConfig *APIConfig) error {
return fmt.Errorf("no such method")
}

View File

@@ -36,7 +36,7 @@ func NewOllamaModel(baseURL map[string]string, urlSuffix URLSuffix) *OllamaModel
}
}
func (o OllamaModel) NewInstance(baseURL map[string]string) ModelDriver {
func (o *OllamaModel) NewInstance(baseURL map[string]string) ModelDriver {
return &OllamaModel{
BaseURL: baseURL,
URLSuffix: o.URLSuffix,
@@ -52,11 +52,11 @@ func (o OllamaModel) NewInstance(baseURL map[string]string) ModelDriver {
}
}
func (o OllamaModel) Name() string {
func (o *OllamaModel) Name() string {
return "ollama"
}
func (o OllamaModel) ChatWithMessages(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig) (*ChatResponse, error) {
func (o *OllamaModel) ChatWithMessages(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig) (*ChatResponse, error) {
if len(messages) == 0 {
return nil, fmt.Errorf("message is nil")
}
@@ -189,7 +189,7 @@ func (o OllamaModel) ChatWithMessages(modelName string, messages []Message, apiC
return chatResponse, nil
}
func (o OllamaModel) ChatStreamlyWithSender(modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, sender func(*string, *string) error) error {
func (o *OllamaModel) ChatStreamlyWithSender(modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, sender func(*string, *string) error) error {
if len(messages) == 0 {
return fmt.Errorf("messages is empty")
}
@@ -350,15 +350,15 @@ func (o OllamaModel) ChatStreamlyWithSender(modelName string, messages []Message
return scanner.Err()
}
func (o OllamaModel) Encode(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([][]float64, error) {
func (o *OllamaModel) Encode(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([][]float64, error) {
return nil, fmt.Errorf("no such method")
}
func (o OllamaModel) Rerank(modelName *string, query string, texts []string, apiConfig *APIConfig) ([]float64, error) {
func (o *OllamaModel) Rerank(modelName *string, query string, texts []string, apiConfig *APIConfig) ([]float64, error) {
return nil, fmt.Errorf("no such method")
}
func (o OllamaModel) ListModels(apiConfig *APIConfig) ([]string, error) {
func (o *OllamaModel) ListModels(apiConfig *APIConfig) ([]string, error) {
var region = "default"
if apiConfig.Region != nil {
@@ -414,10 +414,10 @@ func (o OllamaModel) ListModels(apiConfig *APIConfig) ([]string, error) {
return models, nil
}
func (o OllamaModel) Balance(apiConfig *APIConfig) (map[string]interface{}, error) {
func (o *OllamaModel) Balance(apiConfig *APIConfig) (map[string]interface{}, error) {
return nil, fmt.Errorf("no such method")
}
func (o OllamaModel) CheckConnection(apiConfig *APIConfig) error {
func (o *OllamaModel) CheckConnection(apiConfig *APIConfig) error {
return fmt.Errorf("no such method")
}