mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-12 14:45:42 +08:00
### What problem does this PR solve? The LocalAI Go driver landed in #14809 and Embed landed in #14811. `Rerank` was left as a stub that returns `"not implemented"`. This PR fills the gap. LocalAI exposes a public rerank endpoint at `<tenant-url>/v1/rerank` with a Cohere-shaped request and response (`{model, query, documents, top_n}` → `{results: [{index, relevance_score}]}`). The Python side has had `LocalAIRerank` in `rag/llm/rerank_model.py` for a long time. Until this PR, a tenant who wanted to use LocalAI for reranking in the Go layer got `"not implemented"`. ### What this PR includes - `conf/models/localai.json`: add `"rerank": "rerank"` under `url_suffix` so the driver can build the URL from config. This matches the `URLSuffix.Rerank` field already used by aliyun and siliconflow. - `internal/entity/models/localai.go`: replace the `Rerank` stub with a real implementation that POSTs to `/v1/rerank`. Adds local request/response types `localAIRerankRequest` and `localAIRerankResponse`. No factory change. No interface change. ### How the implementation works - Validate the model name and resolve the tenant-supplied base URL with the existing `resolveBaseURL` helper. - Wrap the request with `context.WithTimeout(nonStreamCallTimeout)` so the call has a clear deadline. Same pattern `ChatWithMessages`, `ListModels`, and `Embed` already use in this file. - Only set the `Authorization` header when a non-empty API key was supplied. LocalAI accepts an empty key by default, so this preserves the optional-auth contract. - Default `top_n` to `len(documents)` when `rerankConfig.TopN == 0`, matching the existing Aliyun and SiliconFlow rerank implementations. - Validate every `results[].index` against `len(documents)`. If the upstream returns an out-of-range index, fail clearly instead of silently writing past the slice. - An empty `documents` slice returns `&RerankResponse{}` with no HTTP call. - Non-200 responses propagate the upstream status line and body. ### Note on stacking This PR builds on #14809 (LocalAI driver) and #14811 (LocalAI Embed). Until both merge, this PR's diff on GitHub will include all three commits. After #14809 and #14811 land on `main`, GitHub will auto-reduce this PR to only the `Rerank` changes (one commit, ~99 line diff in `localai.go` plus 1 line in `localai.json`). ### Type of change - [x] New Feature (non-breaking change which adds functionality) ### How was this tested? - `go build ./internal/entity/models/...` returns exit 0 on go 1.25 (the `go.mod` minimum). - The full method set on `LocalAIModel` still matches the `ModelDriver` interface. - Pattern parity with the existing Aliyun Rerank (`internal/entity/models/aliyun.go`) and SiliconFlow Rerank (`internal/entity/models/siliconflow.go`) implementations. Closes #14812 Depends on #14809, #14811 Tracking: #14736 Co-authored-by: Jin Hai <haijin.chn@gmail.com>
96 lines
3.0 KiB
Go
96 lines
3.0 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 (
|
|
"strings"
|
|
)
|
|
|
|
// ModelFactory creates ModelDriver instances based on provider name
|
|
type ModelFactory struct {
|
|
}
|
|
|
|
// NewModelFactory creates a new ModelFactory
|
|
func NewModelFactory() *ModelFactory {
|
|
return &ModelFactory{}
|
|
}
|
|
|
|
// CreateModelDriver creates a ModelDriver for the given provider and model
|
|
func (f *ModelFactory) CreateModelDriver(providerName string, baseURL map[string]string, urlSuffix URLSuffix) (ModelDriver, error) {
|
|
providerLower := strings.ToLower(providerName)
|
|
switch providerLower {
|
|
case "zhipu-ai":
|
|
return NewZhipuAIModel(baseURL, urlSuffix), nil
|
|
case "deepseek":
|
|
return NewDeepSeekModel(baseURL, urlSuffix), nil
|
|
case "moonshot":
|
|
return NewMoonshotModel(baseURL, urlSuffix), nil
|
|
case "minimax":
|
|
return NewMinimaxModel(baseURL, urlSuffix), nil
|
|
case "gitee":
|
|
return NewGiteeModel(baseURL, urlSuffix), nil
|
|
case "siliconflow":
|
|
return NewSiliconflowModel(baseURL, urlSuffix), nil
|
|
case "google":
|
|
return NewGoogleModel(baseURL, urlSuffix), nil
|
|
case "aliyun":
|
|
return NewAliyunModel(baseURL, urlSuffix), nil
|
|
case "volcengine":
|
|
return NewVolcEngine(baseURL, urlSuffix), nil
|
|
case "vllm":
|
|
return NewVllmModel(baseURL, urlSuffix), nil
|
|
case "xai":
|
|
return NewXAIModel(baseURL, urlSuffix), nil
|
|
case "lmstudio":
|
|
return NewLmStudioModel(baseURL, urlSuffix), nil
|
|
case "ollama":
|
|
return NewOllamaModel(baseURL, urlSuffix), nil
|
|
case "openai":
|
|
return NewOpenAIModel(baseURL, urlSuffix), nil
|
|
case "nvidia":
|
|
return NewNvidiaModel(baseURL, urlSuffix), nil
|
|
case "openrouter":
|
|
return NewOpenRouterModel(baseURL, urlSuffix), nil
|
|
case "huggingface":
|
|
return NewHuggingFaceModel(baseURL, urlSuffix), nil
|
|
case "baidu":
|
|
return NewBaiduModel(baseURL, urlSuffix), nil
|
|
case "cohere":
|
|
return NewCoHereModel(baseURL, urlSuffix), nil
|
|
case "fishaudio":
|
|
return NewFishAudioModel(baseURL, urlSuffix), nil
|
|
case "mistral":
|
|
return NewMistralModel(baseURL, urlSuffix), nil
|
|
case "upstage":
|
|
return NewUpstageModel(baseURL, urlSuffix), nil
|
|
case "stepfun":
|
|
return NewStepFunModel(baseURL, urlSuffix), nil
|
|
case "baichuan":
|
|
return NewBaichuanModel(baseURL, urlSuffix), nil
|
|
case "jina":
|
|
return NewJinaModel(baseURL, urlSuffix), nil
|
|
case "localai":
|
|
return NewLocalAIModel(baseURL, urlSuffix), nil
|
|
case "longcat":
|
|
return NewLongCatModel(baseURL, urlSuffix), nil
|
|
case "novita":
|
|
return NewNovitaModel(baseURL, urlSuffix), nil
|
|
default:
|
|
return NewDummyModel(baseURL, urlSuffix), nil
|
|
}
|
|
}
|