From a6857c3e2131d333515e54469ef25feff97320ea Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Sat, 24 May 2025 11:29:45 +0800 Subject: [PATCH] style: fix ruff check --- pyproject.toml | 4 +- requirements.txt | 35 ----- scripts/vllm/main.py | 2 +- scripts/vllm/test_vllm.py | 2 +- server/db_manager.py | 2 +- server/models/__init__.py | 2 +- server/models/kb_models.py | 2 +- server/models/thread_model.py | 2 +- server/models/user_model.py | 2 +- server/routers/auth_router.py | 15 +- server/routers/base_router.py | 8 +- server/routers/chat_router.py | 21 ++- server/routers/data_router.py | 13 +- server/utils/__init__.py | 2 +- server/utils/auth_middleware.py | 7 +- src/__init__.py | 12 +- src/agents/__init__.py | 2 +- src/agents/chatbot/__init__.py | 2 +- src/agents/chatbot/graph.py | 1 - src/agents/react/__init__.py | 2 +- src/agents/react/workflows.py | 14 +- src/agents/registry.py | 25 ++-- src/agents/utils.py | 4 +- src/config/__init__.py | 15 +- src/core/__init__.py | 2 +- src/core/graphbase.py | 14 +- src/core/history.py | 5 +- src/core/knowledgebase.py | 23 ++- src/core/migrate_kb_to_sqlite.py | 6 +- src/core/retriever.py | 4 +- src/models/chat_model.py | 47 +----- src/models/embedding.py | 26 ++-- src/models/rerank_model.py | 15 +- src/plugins/__init__.py | 2 +- src/plugins/_ocr.py | 6 +- src/plugins/oneke.py | 249 ------------------------------- src/utils/__init__.py | 2 +- src/utils/db_migration.py | 2 +- test/test_concurrency.py | 13 +- 39 files changed, 134 insertions(+), 478 deletions(-) delete mode 100644 requirements.txt delete mode 100644 src/plugins/oneke.py diff --git a/pyproject.toml b/pyproject.toml index 87910759..0227f596 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,10 +41,10 @@ dependencies = [ ] [tool.ruff] line-length = 140 # 代码最大行宽 -select = [ # 选择的规则 +lint.select = [ # 选择的规则 "F", "E", "W", "UP", ] -ignore = ["F401"] # 忽略的规则 \ No newline at end of file +lint.ignore = ["F401"] # 忽略的规则 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index e0a7a137..00000000 --- a/requirements.txt +++ /dev/null @@ -1,35 +0,0 @@ -dashscope>=1.20.3 -PyMuPDF==1.23.26 -peft>=0.11.1 -FlagEmbedding>=1.3.2 -Flask>=3.0.3 -Flask_Cors>=4.0.1 -llama_index>=0.11.8 -openai>=1.44.1 -paddleocr>=2.8.1 -paddlepaddle>=2.4.0 -pandas>=2.2.2 -Pillow>=10.4.0 -pymilvus>=2.4.4 -python-dotenv>=1.0.1 -PyYAML>=6.0.2 -qianfan>=0.4.7 -torch>=2.4.0 -tqdm>=4.66.1 -zhipuai>=2.1.2 -neo4j>=5.22.0 -sentencepiece>=0.1.99 -llama-index-readers-file -opencv-python-headless -docx2txt -uvicorn[standard] -fastapi -python-multipart -tavily-python -rapidocr_onnxruntime -langchain -langsmith -langgraph -langchain-openai -langchain-community -PyJWT>=2.10.1 \ No newline at end of file diff --git a/scripts/vllm/main.py b/scripts/vllm/main.py index 40efaeaa..b9142ae8 100644 --- a/scripts/vllm/main.py +++ b/scripts/vllm/main.py @@ -17,4 +17,4 @@ outputs = llm.generate(prompts, sampling_params) for output in outputs: prompt = output.prompt generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") \ No newline at end of file + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") diff --git a/scripts/vllm/test_vllm.py b/scripts/vllm/test_vllm.py index 582ae272..204a4a76 100644 --- a/scripts/vllm/test_vllm.py +++ b/scripts/vllm/test_vllm.py @@ -15,4 +15,4 @@ chat_response = client.chat.completions.create( {"role": "user", "content": "Tell me a joke."}, ] ) -print("Chat response:", chat_response) \ No newline at end of file +print("Chat response:", chat_response) diff --git a/server/db_manager.py b/server/db_manager.py index 8e3fae30..ad7b4771 100644 --- a/server/db_manager.py +++ b/server/db_manager.py @@ -65,4 +65,4 @@ class DBManager: session.close() # 创建全局数据库管理器实例 -db_manager = DBManager() \ No newline at end of file +db_manager = DBManager() diff --git a/server/models/__init__.py b/server/models/__init__.py index 7c2377ae..860e5425 100644 --- a/server/models/__init__.py +++ b/server/models/__init__.py @@ -1,3 +1,3 @@ from sqlalchemy.ext.declarative import declarative_base -Base = declarative_base() \ No newline at end of file +Base = declarative_base() diff --git a/server/models/kb_models.py b/server/models/kb_models.py index eaa1c874..5ea399d3 100644 --- a/server/models/kb_models.py +++ b/server/models/kb_models.py @@ -103,4 +103,4 @@ class KnowledgeNode(Base): "start_char_idx": self.start_char_idx, "end_char_idx": self.end_char_idx, "metadata": self.meta_info or {} # 确保映射正确 - } \ No newline at end of file + } diff --git a/server/models/thread_model.py b/server/models/thread_model.py index 9bf45a75..10008373 100644 --- a/server/models/thread_model.py +++ b/server/models/thread_model.py @@ -18,4 +18,4 @@ class Thread(Base): update_at = Column(DateTime, default=func.now(), onupdate=func.now(), comment="更新时间") description = Column(String(255), nullable=True, comment="描述") - status = Column(Integer, default=1, comment="状态") \ No newline at end of file + status = Column(Integer, default=1, comment="状态") diff --git a/server/models/user_model.py b/server/models/user_model.py index 527d1f81..32a5ae88 100644 --- a/server/models/user_model.py +++ b/server/models/user_model.py @@ -52,4 +52,4 @@ class OperationLog(Base): "details": self.details, "ip_address": self.ip_address, "timestamp": self.timestamp.isoformat() if self.timestamp else None - } \ No newline at end of file + } diff --git a/server/routers/auth_router.py b/server/routers/auth_router.py index 08654bf8..179fb741 100644 --- a/server/routers/auth_router.py +++ b/server/routers/auth_router.py @@ -1,9 +1,8 @@ from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.security import OAuth2PasswordRequestForm from pydantic import BaseModel -from typing import List, Optional from sqlalchemy.orm import Session -from datetime import datetime, timedelta +from datetime import datetime from server.db_manager import db_manager from server.models.user_model import User, OperationLog @@ -27,16 +26,16 @@ class UserCreate(BaseModel): role: str = "user" class UserUpdate(BaseModel): - username: Optional[str] = None - password: Optional[str] = None - role: Optional[str] = None + username: str | None = None + password: str | None = None + role: str | None = None class UserResponse(BaseModel): id: int username: str role: str created_at: str - last_login: Optional[str] = None + last_login: str | None = None class InitializeAdmin(BaseModel): username: str @@ -202,7 +201,7 @@ async def create_user( return new_user.to_dict() # 路由：获取所有用户（管理员权限） -@auth.get("/users", response_model=List[UserResponse]) +@auth.get("/users", response_model=list[UserResponse]) async def read_users( skip: int = 0, limit: int = 100, @@ -344,4 +343,4 @@ async def delete_user( db.delete(user) db.commit() - return {"success": True, "message": "用户已删除"} \ No newline at end of file + return {"success": True, "message": "用户已删除"} diff --git a/server/routers/base_router.py b/server/routers/base_router.py index ea70877b..2cf61851 100644 --- a/server/routers/base_router.py +++ b/server/routers/base_router.py @@ -1,15 +1,13 @@ from fastapi import Request, Body, Depends from fastapi import APIRouter -from fastapi import Request, Body - -base = APIRouter() from src import config, retriever, knowledge_base, graph_base -from src.utils import logger from server.utils.auth_middleware import get_admin_user, get_superadmin_user from server.models.user_model import User +base = APIRouter() + @base.get("/") async def route_index(): return {"message": "You Got It!"} @@ -49,7 +47,7 @@ def get_log(current_user: User = Depends(get_admin_user)): from src.utils.logging_config import LOG_FILE from collections import deque - with open(LOG_FILE, 'r') as f: + with open(LOG_FILE) as f: last_lines = deque(f, maxlen=1000) log = ''.join(last_lines) diff --git a/server/routers/chat_router.py b/server/routers/chat_router.py index bcdded99..ad1e83ca 100644 --- a/server/routers/chat_router.py +++ b/server/routers/chat_router.py @@ -8,7 +8,6 @@ from fastapi import APIRouter, Body, Depends, HTTPException, Query from fastapi.responses import StreamingResponse from langchain_core.messages import AIMessageChunk, HumanMessage from sqlalchemy.orm import Session -from typing import List, Optional from pydantic import BaseModel from src import executor, config, retriever @@ -270,7 +269,7 @@ async def save_agent_config( if result: return {"success": True, "message": f"智能体 {agent_name} 配置已保存"} else: - raise HTTPException(status_code=500, detail=f"保存智能体配置失败") + raise HTTPException(status_code=500, detail="保存智能体配置失败") except Exception as e: logger.error(f"保存智能体配置出错: {e}, {traceback.format_exc()}") @@ -318,18 +317,18 @@ async def get_agent_config( # ==================== 线程管理 API ==================== class ThreadCreate(BaseModel): - title: Optional[str] = None + title: str | None = None agent_id: str - description: Optional[str] = None - metadata: Optional[dict] = None + description: str | None = None + metadata: dict | None = None class ThreadResponse(BaseModel): id: str user_id: str agent_id: str - title: Optional[str] - description: Optional[str] + title: str | None = None + description: str | None = None create_at: str update_at: str @@ -366,9 +365,9 @@ async def create_thread( } -@chat.get("/threads", response_model=List[ThreadResponse]) +@chat.get("/threads", response_model=list[ThreadResponse]) async def list_threads( - agent_id: Optional[str] = None, + agent_id: str | None = None, db: Session = Depends(get_db), current_user: User = Depends(get_required_user) ): @@ -420,8 +419,8 @@ async def delete_thread( class ThreadUpdate(BaseModel): - title: Optional[str] = None - description: Optional[str] = None + title: str | None = None + description: str | None = None @chat.put("/thread/{thread_id}", response_model=ThreadResponse) diff --git a/server/routers/data_router.py b/server/routers/data_router.py index 29584608..7c556b6b 100644 --- a/server/routers/data_router.py +++ b/server/routers/data_router.py @@ -1,7 +1,6 @@ import os import asyncio import traceback -from typing import List, Optional from fastapi import APIRouter, File, UploadFile, HTTPException, Depends, Body, Form, Query from src.utils import logger, hashstr @@ -25,7 +24,7 @@ async def get_databases(current_user: User = Depends(get_admin_user)): async def create_database( database_name: str = Body(...), description: str = Body(...), - dimension: Optional[int] = Body(None), + dimension: int | None = Body(None), current_user: User = Depends(get_admin_user) ): logger.debug(f"Create database {database_name}") @@ -53,19 +52,19 @@ async def query_test(query: str = Body(...), meta: dict = Body(...), current_use return result @data.post("/file-to-chunk") -async def file_to_chunk(files: List[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)): +async def file_to_chunk(files: list[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)): logger.debug(f"File to chunk: {files} {params=}") result = await knowledge_base.file_to_chunk(files, params=params) return result @data.post("/url-to-chunk") -async def url_to_chunk(urls: List[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)): +async def url_to_chunk(urls: list[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)): logger.debug(f"Url to chunk: {urls}") result = await knowledge_base.url_to_chunk(urls, params=params) return result @data.post("/add-by-file") -async def create_document_by_file(db_id: str = Body(...), files: List[str] = Body(...), current_user: User = Depends(get_admin_user)): +async def create_document_by_file(db_id: str = Body(...), files: list[str] = Body(...), current_user: User = Depends(get_admin_user)): logger.debug(f"Add document in {db_id} by file: {files}") try: await knowledge_base.add_files(db_id, files) @@ -113,7 +112,7 @@ async def get_document_info(db_id: str, file_id: str, current_user: User = Depen @data.post("/upload") async def upload_file( file: UploadFile = File(...), - db_id: Optional[str] = Query(None), + db_id: str | None = Query(None), current_user: User = Depends(get_admin_user) ): if not file.filename: @@ -170,7 +169,7 @@ async def get_graph_nodes(kgdb_name: str, num: int, current_user: User = Depends return {"result": graph_base.format_general_results(result), "message": "success"} @data.post("/graph/add-by-jsonl") -async def add_graph_entity(file_path: str = Body(...), kgdb_name: Optional[str] = Body(None), current_user: User = Depends(get_admin_user)): +async def add_graph_entity(file_path: str = Body(...), kgdb_name: str | None = Body(None), current_user: User = Depends(get_admin_user)): if not config.enable_knowledge_graph: return {"message": "知识图谱未启用", "status": "failed"} diff --git a/server/utils/__init__.py b/server/utils/__init__.py index 4901b718..00422248 100644 --- a/server/utils/__init__.py +++ b/server/utils/__init__.py @@ -1 +1 @@ -# utils包初始化文件 \ No newline at end of file +# utils包初始化文件 diff --git a/server/utils/auth_middleware.py b/server/utils/auth_middleware.py index 6bc69d04..6e917a12 100644 --- a/server/utils/auth_middleware.py +++ b/server/utils/auth_middleware.py @@ -1,5 +1,4 @@ -from typing import Optional, List, Callable -from fastapi import Depends, HTTPException, status, APIRouter +from fastapi import Depends, HTTPException, status from fastapi.security import OAuth2PasswordBearer from sqlalchemy.orm import Session from jose import JWTError, jwt @@ -30,7 +29,7 @@ def get_db(): db.close() # 获取当前用户 -async def get_current_user(token: Optional[str] = Depends(oauth2_scheme), db: Session = Depends(get_db)): +async def get_current_user(token: str | None = Depends(oauth2_scheme), db: Session = Depends(get_db)): credentials_exception = HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="无效的凭证", @@ -66,7 +65,7 @@ async def get_current_user(token: Optional[str] = Depends(oauth2_scheme), db: Se return user # 获取已登录用户（抛出401如果未登录） -async def get_required_user(user: Optional[User] = Depends(get_current_user)): +async def get_required_user(user: User | None = Depends(get_current_user)): if user is None: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, diff --git a/src/__init__.py b/src/__init__.py index c90b868d..ea1782fc 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -2,17 +2,17 @@ from dotenv import load_dotenv load_dotenv("src/.env") -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import ThreadPoolExecutor # noqa: E402 executor = ThreadPoolExecutor() -from src.config import Config +from src.config import Config # noqa: E402 config = Config() -from src.core import KnowledgeBase +from src.core import KnowledgeBase # noqa: E402 knowledge_base = KnowledgeBase() -from src.core import GraphDatabase +from src.core import GraphDatabase # noqa: E402 graph_base = GraphDatabase() -from src.core.retriever import Retriever -retriever = Retriever() \ No newline at end of file +from src.core.retriever import Retriever # noqa: E402 +retriever = Retriever() diff --git a/src/agents/__init__.py b/src/agents/__init__.py index 389b0049..9448ff6f 100644 --- a/src/agents/__init__.py +++ b/src/agents/__init__.py @@ -39,4 +39,4 @@ __all__ = ["agent_manager"] if __name__ == "__main__": - pass \ No newline at end of file + pass diff --git a/src/agents/chatbot/__init__.py b/src/agents/chatbot/__init__.py index 8ebda7da..ff0e6f84 100644 --- a/src/agents/chatbot/__init__.py +++ b/src/agents/chatbot/__init__.py @@ -1,4 +1,4 @@ from .graph import ChatbotAgent from .configuration import ChatbotConfiguration -__all__ = ["ChatbotAgent", "ChatbotConfiguration"] \ No newline at end of file +__all__ = ["ChatbotAgent", "ChatbotConfiguration"] diff --git a/src/agents/chatbot/graph.py b/src/agents/chatbot/graph.py index 9daf6c35..9d4ab4cd 100644 --- a/src/agents/chatbot/graph.py +++ b/src/agents/chatbot/graph.py @@ -63,7 +63,6 @@ class ChatbotAgent(BaseAgent): if self.graph: return self.graph - conf = self.config_schema.from_runnable_config(config_schema, agent_name=self.name) workflow = StateGraph(State, config_schema=self.config_schema) workflow.add_node("chatbot", self.llm_call) workflow.add_node("tools", ToolNode(tools=list(get_all_tools().values()))) diff --git a/src/agents/react/__init__.py b/src/agents/react/__init__.py index 3d2b2444..1fc79cca 100644 --- a/src/agents/react/__init__.py +++ b/src/agents/react/__init__.py @@ -1,4 +1,4 @@ from .graph import ReActAgent from .configuration import ReActConfiguration -__all__ = ["ReActAgent", "ReActConfiguration"] \ No newline at end of file +__all__ = ["ReActAgent", "ReActConfiguration"] diff --git a/src/agents/react/workflows.py b/src/agents/react/workflows.py index 5588d3a6..3356cc1c 100644 --- a/src/agents/react/workflows.py +++ b/src/agents/react/workflows.py @@ -2,24 +2,12 @@ import os from langchain_openai import ChatOpenAI from langgraph.checkpoint.memory import InMemorySaver +from langgraph.prebuilt import create_react_agent model = ChatOpenAI(model="glm-4-plus", api_key=os.getenv("ZHIPUAI_API_KEY"), base_url="https://open.bigmodel.cn/api/paas/v4/", temperature=0) - -# For this tutorial we will use custom tool that returns pre-defined values for weather in two cities (NYC & SF) - -from typing import Literal, Annotated - -from langchain_core.tools import tool, StructuredTool - - - tools = [] - - -from langgraph.prebuilt import create_react_agent - graph = create_react_agent(model, tools=tools, checkpointer=InMemorySaver()) diff --git a/src/agents/registry.py b/src/agents/registry.py index d7c0236e..5ef21964 100644 --- a/src/agents/registry.py +++ b/src/agents/registry.py @@ -4,8 +4,7 @@ import os import yaml import uuid from pathlib import Path -from typing import Type, Annotated, Optional, TypedDict -from enum import Enum +from typing import Annotated, TypedDict from abc import abstractmethod from dataclasses import dataclass, fields, field @@ -33,7 +32,7 @@ class Configuration(dict): @classmethod def from_runnable_config( - cls, config: Optional[RunnableConfig] = None, agent_name: str = None + cls, config: RunnableConfig | None = None, agent_name: str | None = None ) -> Configuration: """Create a Configuration instance from a RunnableConfig object. @@ -58,18 +57,18 @@ class Configuration(dict): # 合并三级配置，注意优先级 merged_config = {} - for field in _fields: + for config_field in _fields: # 1. 默认使用类默认值 - if hasattr(instance, field): - merged_config[field] = getattr(instance, field) + if hasattr(instance, config_field): + merged_config[config_field] = getattr(instance, config_field) # 2. 如果文件配置中有此字段，则覆盖 - if field in file_config: - merged_config[field] = file_config[field] + if config_field in file_config: + merged_config[config_field] = file_config[config_field] # 3. 如果运行时配置中有此字段，则覆盖 - if field in configurable: - merged_config[field] = configurable[field] + if config_field in configurable: + merged_config[config_field] = configurable[config_field] # 创建并返回配置实例 # logger.debug(f"合并配置: {merged_config}") @@ -82,7 +81,7 @@ class Configuration(dict): file_config = {} if os.path.exists(config_file_path): try: - with open(config_file_path, 'r', encoding='utf-8') as f: + with open(config_file_path, encoding='utf-8') as f: file_config = yaml.safe_load(f) or {} # logger.info(f"从文件加载智能体 {agent_name} 配置: {file_config}") except Exception as e: @@ -158,7 +157,7 @@ class Configuration(dict): }, ) -class BaseAgent(): +class BaseAgent: """ 定义一个基础 Agent 供 各类 graph 继承 @@ -245,4 +244,4 @@ class BaseAgent(): 必须确保在编译时设置 checkpointer，否则将无法获取历史记录。 例如: graph = workflow.compile(checkpointer=sqlite_checkpointer) """ - pass \ No newline at end of file + pass diff --git a/src/agents/utils.py b/src/agents/utils.py index f49168df..9872fbf7 100644 --- a/src/agents/utils.py +++ b/src/agents/utils.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC import asyncio from src.models import select_model @@ -57,5 +57,5 @@ async def agent_cli(agent: BaseAgent, config: RunnableConfig = None): print(f"Tool: {msg.content}") def get_cur_time_with_utc(): - return datetime.now(tz=timezone.utc).isoformat() + return datetime.now(tz=UTC).isoformat() diff --git a/src/config/__init__.py b/src/config/__init__.py index 5c47911a..c4e51b50 100644 --- a/src/config/__init__.py +++ b/src/config/__init__.py @@ -50,7 +50,7 @@ class Config(SimpleConfig): self.add_item("enable_reranker", default=False, des="是否开启重排序") self.add_item("enable_knowledge_base", default=False, des="是否开启知识库") self.add_item("enable_knowledge_graph", default=False, des="是否开启知识图谱") - self.add_item("enable_web_search", default=False, des="是否开启网页搜索（注：现阶段会根据 TAVILY_API_KEY 自动开启，无法手动配置，将会在下个版本移除此配置项）") + self.add_item("enable_web_search", default=False, des="是否开启网页搜索（注：现阶段会根据 TAVILY_API_KEY 自动开启，无法手动配置，将会在下个版本移除此配置项）") # noqa: E501 # 默认智能体配置 self.add_item("default_agent_id", default="", des="默认智能体ID") # 模型配置 @@ -60,7 +60,7 @@ class Config(SimpleConfig): self.add_item("model_name", default="Qwen/Qwen2.5-7B-Instruct", des="模型名称") self.add_item("embed_model", default="siliconflow/BAAI/bge-m3", des="Embedding 模型", choices=list(self.embed_model_names.keys())) - self.add_item("reranker", default="siliconflow/BAAI/bge-reranker-v2-m3", des="Re-Ranker 模型", choices=list(self.reranker_names.keys())) + self.add_item("reranker", default="siliconflow/BAAI/bge-reranker-v2-m3", des="Re-Ranker 模型", choices=list(self.reranker_names.keys())) # noqa: E501 self.add_item("model_local_paths", default={}, des="本地模型路径") self.add_item("use_rewrite_query", default="on", des="重写查询", choices=["off", "on", "hyde"]) self.add_item("device", default="cuda", des="运行本地模型的设备", choices=["cpu", "cuda"]) @@ -92,12 +92,12 @@ class Config(SimpleConfig): 从 models.yaml 和 models.private.yml 中更新 MODEL_NAMES """ - with open(Path("src/static/models.yaml"), 'r', encoding='utf-8') as f: + with open(Path("src/static/models.yaml"), encoding='utf-8') as f: _models = yaml.safe_load(f) # 尝试打开一个 models.private.yml 文件，用来覆盖 models.yaml 中的配置 try: - with open(Path("src/static/models.private.yml"), 'r', encoding='utf-8') as f: + with open(Path("src/static/models.private.yml"), encoding='utf-8') as f: _models_private = yaml.safe_load(f) except FileNotFoundError: _models_private = {} @@ -129,7 +129,8 @@ class Config(SimpleConfig): if os.path.exists(self.model_dir): logger.debug(f"MODEL_DIR （{self.model_dir}） 下面的文件夹: {os.listdir(self.model_dir)}") else: - logger.warning(f"提醒：MODEL_DIR （{self.model_dir}） 不存在，如果未配置，请忽略，如果配置了，请检查是否配置正确，比如 docker-compose 文件中的映射") + logger.warning(f"提醒：MODEL_DIR （{self.model_dir}） 不存在，如果未配置，请忽略，如果配置了，请检查是否配置正确;" + "比如 docker-compose 文件中的映射") # 检查模型提供商是否存在 @@ -177,7 +178,7 @@ class Config(SimpleConfig): if self.filename is not None and os.path.exists(self.filename): if self.filename.endswith(".json"): - with open(self.filename, 'r') as f: + with open(self.filename) as f: content = f.read() if content: local_config = json.loads(content) @@ -186,7 +187,7 @@ class Config(SimpleConfig): print(f"{self.filename} is empty.") elif self.filename.endswith(".yaml"): - with open(self.filename, 'r') as f: + with open(self.filename) as f: content = f.read() if content: local_config = yaml.safe_load(content) diff --git a/src/core/__init__.py b/src/core/__init__.py index 3a868409..b603d303 100644 --- a/src/core/__init__.py +++ b/src/core/__init__.py @@ -1,3 +1,3 @@ -from .history import * +from .history import HistoryManager from .knowledgebase import KnowledgeBase from .graphbase import GraphDatabase diff --git a/src/core/graphbase.py b/src/core/graphbase.py index 073c068a..9f923aed 100644 --- a/src/core/graphbase.py +++ b/src/core/graphbase.py @@ -189,7 +189,7 @@ class GraphDatabase: # 筛选出没有embedding的节点 nodes_without_embedding = session.execute_read(_get_nodes_without_embedding, all_entities) if not nodes_without_embedding: - logger.info(f"所有实体已有embedding，无需重新计算") + logger.info("所有实体已有embedding，无需重新计算") return logger.info(f"需要为{len(nodes_without_embedding)}/{len(all_entities)}个实体计算embedding") @@ -200,7 +200,11 @@ class GraphDatabase: for i in range(0, total_entities, max_batch_size): batch_entities = nodes_without_embedding[i:i+max_batch_size] - logger.debug(f"Processing entities batch {i//max_batch_size + 1}/{(total_entities-1)//max_batch_size + 1} ({len(batch_entities)} entities)") + logger.debug( + f"Processing entities batch " + f"{i//max_batch_size + 1}/{(total_entities-1)//max_batch_size + 1} " + f"({len(batch_entities)} entities)" + ) # 批量获取嵌入向量 batch_embeddings = await self.aget_embedding(batch_entities) @@ -221,7 +225,7 @@ class GraphDatabase: logger.info(f"Start adding entity to {kgdb_name} with {file_path}") def read_triples(file_path): - with open(file_path, 'r', encoding='utf-8') as file: + with open(file_path, encoding='utf-8') as file: for line in file: if line.strip(): yield json.loads(line.strip()) @@ -478,7 +482,7 @@ class GraphDatabase: try: graph_info = self.get_graph_info(graph_name) if graph_info is None: - logger.error(f"图数据库信息为空，无法保存") + logger.error("图数据库信息为空，无法保存") return False info_file_path = os.path.join(self.work_dir, "graph_info.json") @@ -521,7 +525,7 @@ class GraphDatabase: logger.debug(f"图数据库信息文件不存在：{info_file_path}") return False - with open(info_file_path, 'r', encoding='utf-8') as f: + with open(info_file_path, encoding='utf-8') as f: graph_info = json.load(f) # 更新对象属性 diff --git a/src/core/history.py b/src/core/history.py index 6e5c597e..17676145 100644 --- a/src/core/history.py +++ b/src/core/history.py @@ -1,7 +1,6 @@ from src.utils.prompts import get_system_prompt -from src.utils.logging_config import logger -class HistoryManager(): +class HistoryManager: def __init__(self, history=None, system_prompt=None): self.messages = [] @@ -47,4 +46,4 @@ class HistoryManager(): for message in self.messages: msg = message["content"].replace('\n', ' ') history_str += f"\n{message['role']}: {msg}" - return history_str \ No newline at end of file + return history_str diff --git a/src/core/knowledgebase.py b/src/core/knowledgebase.py index 951d75ee..30fac731 100644 --- a/src/core/knowledgebase.py +++ b/src/core/knowledgebase.py @@ -73,11 +73,11 @@ class KnowledgeBase: return from src.models.embedding import get_embedding_model - self.embed_model = get_embedding_model(config) + self.embed_model = get_embedding_model() if config.enable_reranker: from src.models.rerank_model import get_reranker - self.reranker = get_reranker(config) + self.reranker = get_reranker() if not self.connect_to_milvus(): raise ConnectionError("Failed to connect to Milvus") @@ -472,10 +472,7 @@ class KnowledgeBase: async def add_chunks(self, db_id, file_chunks): """添加分块""" db = self.get_kb_by_id(db_id) - - if db["embed_model"] != self.embed_model.embed_model_fullname: - logger.error(f"Embed model not match, {db['embed_model']} != {self.embed_model.embed_model_fullname}") - return {"message": f"Embed model not match, cur: {self.embed_model.embed_model_fullname}, req: {db['embed_model']}", "status": "failed"} + assert self.check_embed_model(db_id), f"Embed model not match, {db['embed_model']} != {self.embed_model.embed_model_fullname}" for file_id, chunk_info in file_chunks.items(): # 在数据库中创建文件记录 @@ -505,10 +502,7 @@ class KnowledgeBase: async def add_files(self, db_id, files, params=None): db = self.get_kb_by_id(db_id) - - if not self.check_embed_model(db_id): - logger.error(f"Embed model not match, {db['embed_model']} != {self.embed_model.embed_model_fullname}") - return {"message": f"Embed model not match, cur: {self.embed_model.embed_model_fullname}, req: {db['embed_model']}", "status": "failed"} + assert self.check_embed_model(db_id), f"Embed model not match, {db['embed_model']} != {self.embed_model.embed_model_fullname}" # Preprocessing the files to the queue new_files = await self.file_to_chunk(files, params=params) @@ -625,11 +619,11 @@ class KnowledgeBase: "embed_model": db["embed_model"], } else: - logger.warning(( + logger.warning( f"无法将知识库 {db['name']} 转换为 Tools, 因为向量模型不匹配，" f"当前向量模型: {self.embed_model.embed_model_fullname}，" f"知识库向量模型: {db['embed_model']}。" - )) + ) return retrievers ################################ @@ -648,7 +642,8 @@ class KnowledgeBase: logger.info(f"Successfully connected to Milvus at {uri}") return True except MilvusException as e: - logger.error(f"Failed to connect to Milvus: {e}，请检查 milvus 的容器是否正常运行，如果已退出，请重新启动 `docker restart milvus-standalone-dev`") + logger.error(f"Failed to connect to Milvus: {e}，请检查 milvus 的容器是否正常运行。") + logger.error("如果已退出，请重新启动 `docker restart milvus-standalone-dev`。") return False def get_collection_names(self): @@ -781,4 +776,4 @@ def gen_filename_from_url(url): if len(filename) > 100: filename = filename[:97] + "..." filename = filename.replace('/', '_') - return filename \ No newline at end of file + return filename diff --git a/src/core/migrate_kb_to_sqlite.py b/src/core/migrate_kb_to_sqlite.py index 6b557642..6e3686d5 100644 --- a/src/core/migrate_kb_to_sqlite.py +++ b/src/core/migrate_kb_to_sqlite.py @@ -21,7 +21,7 @@ def migrate_json_to_sqlite(): try: # 读取旧的JSON文件 - with open(json_path, 'r', encoding='utf-8') as f: + with open(json_path, encoding='utf-8') as f: old_data = json.load(f) if not old_data or not isinstance(old_data, dict) or not old_data.get("databases"): @@ -36,7 +36,7 @@ def migrate_json_to_sqlite(): logger.info(f"正在迁移知识库: {db_info.get('name')} (ID: {db_id})") # 创建知识库 - db_dict = db_manager.create_database( + db_manager.create_database( db_id=db_id, name=db_info.get('name', '未命名知识库'), description=db_info.get('description', ''), @@ -81,4 +81,4 @@ def migrate_json_to_sqlite(): except Exception as e: logger.error(f"迁移过程中出错: {e}") - return False \ No newline at end of file + return False diff --git a/src/core/retriever.py b/src/core/retriever.py index 0addb0c6..a399979a 100644 --- a/src/core/retriever.py +++ b/src/core/retriever.py @@ -13,7 +13,7 @@ class Retriever: def _load_models(self): if config.enable_reranker: - self.reranker = get_reranker(config) + self.reranker = get_reranker() if config.enable_web_search: from src.utils.web_search import WebSearcher @@ -168,7 +168,7 @@ class Retriever: entities = [] if refs["meta"].get("use_graph"): from src.utils.prompts import entity_extraction_prompt_template as entity_template - from src.utils.prompts import keywords_prompt_template as entity_template + # from src.utils.prompts import keywords_prompt_template as entity_templat|e entity_extraction_prompt = entity_template.format(text=query) entities = model.predict(entity_extraction_prompt).content.split("<->") diff --git a/src/models/chat_model.py b/src/models/chat_model.py index 5699a20d..5fca3fdb 100644 --- a/src/models/chat_model.py +++ b/src/models/chat_model.py @@ -4,7 +4,7 @@ from openai import OpenAI from src.utils import logger, get_docker_safe_url from langchain_openai import ChatOpenAI -class OpenAIBase(): +class OpenAIBase: def __init__(self, api_key, base_url, model_name, chat_open_ai=None, **kwargs): self.api_key = api_key self.base_url = base_url @@ -61,49 +61,6 @@ class OpenAIBase(): logger.error(f"Error getting models: {e}") return [] - def _get_model_by_model_url(self, model_url): - """ - Refs: https://docs.together.ai/reference/models-1 - - Return: [ - { - "id": "meta-llama/Meta-Llama-3-70B-Instruct-Turbo", - "object": "model", - "created": 0, - "type": "chat", - "running": false, - "display_name": "Meta Llama 3 70B Instruct Turbo", - "organization": "Meta", - "link": "https://huggingface.co/meta-llama/Meta-Llama-3-70B-Instruct", - "license": "Llama-3 (Other)", - "context_length": 8192, - "config": { - "chat_template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}", - "stop": [ - "<|eot_id|>" - ], - "bos_token": "<|begin_of_text|>", - "eos_token": "<|end_of_text|>" - }, - "pricing": { - "hourly": 0, - "input": 0.88, - "output": 0.88, - "base": 0, - "finetune": 0 - } - }, - ] - - """ - headers = { - "accept": "application/json", - "authorization": f"Bearer {self.api_key}" - } - response = requests.get(model_url, headers=headers) - return response.json() - - class OpenModel(OpenAIBase): def __init__(self, model_name=None): @@ -214,4 +171,4 @@ class DashScope(OpenAIBase): if __name__ == "__main__": - pass \ No newline at end of file + pass diff --git a/src/models/embedding.py b/src/models/embedding.py index 4fff41fb..7fb16b02 100644 --- a/src/models/embedding.py +++ b/src/models/embedding.py @@ -2,7 +2,6 @@ import os import json import requests import asyncio -import json from abc import abstractmethod from zhipuai import ZhipuAI from langchain_huggingface import HuggingFaceEmbeddings @@ -68,7 +67,7 @@ class BaseEmbeddingModel: return data class LocalEmbeddingModel(BaseEmbeddingModel): - def __init__(self, config, **kwargs): + def __init__(self, **kwargs): info = config.embed_model_names[config.embed_model] self.model = config.model_local_paths.get(info["name"], info.get("local_path")) @@ -104,7 +103,7 @@ class LocalEmbeddingModel(BaseEmbeddingModel): return await self.model.aembed_documents(message) def encode_queries(self, queries): - logger.warning(f"Huggingface Model 不支持批量 encode queries，因此使用训练实现") + logger.warning("Huggingface Model 不支持批量 encode queries，因此使用训练实现") data = [] for q in queries: data.append(self.predict(q)) @@ -114,7 +113,7 @@ class LocalEmbeddingModel(BaseEmbeddingModel): class ZhipuEmbedding(BaseEmbeddingModel): - def __init__(self, config) -> None: + def __init__(self) -> None: self.config = config self.model = config.embed_model_names[config.embed_model]["name"] self.dimension = config.embed_model_names[config.embed_model]["dimension"] @@ -131,7 +130,7 @@ class ZhipuEmbedding(BaseEmbeddingModel): class OllamaEmbedding(BaseEmbeddingModel): - def __init__(self, config) -> None: + 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") @@ -155,7 +154,7 @@ class OllamaEmbedding(BaseEmbeddingModel): class OtherEmbedding(BaseEmbeddingModel): - def __init__(self, config) -> None: + def __init__(self) -> None: self.info = config.embed_model_names[config.embed_model] self.embed_model_fullname = config.embed_model self.dimension = self.info.get("dimension", None) @@ -182,27 +181,28 @@ class OtherEmbedding(BaseEmbeddingModel): "input": message, } -def get_embedding_model(config): +def get_embedding_model(): if not config.enable_knowledge_base: return None provider, model_name = config.embed_model.split('/', 1) - assert config.embed_model in config.embed_model_names.keys(), f"Unsupported embed model: {config.embed_model}, only support {config.embed_model_names.keys()}" + support_embed_models = config.embed_model_names.keys() + assert config.embed_model in support_embed_models, f"Unsupported embed model: {config.embed_model}, only support {support_embed_models}" logger.debug(f"Loading embedding model {config.embed_model}") if provider == "local": - model = LocalEmbeddingModel(config) + model = LocalEmbeddingModel() elif provider == "zhipu": - model = ZhipuEmbedding(config) + model = ZhipuEmbedding() elif provider == "ollama": - model = OllamaEmbedding(config) + model = OllamaEmbedding() else: - model = OtherEmbedding(config) + model = OtherEmbedding() return model def handle_local_model(paths, model_name, default_path): model_path = paths.get(model_name, default_path) - return model_path \ No newline at end of file + return model_path diff --git a/src/models/rerank_model.py b/src/models/rerank_model.py index 145ad2b4..60f6948b 100644 --- a/src/models/rerank_model.py +++ b/src/models/rerank_model.py @@ -9,7 +9,7 @@ from src.utils.logging_config import logger class LocalReranker(FlagReranker): - def __init__(self, config, **kwargs): + def __init__(self, **kwargs): model_info = config.reranker_names[config.reranker] model_name_or_path = config.model_local_paths.get(model_info["name"], model_info.get("local_path")) model_name_or_path = model_name_or_path or model_info["name"] @@ -22,8 +22,8 @@ class LocalReranker(FlagReranker): def sigmoid(x): return 1 / (1 + np.exp(-x)) -class SiliconFlowReranker(): - def __init__(self, config, **kwargs): +class SiliconFlowReranker: + def __init__(self, **kwargs): self.url = "https://api.siliconflow.cn/v1/rerank" self.model = config.reranker_names[config.reranker]["name"] @@ -58,13 +58,14 @@ class SiliconFlowReranker(): "max_chunks_per_doc": max_length, } -def get_reranker(config): - assert config.reranker in config.reranker_names.keys(), f"Unsupported Reranker: {config.reranker}, only support {config.reranker_names.keys()}" +def get_reranker(): + support_rerankers = config.reranker_names.keys() + assert config.reranker in support_rerankers, f"Unsupported Reranker: {config.reranker}, only support {support_rerankers}" provider, model_name = config.reranker.split('/', 1) if provider == "local": - return LocalReranker(config) + return LocalReranker() elif provider == "siliconflow": - return SiliconFlowReranker(config) + return SiliconFlowReranker() else: raise ValueError(f"Unsupported Reranker: {config.reranker}, only support {config.reranker_names.keys()}") diff --git a/src/plugins/__init__.py b/src/plugins/__init__.py index 694cab87..5c02ab3c 100644 --- a/src/plugins/__init__.py +++ b/src/plugins/__init__.py @@ -2,4 +2,4 @@ from src.plugins._ocr import OCRPlugin ocr = OCRPlugin() -__all__ = ["ocr"] \ No newline at end of file +__all__ = ["ocr"] diff --git a/src/plugins/_ocr.py b/src/plugins/_ocr.py index fc42840e..b19aea36 100644 --- a/src/plugins/_ocr.py +++ b/src/plugins/_ocr.py @@ -24,7 +24,7 @@ class OCRPlugin: def load_model(self): """加载 OCR 模型""" - logger.info(f"加载 OCR 模型，仅在第一次调用时加载") + logger.info("加载 OCR 模型，仅在第一次调用时加载") model_dir = os.path.join(os.getenv("MODEL_DIR", ""), "SWHL/RapidOCR") det_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_det_infer.onnx") rec_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx") @@ -75,7 +75,7 @@ class OCRPlugin: text = '\n'.join([line[1] for line in result]) return text else: - logger.warning(f"OCR未能识别出文本内容") + logger.warning("OCR未能识别出文本内容") return "" except Exception as e: @@ -214,7 +214,7 @@ def plainreader(file_path): """读取普通文本文件并返回text文本""" assert os.path.exists(file_path), "File not found" - with open(file_path, "r") as f: + with open(file_path) as f: text = f.read() return text diff --git a/src/plugins/oneke.py b/src/plugins/oneke.py deleted file mode 100644 index 03bceacf..00000000 --- a/src/plugins/oneke.py +++ /dev/null @@ -1,249 +0,0 @@ -# readmore: https://github.com/zjunlp/DeepKE/blob/main/example/llm/OneKE.md -import os -import json -import torch -import dotenv -from transformers import ( - AutoConfig, - AutoTokenizer, - AutoModelForCausalLM, - GenerationConfig, - BitsAndBytesConfig -) - -from src.utils import logger - -dotenv.load_dotenv() - -instruction_mapper = { - 'NERzh': "你是专门进行实体抽取的专家。请从input中抽取出符合schema定义的实体，不存在的实体类型返回空列表。请按照JSON字符串的格式回答。", - 'REzh': "你是专门进行关系抽取的专家。请从input中抽取出符合schema定义的关系三元组。请按照JSON字符串的格式回答。", - 'EEzh': "你是专门进行事件提取的专家。请从input中抽取出符合schema定义的事件，不存在的事件返回空列表，不存在的论元返回NAN，如果论元存在多值请返回列表。请按照JSON字符串的格式回答。", - 'EETzh': "你是专门进行事件提取的专家。请从input中抽取出符合schema定义的事件类型及事件触发词，不存在的事件返回空列表。请按照JSON字符串的格式回答。", - 'EEAzh': "你是专门进行事件论元提取的专家。请从input中抽取出符合schema定义的事件论元及论元角色，不存在的论元返回NAN或空字典，如果论元存在多值请返回列表。请按照JSON字符串的格式回答。", - 'KGzh': '你是一个图谱实体知识结构化专家。根据输入实体类型(entity type)的schema描述，从文本中抽取出相应的实体实例和其属性信息，不存在的属性不输出, 属性存在多值就返回列表，并输出为可解析的json格式。', - 'NERen': "You are an expert in named entity recognition. Please extract entities that match the schema definition from the input. Return an empty list if the entity type does not exist. Please respond in the format of a JSON string.", - 'REen': "You are an expert in relationship extraction. Please extract relationship triples that match the schema definition from the input. Return an empty list for relationships that do not exist. Please respond in the format of a JSON string.", - 'EEen': "You are an expert in event extraction. Please extract events from the input that conform to the schema definition. Return an empty list for events that do not exist, and return NAN for arguments that do not exist. If an argument has multiple values, please return a list. Respond in the format of a JSON string.", - 'EETen': "You are an expert in event extraction. Please extract event types and event trigger words from the input that conform to the schema definition. Return an empty list for non-existent events. Please respond in the format of a JSON string.", - 'EEAen': "You are an expert in event argument extraction. Please extract event arguments and their roles from the input that conform to the schema definition, which already includes event trigger words. If an argument does not exist, return NAN or an empty dictionary. Please respond in the format of a JSON string.", - 'KGen': 'You are an expert in structured knowledge systems for graph entities. Based on the schema description of the input entity type, you extract the corresponding entity instances and their attribute information from the text. Attributes that do not exist should not be output. If an attribute has multiple values, a list should be returned. The results should be output in a parsable JSON format.', -} - -split_num_mapper = { - 'NER':6, 'RE':4, 'EE':4, 'EET':4, 'EEA':4, 'KG':1 -} - -class OneKE: - - def __init__(self, config=None): - - self.config = config - model_name_or_path = config.model_local_paths.get('zjunlp/OneKE', "zjunlp/OneKE") - logger.info(f"Loading KGC model OneKE from {model_name_or_path}") - - model_config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True) - self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True) - - self.quantization_config=BitsAndBytesConfig( - load_in_4bit=True, - llm_int8_threshold=6.0, - llm_int8_has_fp16_weight=False, - bnb_4bit_compute_dtype=torch.bfloat16, - bnb_4bit_use_double_quant=True, - bnb_4bit_quant_type="nf4", - ) - - self.generate_config = GenerationConfig( - max_length=1024, - max_new_tokens=512, - return_dict_in_generate=True - ) - - self.model = AutoModelForCausalLM.from_pretrained( - model_name_or_path, - config=model_config, - device_map="auto", - quantization_config=self.quantization_config, - torch_dtype=torch.bfloat16, - trust_remote_code=True, - ) - self.model.eval() - - def construct_input(self, text, schema, task, language="zh", use_split=False): - if use_split: - split_num = split_num_mapper[task] - if isinstance(schema, dict): - schema_keys = list(schema.keys()) - schema_keys = [schema_keys[i:i+split_num] for i in range(0, len(schema_keys), split_num)] - schema = [{k: schema[k] for k in keys} for keys in schema_keys] - else: - schema = [schema[i:i+split_num] for i in range(0, len(schema), split_num)] - else: - schema = [schema] - - sintructs = [] - for s in schema: - sintruct = json.dumps({ - "instruction": instruction_mapper[task+language], - "schema": s, - "input": text, - }, ensure_ascii=False) - sintructs.append(sintruct) - - return sintructs - - def predict(self, text, schema, task, language="zh", use_split=False): - sintructs = self.construct_input(text, schema, task, language, use_split) - outputs = [] - for sintruct in sintructs: - input_ids = self.tokenizer.encode(sintruct, return_tensors="pt").to(self.model.device) - input_length = input_ids.size(1) - - generation_output = self.model.generate( - input_ids=input_ids, - generation_config=self.generate_config, - pad_token_id=self.tokenizer.eos_token_id - ) - generation_output = generation_output.sequences[0] - generation_output = generation_output[input_length:] - output = self.tokenizer.decode(generation_output, skip_special_tokens=True) - - outputs.append(output) - - return outputs - - def processing_text_to_kg(self, text_or_path, output_path): - for chunk in read_and_process_chars(text_or_path): - - text = chunk - schema = [ - { - "entity_type": "食品", - "attributes": { - "名称": "食品的名称，包括品牌名、通用名称或专业化学名", - "分类": "食品所属的类型，例如水果、蔬菜、肉类、谷物、调料、添加剂、益生菌等", - "成分": "食品的主要成分，详细列出包括天然成分、添加剂、保鲜剂、营养强化剂等", - "营养价值": "食品的营养成分，概括其提供的能量和主要营养素，如蛋白质、脂肪、碳水化合物、维生素和矿物质", - "加工方式": "食品的处理或制备方法，包括日常烹饪、加工处理及实验室制备方式等", - "作用或食用效果": "食品对健康或身体的影响，可能的功效或用途" - } - } - ] - task = "KG" - output = self.predict(text=text, schema=schema, task=task, language="zh") - formatted_output = parse_and_format_output(output=output, task_type=task) - - with open(output_path, 'a+', encoding='utf-8') as f: - for entry in formatted_output: - f.write(json.dumps(entry, ensure_ascii=False) + '\n') - - print(f"预测结果已添加到 {output_path} 文件中。") - return output_path - -def read_and_process_chars(file_path, char_size=512, overlap_size=100): - buffer = "" - with open(file_path, 'r', encoding='utf-8') as file: - while True: - chunk = file.read(char_size) - if not chunk: # 文件读取完毕 - if buffer: - yield buffer - break - chunk = chunk.replace('\n', '').replace('\r', '') # 去除换行符 - buffer += chunk - while len(buffer) >= char_size: - yield buffer[:char_size] - buffer = buffer[char_size - overlap_size:] - -def parse_and_format_output(output, task_type): - formatted_output = [] - - for entry in output: - try: - # Check if the entry is a valid JSON string - if isinstance(entry, str) and entry.strip().startswith('{') and entry.strip().endswith('}'): - parsed_entry = json.loads(entry) - - if task_type == "KG": - for entity_type, entities in parsed_entry.items(): - for entity_name, attributes in entities.items(): - for attribute_name, attribute_values in attributes.items(): - if isinstance(attribute_values, list): - for value in attribute_values: - formatted_output.append({ - "h": entity_name, - "t": value, - "r": attribute_name - }) - else: - formatted_output.append({ - "h": entity_name, - "t": attribute_values, - "r": attribute_name - }) - elif task_type == "RE": - for relation_type, pairs in parsed_entry.items(): - for pair in pairs: - formatted_output.append({ - "h": pair["subject"], - "t": pair["object"], - "r": relation_type - }) - else: - raise json.JSONDecodeError("Invalid JSON format", entry, 0) - except json.JSONDecodeError as e: - print(f"JSONDecodeError: {e} - Skipping entry") - continue - except TypeError as e: - print(f"TypeError: {e} - Skipping entry") - continue - except AttributeError as e: - print(f"AttributeError: {e} - Skipping entry") - continue - - return formatted_output - - -if __name__ == "__main__": - oneke = OneKE() - oneke.processing_text_to_kg("asdasdadadsad", 'kg.jsonl') - - file_path = '' - output_path = '' - text = "" - task = "KG" - for chunk in read_and_process_chars(file_path): - - text = chunk - - # schema = { - # "定义": "描述食品的起源、传统制作方法、文化象征意义或者描述食品或相关事物的定义或含义。包括食品的来源、特点及其在特定文化或背景下的意义。", - # "组成成分": "描述食品的组成部分或成分，包括主要成分、微量成分、添加剂等，揭示食品的化学或物理组成。", - # "功能": "描述食品或其成分的功能或作用，包括其对人体健康的影响、在烹饪中的用途、药用价值等。", - # "属性": "描述食品或其成分的特性或属性，揭示其独特的营养价值、口感特点、保存方式等，包括食品的营养价值、口感特点（如口感丰富、清淡等）、保存方式（如冷藏、冷冻、干燥等）", - # "种类": "描述食品或相关事物的种类或类别，揭示其分类体系、不同类型的特点及其在具体应用中的区别。" - # } - - schema = [ - { - "entity_type": "食品", - "attributes": { - "名称": "食品的名称，包括品牌名、通用名称或专业化学名", - "分类": "食品所属的类型，例如水果、蔬菜、肉类、谷物、调料、添加剂、益生菌等", - "成分": "食品的主要成分，详细列出包括天然成分、添加剂、保鲜剂、营养强化剂等", - "营养价值": "食品的营养成分，概括其提供的能量和主要营养素，如蛋白质、脂肪、碳水化合物、维生素和矿物质", - "加工方式": "食品的处理或制备方法，包括日常烹饪、加工处理及实验室制备方式等", - "作用或食用效果": "食品对健康或身体的影响，可能的功效或用途" - } - } - ] - - output = oneke.predict(text=text, schema=schema, task=task, language="zh") - formatted_output = parse_and_format_output(output=output, task_type=task) - - with open(output_path, 'a+', encoding='utf-8') as f: - for entry in formatted_output: - f.write(json.dumps(entry, ensure_ascii=False) + '\n') - - print(f"预测结果已添加到 {output_path} 文件中。") \ No newline at end of file diff --git a/src/utils/__init__.py b/src/utils/__init__.py index c9dc555d..eb416c3b 100644 --- a/src/utils/__init__.py +++ b/src/utils/__init__.py @@ -38,4 +38,4 @@ def get_docker_safe_url(base_url): base_url = base_url.replace("http://localhost", "http://host.docker.internal") base_url = base_url.replace("http://127.0.0.1", "http://host.docker.internal") logger.info(f"Running in docker, using {base_url} as base url") - return base_url \ No newline at end of file + return base_url diff --git a/src/utils/db_migration.py b/src/utils/db_migration.py index c64805ee..40c966ff 100644 --- a/src/utils/db_migration.py +++ b/src/utils/db_migration.py @@ -99,4 +99,4 @@ def migrate_knowledge_db(): if __name__ == "__main__": # 当作为脚本运行时执行迁移 result = migrate_knowledge_db() - print(f"迁移{'成功' if result else '失败'}") \ No newline at end of file + print(f"迁移{'成功' if result else '失败'}") diff --git a/test/test_concurrency.py b/test/test_concurrency.py index 2c4c4b08..0061d3bb 100644 --- a/test/test_concurrency.py +++ b/test/test_concurrency.py @@ -1,8 +1,7 @@ import asyncio import aiohttp import time -import json -from typing import List + async def make_request(session: aiohttp.ClientSession, request_id: int) -> dict: """发送单个请求到API""" @@ -17,6 +16,7 @@ async def make_request(session: aiohttp.ClientSession, request_id: int) -> dict: try: async with session.post(url, json=payload) as response: result = await response.json() + print(f"请求 {request_id} 结果: {result}") end_time = time.time() duration = end_time - start_time print(f"请求 {request_id} 完成时间: {time.strftime('%H:%M:%S', time.localtime(end_time))} (耗时: {duration:.2f}秒)") @@ -42,13 +42,13 @@ async def make_request(session: aiohttp.ClientSession, request_id: int) -> dict: "error": str(e) } -async def run_concurrent_test(num_requests: int = 10) -> List[dict]: +async def run_concurrent_test(num_requests: int = 10) -> list[dict]: """运行并发测试""" async with aiohttp.ClientSession() as session: tasks = [make_request(session, i) for i in range(num_requests)] return await asyncio.gather(*tasks) -def analyze_results(results: List[dict]) -> None: +def analyze_results(results: list[dict]) -> None: """分析并打印测试结果""" total_requests = len(results) successful_requests = sum(1 for r in results if r["success"]) @@ -62,7 +62,7 @@ def analyze_results(results: List[dict]) -> None: else: avg_time = max_time = min_time = 0 - print(f"\n=== 并发测试结果 ===") + print("\n=== 并发测试结果 ===") print(f"总请求数: {total_requests}") print(f"成功请求: {successful_requests}") print(f"失败请求: {failed_requests}") @@ -80,6 +80,9 @@ def analyze_results(results: List[dict]) -> None: print("\n=== 请求时间线分析 ===") sorted_results = sorted(results, key=lambda x: x["start_time"]) test_start_time = sorted_results[0]["start_time"] + test_end_time = sorted_results[-1]["end_time"] + print(f"测试开始时间: {time.strftime('%H:%M:%S', time.localtime(test_start_time))}") + print(f"测试结束时间: {time.strftime('%H:%M:%S', time.localtime(test_end_time))}") print("\n时间线详情:") print("请求ID 开始时间 结束时间 耗时(秒) 重叠请求数")