fix: 更新嵌入模型配置，修复相关信息获取逻辑
- 修改数据路由器中的嵌入模型信息获取方式，确保使用正确的配置项。 - 在 LightRagBasedKB 类中，增强 LLM 和嵌入函数的获取逻辑，支持动态配置。 - 更新嵌入模型和重排序模型的 URL 获取方式，统一为 base_url。 - 前端数据库视图中，调整嵌入模型选择的样式和信息展示，提升用户体验。
This commit is contained in:
Wenjie Zhang
parent
30163baed4
commit
a91d57f274
@ -29,7 +29,7 @@ async def create_database(
|
||||
):
|
||||
logger.debug(f"Create database {database_name}")
|
||||
try:
|
||||
embed_info = config.embed_model_info[embed_model_name]
|
||||
embed_info = config.embed_model_names[embed_model_name]
|
||||
database_info = knowledge_base.create_database(
|
||||
database_name,
|
||||
description,
|
||||
|
||||
@ -76,6 +76,11 @@ class LightRagBasedKB:
|
||||
if db_id not in self.databases_meta:
|
||||
return None
|
||||
|
||||
llm_info = self.databases_meta[db_id].get("llm_info", {})
|
||||
embed_info = self.databases_meta[db_id].get("embed_info", {})
|
||||
logger.info(f"LLM info: {llm_info}")
|
||||
logger.info(f"Embed info: {embed_info}")
|
||||
|
||||
# 创建 LightRAG 实例
|
||||
working_dir = os.path.join(self.work_dir, db_id)
|
||||
os.makedirs(working_dir, exist_ok=True)
|
||||
@ -84,8 +89,8 @@ class LightRagBasedKB:
|
||||
# 使用配置的 LLM 和 embedding 函数
|
||||
rag = LightRAG(
|
||||
working_dir=working_dir,
|
||||
llm_model_func=self._get_llm_func(),
|
||||
embedding_func=self._get_embedding_func(),
|
||||
llm_model_func=self._get_llm_func(llm_info),
|
||||
embedding_func=self._get_embedding_func(embed_info),
|
||||
vector_storage="MilvusVectorDBStorage",
|
||||
kv_storage="JsonKVStorage",
|
||||
graph_storage="PGGraphStorage",
|
||||
@ -110,30 +115,40 @@ class LightRagBasedKB:
|
||||
await rag.initialize_storages()
|
||||
await initialize_pipeline_status()
|
||||
|
||||
def _get_llm_func(self):
|
||||
def _get_llm_func(self, llm_info: dict):
|
||||
"""获取 LLM 函数"""
|
||||
llm_info = llm_info | {
|
||||
"model_name": "qwen3-1.7b",
|
||||
"provider": "dashscope"
|
||||
}
|
||||
provider_info = config.model_names[llm_info.get("provider")]
|
||||
api_key = os.getenv(provider_info.get("env")[0] or "OPENAI_API_KEY") or "no_api_key"
|
||||
base_url = get_docker_safe_url(provider_info.get("base_url", "http://localhost:8081/v1"))
|
||||
async def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
|
||||
return await openai_complete_if_cache(
|
||||
"qwen3:32b",
|
||||
llm_info.get("model_name"),
|
||||
prompt,
|
||||
system_prompt=system_prompt,
|
||||
history_messages=history_messages,
|
||||
api_key="no_api_key",
|
||||
base_url="http://172.19.13.7:8080/v1",
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
extra_body={"enable_thinking": False},
|
||||
**kwargs,
|
||||
)
|
||||
return llm_model_func
|
||||
|
||||
def _get_embedding_func(self):
|
||||
def _get_embedding_func(self, embed_info: dict):
|
||||
"""获取 embedding 函数"""
|
||||
api_key = os.getenv(embed_info.get("api_key", "OPENAI_API_KEY")) or "no_api_key"
|
||||
base_url = embed_info.get("base_url", "http://localhost:8081/v1").replace("/embeddings", "")
|
||||
return EmbeddingFunc(
|
||||
embedding_dim=1024,
|
||||
embedding_dim=embed_info.get("dimension") or 1024,
|
||||
max_token_size=4096,
|
||||
func=lambda texts: openai_embed(
|
||||
texts=texts,
|
||||
model="Qwen3-Embedding-0.6B",
|
||||
api_key="no_api_key",
|
||||
base_url=get_docker_safe_url("http://localhost:8081/v1")
|
||||
model=embed_info.get("model_name") or "Qwen3-Embedding-0.6B",
|
||||
api_key=api_key,
|
||||
base_url=get_docker_safe_url(base_url)
|
||||
),
|
||||
)
|
||||
|
||||
@ -248,6 +263,7 @@ class LightRagBasedKB:
|
||||
|
||||
def delete_database(self, db_id):
|
||||
"""删除数据库 - data_router.py 使用"""
|
||||
# TODO 删除数据库时，需要删除文件记录，并删除 LightRAG 中的文件
|
||||
if db_id in self.databases_meta:
|
||||
# 删除相关文件记录
|
||||
files_to_delete = [fid for fid, finfo in self.files_meta.items()
|
||||
@ -377,6 +393,7 @@ class LightRagBasedKB:
|
||||
|
||||
async def delete_file(self, db_id, file_id):
|
||||
"""删除文件 - data_router.py 使用"""
|
||||
# TODO 删除文件时，需要删除文件记录，并删除 LightRAG 中的文件
|
||||
rag = await self._get_lightrag_instance(db_id)
|
||||
if rag:
|
||||
try:
|
||||
|
||||
@ -133,7 +133,7 @@ class OllamaEmbedding(BaseEmbeddingModel):
|
||||
def __init__(self) -> None:
|
||||
self.info = config.embed_model_names[config.embed_model]
|
||||
self.model = self.info["name"]
|
||||
self.url = self.info.get("url", "http://localhost:11434/api/embed")
|
||||
self.url = self.info.get("base_url", "http://localhost:11434/api/embed")
|
||||
self.url = get_docker_safe_url(self.url)
|
||||
self.dimension = self.info.get("dimension", None)
|
||||
self.embed_model_fullname = config.embed_model
|
||||
@ -160,7 +160,7 @@ class OtherEmbedding(BaseEmbeddingModel):
|
||||
self.dimension = self.info.get("dimension", None)
|
||||
self.model = self.info["name"]
|
||||
self.api_key = os.getenv(self.info["api_key"], self.info["api_key"])
|
||||
self.url = get_docker_safe_url(self.info["url"])
|
||||
self.url = get_docker_safe_url(self.info["base_url"])
|
||||
assert self.url and self.model, f"URL and model are required. Cur embed model: {config.embed_model}"
|
||||
self.headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
|
||||
@ -25,7 +25,7 @@ def sigmoid(x):
|
||||
class OnlineRerank:
|
||||
def __init__(self, **kwargs):
|
||||
model_info = config.reranker_names[config.reranker]
|
||||
self.url = get_docker_safe_url(model_info["url"])
|
||||
self.url = get_docker_safe_url(model_info["base_url"])
|
||||
self.model = model_info["name"]
|
||||
|
||||
api_key = os.getenv(model_info["api_key"], model_info["api_key"])
|
||||
|
||||
@ -111,44 +111,36 @@ MODEL_NAMES:
|
||||
|
||||
|
||||
EMBED_MODEL_INFO:
|
||||
zhipu/zhipu-embedding-2:
|
||||
name: embedding-2
|
||||
dimension: 1024
|
||||
|
||||
zhipu/zhipu-embedding-3:
|
||||
name: embedding-3
|
||||
dimension: 2048
|
||||
|
||||
siliconflow/BAAI/bge-m3:
|
||||
name: BAAI/bge-m3
|
||||
dimension: 1024
|
||||
url: https://api.siliconflow.cn/v1/embeddings
|
||||
base_url: https://api.siliconflow.cn/v1/embeddings
|
||||
api_key: SILICONFLOW_API_KEY
|
||||
|
||||
vllm/Qwen/Qwen3-Embedding-0.6B:
|
||||
name: Qwen3-Embedding-0.6B
|
||||
dimension: 1024
|
||||
url: http://localhost:8081/v1/embeddings
|
||||
base_url: http://localhost:8081/v1/embeddings
|
||||
api_key: no_api_key
|
||||
|
||||
ollama/nomic-embed-text:
|
||||
name: nomic-embed-text
|
||||
url: http://localhost:11434/api/embed
|
||||
base_url: http://localhost:11434/api/embed
|
||||
dimension: 768
|
||||
|
||||
ollama/bge-m3:
|
||||
name: bge-m3
|
||||
url: http://localhost:11434/api/embed
|
||||
base_url: http://localhost:11434/api/embed
|
||||
dimension: 1024
|
||||
|
||||
RERANKER_LIST:
|
||||
|
||||
siliconflow/BAAI/bge-reranker-v2-m3:
|
||||
name: BAAI/bge-reranker-v2-m3
|
||||
url: https://api.siliconflow.cn/v1/rerank
|
||||
base_url: https://api.siliconflow.cn/v1/rerank
|
||||
api_key: SILICONFLOW_API_KEY
|
||||
|
||||
vllm/Qwen/Qwen3-Reranker-0.6B:
|
||||
name: Qwen/Qwen3-Reranker-0.6B
|
||||
url: http://localhost:8081/v1/rerank
|
||||
base_url: http://localhost:8081/v1/rerank
|
||||
api_key: no_api_key
|
||||
|
||||
@ -8,11 +8,11 @@
|
||||
</template>
|
||||
</HeaderComponent>
|
||||
|
||||
<a-modal :open="state.openNewDatabaseModel" title="新建知识库" @ok="createDatabase" @cancel="cancelCreateDatabase">
|
||||
<a-modal :open="state.openNewDatabaseModel" title="新建知识库" @ok="createDatabase" @cancel="cancelCreateDatabase" class="new-database-modal">
|
||||
<h3>知识库名称<span style="color: var(--error-color)">*</span></h3>
|
||||
<a-input v-model:value="newDatabase.name" placeholder="新建知识库名称" />
|
||||
<h3>嵌入模型</h3>
|
||||
<a-select v-model:value="newDatabase.embed_model_name" :options="embedModelOptions" />
|
||||
<a-select v-model:value="newDatabase.embed_model_name" :options="embedModelOptions" style="width: 100%;" />
|
||||
<h3 style="margin-top: 20px;">知识库描述</h3>
|
||||
<p style="color: var(--gray-700); font-size: 14px;">在智能体流程中，这里的描述会作为工具的描述。智能体会根据知识库的标题和描述来选择合适的工具。所以这里描述的越详细，智能体越容易选择到合适的工具。</p>
|
||||
<a-textarea
|
||||
@ -85,8 +85,8 @@ const state = reactive({
|
||||
})
|
||||
|
||||
const embedModelOptions = computed(() => {
|
||||
return Object.keys(configStore.config?.embed_model_info || {}).map(key => ({
|
||||
label: configStore.config?.embed_model_info[key]?.name || key,
|
||||
return Object.keys(configStore.config?.embed_model_names || {}).map(key => ({
|
||||
label: `${key} (${configStore.config?.embed_model_names[key]?.dimension})`,
|
||||
value: key,
|
||||
}))
|
||||
})
|
||||
@ -94,7 +94,7 @@ const embedModelOptions = computed(() => {
|
||||
const emptyEmbedInfo = {
|
||||
name: '',
|
||||
description: '',
|
||||
embed_model_name: '',
|
||||
embed_model_name: configStore.config?.embed_model,
|
||||
}
|
||||
|
||||
const newDatabase = reactive({
|
||||
@ -270,4 +270,10 @@ onMounted(() => {
|
||||
.database-container {
|
||||
padding: 0;
|
||||
}
|
||||
|
||||
.new-database-modal {
|
||||
h3 {
|
||||
margin-top: 10px;
|
||||
}
|
||||
}
|
||||
</style>
|
||||
|
||||
Loading…
Reference in New Issue
Block a user