Merge branch 'main' of https://github.com/xerrors/Yuxi-Know

Author: xerrors

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

src/knowledge/base.py

@@ -143,11 +143,12 @@ def create_database(
 
         # 创建数据库记录
         # 确保 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(),

src/knowledge/implementations/lightrag.py

@@ -201,9 +201,27 @@ def _get_embedding_func(self, embed_info: dict):
         """获取 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 @@ async def aquery(self, query_text: str, db_id: str, **kwargs) -> str:
             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)
 
             # 执行查询

src/knowledge/utils/kb_utils.py

@@ -7,6 +7,7 @@
 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,33 +294,17 @@ 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"]
                 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.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}")

src/storage/minio/__init__.py

@@ -4,7 +4,7 @@
 """
 
 # 导出核心功能
-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
 
 # 为了向后兼容，导出常用的函数

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