Merge branch 'main' of https://github.com/xerrors/Yuxi-Know
This commit is contained in:
Wenjie Zhang
commit
fd5fca0076
@ -42,10 +42,7 @@ async def text_to_img_qwen(text: str) -> str:
|
||||
|
||||
file_name = f"{uuid.uuid4()}.jpg"
|
||||
image_url = await aupload_file_to_minio(
|
||||
bucket_name="generated-images",
|
||||
file_name=file_name,
|
||||
data=file_data,
|
||||
file_extension="jpg"
|
||||
bucket_name="generated-images", file_name=file_name, data=file_data, file_extension="jpg"
|
||||
)
|
||||
logger.info(f"Image uploaded. URL: {image_url}")
|
||||
return image_url
|
||||
|
||||
@ -143,11 +143,12 @@ class KnowledgeBase(ABC):
|
||||
|
||||
# 创建数据库记录
|
||||
# 确保 Pydantic 模型被转换为字典，以便 JSON 序列化
|
||||
embed_info_dump = embed_info.model_dump() if hasattr(embed_info, "model_dump") else embed_info
|
||||
self.databases_meta[db_id] = {
|
||||
"name": database_name,
|
||||
"description": description,
|
||||
"kb_type": self.kb_type,
|
||||
"embed_info": embed_info.model_dump() if hasattr(embed_info, "model_dump") else embed_info,
|
||||
"embed_info": embed_info_dump,
|
||||
"llm_info": llm_info.model_dump() if hasattr(llm_info, "model_dump") else llm_info,
|
||||
"metadata": kwargs,
|
||||
"created_at": utc_isoformat(),
|
||||
|
||||
@ -201,9 +201,27 @@ class LightRagKB(KnowledgeBase):
|
||||
"""获取 embedding 函数"""
|
||||
config_dict = get_embedding_config(embed_info)
|
||||
|
||||
if config_dict["model_id"].startswith("ollama"):
|
||||
from lightrag.llm.ollama import ollama_embed
|
||||
|
||||
from src.utils import get_docker_safe_url
|
||||
|
||||
host = get_docker_safe_url(config_dict["base_url"].replace("/api/embed", ""))
|
||||
logger.debug(f"Ollama host: {host}")
|
||||
return EmbeddingFunc(
|
||||
embedding_dim=config_dict["dimension"],
|
||||
max_token_size=8192,
|
||||
func=lambda texts: ollama_embed(
|
||||
texts=texts,
|
||||
embed_model=config_dict["name"],
|
||||
api_key=config_dict["api_key"],
|
||||
host=host,
|
||||
),
|
||||
)
|
||||
|
||||
return EmbeddingFunc(
|
||||
embedding_dim=config_dict["dimension"],
|
||||
max_token_size=4096,
|
||||
max_token_size=8192,
|
||||
func=lambda texts: openai_embed(
|
||||
texts=texts,
|
||||
model=config_dict["model"],
|
||||
@ -365,12 +383,37 @@ class LightRagKB(KnowledgeBase):
|
||||
raise ValueError(f"Database {db_id} not found")
|
||||
|
||||
try:
|
||||
# QueryParam 支持的参数列表
|
||||
valid_params = {
|
||||
"mode",
|
||||
"only_need_context",
|
||||
"only_need_prompt",
|
||||
"response_type",
|
||||
"stream",
|
||||
"top_k",
|
||||
"chunk_top_k",
|
||||
"max_entity_tokens",
|
||||
"max_relation_tokens",
|
||||
"max_total_tokens",
|
||||
"hl_keywords",
|
||||
"ll_keywords",
|
||||
"conversation_history",
|
||||
"history_turns",
|
||||
"model_func",
|
||||
"user_prompt",
|
||||
"enable_rerank",
|
||||
"include_references",
|
||||
}
|
||||
|
||||
# 过滤 kwargs，只保留 QueryParam 支持的参数
|
||||
filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
|
||||
|
||||
# 设置查询参数
|
||||
params_dict = {
|
||||
"mode": "mix",
|
||||
"only_need_context": True,
|
||||
"top_k": 10,
|
||||
} | kwargs
|
||||
} | filtered_kwargs
|
||||
param = QueryParam(**params_dict)
|
||||
|
||||
# 执行查询
|
||||
|
||||
@ -7,6 +7,7 @@ import aiofiles
|
||||
from langchain_text_splitters import MarkdownTextSplitter
|
||||
|
||||
from src import config
|
||||
from src.config.static.models import EmbedModelInfo
|
||||
from src.utils import hashstr, logger
|
||||
from src.utils.datetime_utils import utc_isoformat
|
||||
|
||||
@ -293,19 +294,9 @@ def get_embedding_config(embed_info: dict) -> dict:
|
||||
if embed_info:
|
||||
# 优先检查是否有 model_id 字段
|
||||
if "model_id" in embed_info:
|
||||
from src.models.embed import select_embedding_model
|
||||
|
||||
model = select_embedding_model(embed_info["model_id"])
|
||||
config_dict["model"] = model.model
|
||||
config_dict["api_key"] = model.api_key
|
||||
config_dict["base_url"] = model.base_url
|
||||
config_dict["dimension"] = getattr(model, "dimension", 1024)
|
||||
elif hasattr(embed_info, "name"):
|
||||
# EmbedModelInfo 对象
|
||||
config_dict["model"] = embed_info.name
|
||||
config_dict["api_key"] = os.getenv(embed_info.api_key) or embed_info.api_key
|
||||
config_dict["base_url"] = embed_info.base_url
|
||||
config_dict["dimension"] = embed_info.dimension
|
||||
return config.embed_model_names[embed_info["model_id"]].model_dump()
|
||||
elif hasattr(embed_info, "name") and isinstance(embed_info, EmbedModelInfo):
|
||||
return embed_info.model_dump()
|
||||
else:
|
||||
# 字典形式（保持向后兼容）
|
||||
config_dict["model"] = embed_info["name"]
|
||||
@ -313,13 +304,7 @@ def get_embedding_config(embed_info: dict) -> dict:
|
||||
config_dict["base_url"] = embed_info["base_url"]
|
||||
config_dict["dimension"] = embed_info.get("dimension", 1024)
|
||||
else:
|
||||
from src.models import select_embedding_model
|
||||
|
||||
default_model = select_embedding_model(config.embed_model)
|
||||
config_dict["model"] = default_model.model
|
||||
config_dict["api_key"] = default_model.api_key
|
||||
config_dict["base_url"] = default_model.base_url
|
||||
config_dict["dimension"] = getattr(default_model, "dimension", 1024)
|
||||
return config.embed_model_names[config.embed_model].model_dump()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in get_embedding_config: {e}, {embed_info}")
|
||||
|
||||
@ -4,7 +4,7 @@ MinIO 存储模块
|
||||
"""
|
||||
|
||||
# 导出核心功能
|
||||
from .client import MinIOClient, StorageError, UploadResult, get_minio_client, aupload_file_to_minio
|
||||
from .client import MinIOClient, StorageError, UploadResult, aupload_file_to_minio, get_minio_client
|
||||
from .utils import generate_unique_filename, get_file_size
|
||||
|
||||
# 为了向后兼容，导出常用的函数
|
||||
|
||||
@ -288,8 +288,6 @@ def get_minio_client() -> MinIOClient:
|
||||
return _default_client
|
||||
|
||||
|
||||
|
||||
|
||||
async def aupload_file_to_minio(bucket_name: str, file_name: str, data: bytes, file_extension: str) -> str:
|
||||
"""
|
||||
通过字节上传文件到 MinIO的异步接口，根据输入的file_extension确定文件格式，并返回资源url
|
||||
|
||||
Loading…
Reference in New Issue
Block a user