mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-14 00:48:26 +08:00
e48d75987c
### What problem does this PR solve? 1. Supports stream and non-stream chat 2. Supports think and non-think chat 3. List supported models from DeepSeek service. (This command can be used to verify the API validity) ### Type of change - [x] New Feature (non-breaking change which adds functionality) --------- Signed-off-by: Jin Hai <haijin.chn@gmail.com>
119 lines
3.5 KiB
Go
119 lines
3.5 KiB
Go
//
|
|
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
//
|
|
|
|
package models
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"time"
|
|
)
|
|
|
|
// MooshotModel implements ModelDriver for Mooshot
|
|
type MooshotModel struct {
|
|
BaseURL map[string]string
|
|
URLSuffix URLSuffix
|
|
httpClient *http.Client // Reusable HTTP client with connection pool
|
|
}
|
|
|
|
// NewMooshotModel creates a new Mooshot model instance
|
|
func NewMooshotModel(baseURL map[string]string, urlSuffix URLSuffix) *MooshotModel {
|
|
return &MooshotModel{
|
|
BaseURL: baseURL,
|
|
URLSuffix: urlSuffix,
|
|
httpClient: &http.Client{
|
|
Timeout: 120 * time.Second,
|
|
Transport: &http.Transport{
|
|
MaxIdleConns: 100,
|
|
MaxIdleConnsPerHost: 10,
|
|
IdleConnTimeout: 90 * time.Second,
|
|
DisableCompression: false,
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
// Chat sends a message and returns response
|
|
func (z *MooshotModel) Chat(modelName, message *string, apiConfig *APIConfig, chatModelConfig *ChatConfig) (*ChatResponse, error) {
|
|
return nil, fmt.Errorf("not implemented")
|
|
}
|
|
|
|
// ChatStreamlyWithSender sends a message and streams response via sender function (best performance, no channel)
|
|
func (z *MooshotModel) ChatStreamlyWithSender(modelName, message *string, apiConfig *APIConfig, chatModelConfig *ChatConfig, sender func(*string, *string) error) error {
|
|
return fmt.Errorf("not implemented")
|
|
}
|
|
|
|
// EncodeToEmbedding encodes a list of texts into embeddings
|
|
func (z *MooshotModel) EncodeToEmbedding(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([][]float64, error) {
|
|
return nil, fmt.Errorf("not implemented")
|
|
}
|
|
|
|
func (z *MooshotModel) ListModels(apiConfig *APIConfig) ([]string, error) {
|
|
var region = "default"
|
|
if apiConfig.Region != nil {
|
|
region = *apiConfig.Region
|
|
}
|
|
|
|
url := fmt.Sprintf("%s/%s", z.BaseURL[region], z.URLSuffix.Models)
|
|
|
|
// Build request body
|
|
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("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 := 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)
|
|
}
|
|
|
|
models, ok := result["models"].([]string)
|
|
if !ok || len(models) == 0 {
|
|
return nil, fmt.Errorf("no models in response")
|
|
}
|
|
|
|
return models, nil
|
|
}
|