diff --git a/src/agents/chatbot/tools.py b/src/agents/chatbot/tools.py index a8f3cca3..0a7114cb 100644 --- a/src/agents/chatbot/tools.py +++ b/src/agents/chatbot/tools.py @@ -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 diff --git a/src/knowledge/base.py b/src/knowledge/base.py index b98e0ea1..a235d985 100644 --- a/src/knowledge/base.py +++ b/src/knowledge/base.py @@ -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(), diff --git a/src/knowledge/implementations/lightrag.py b/src/knowledge/implementations/lightrag.py index b940d480..09ccca37 100644 --- a/src/knowledge/implementations/lightrag.py +++ b/src/knowledge/implementations/lightrag.py @@ -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) # 执行查询 diff --git a/src/knowledge/utils/kb_utils.py b/src/knowledge/utils/kb_utils.py index 7b8c6bf8..d220d4bf 100644 --- a/src/knowledge/utils/kb_utils.py +++ b/src/knowledge/utils/kb_utils.py @@ -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}") diff --git a/src/storage/minio/__init__.py b/src/storage/minio/__init__.py index 0c8ce9af..bb795b57 100644 --- a/src/storage/minio/__init__.py +++ b/src/storage/minio/__init__.py @@ -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 # 为了向后兼容，导出常用的函数 diff --git a/src/storage/minio/client.py b/src/storage/minio/client.py index 4630c4be..afecf6d0 100644 --- a/src/storage/minio/client.py +++ b/src/storage/minio/client.py @@ -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