From e04add5bc25c846f1fce4055076a054d45f376cd Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Thu, 27 Feb 2025 19:35:25 +0800 Subject: [PATCH] update logger --- src/config/__init__.py | 5 +---- src/core/database.py | 4 +--- src/core/knowledgebase.py | 3 +-- src/core/retriever.py | 4 +--- src/core/startup.py | 4 +--- src/main.py | 5 +---- src/models/__init__.py | 1 + src/models/chat_model.py | 5 +---- src/models/embedding.py | 5 +---- src/models/ollama_embedding.py | 4 +--- src/models/rerank_model.py | 5 +---- src/plugins/oneke.py | 3 +-- src/routers/base_router.py | 8 ++++---- src/routers/chat_router.py | 3 +-- src/routers/data_router.py | 3 +-- src/routers/tool_router.py | 3 +-- src/utils/__init__.py | 2 +- src/utils/web_search.py | 4 +--- 18 files changed, 21 insertions(+), 50 deletions(-) diff --git a/src/config/__init__.py b/src/config/__init__.py index 32a361ea..522788bd 100644 --- a/src/config/__init__.py +++ b/src/config/__init__.py @@ -2,9 +2,7 @@ import os import json import yaml from pathlib import Path -from src.utils.logging_config import setup_logger - -logger = setup_logger("Config") +from src.utils.logging_config import logger with open(Path("src/static/models.yaml"), "r") as f: _models = yaml.safe_load(f) @@ -53,7 +51,6 @@ 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_search_engine", default=False, des="是否开启搜索引擎") self.add_item("enable_web_search", default=False, des="是否开启网页搜索（需配置 TAVILY_API_KEY）") # 模型配置 ## 注意这里是模型名，而不是具体的模型路径，默认使用 HuggingFace 的路径 diff --git a/src/core/database.py b/src/core/database.py index 3498dbb8..99ee2832 100644 --- a/src/core/database.py +++ b/src/core/database.py @@ -2,11 +2,9 @@ import os import json import time from src.plugins import pdf2txt -from src.utils import hashstr, setup_logger, is_text_pdf +from src.utils import hashstr, logger, is_text_pdf from src.models.embedding import get_embedding_model -logger = setup_logger("DataBaseManager") - class DataBaseManager: diff --git a/src/core/knowledgebase.py b/src/core/knowledgebase.py index 99ecb830..6d68087e 100644 --- a/src/core/knowledgebase.py +++ b/src/core/knowledgebase.py @@ -1,8 +1,7 @@ import os from pymilvus import MilvusClient, MilvusException -from src.utils import setup_logger, hashstr -logger = setup_logger("KnowledgeBase") +from src.utils import logger, hashstr class KnowledgeBase: diff --git a/src/core/retriever.py b/src/core/retriever.py index 1cfb862a..fd5771a7 100644 --- a/src/core/retriever.py +++ b/src/core/retriever.py @@ -1,7 +1,5 @@ from src.models.rerank_model import get_reranker -from src.utils.logging_config import setup_logger - -logger = setup_logger("server-common") +from src.utils.logging_config import logger class Retriever: diff --git a/src/core/startup.py b/src/core/startup.py index edab9216..11092a2d 100644 --- a/src/core/startup.py +++ b/src/core/startup.py @@ -3,9 +3,7 @@ from src.core import DataBaseManager from src.core.retriever import Retriever from src.models import select_model from src.config import Config -from src.utils import setup_logger - -logger = setup_logger("Startup") +from src.utils import logger class Startup: diff --git a/src/main.py b/src/main.py index d831e3d0..9266eb32 100644 --- a/src/main.py +++ b/src/main.py @@ -3,7 +3,7 @@ from dotenv import load_dotenv from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from src.routers import router -from src.utils.logging_config import setup_logger +from src.utils.logging_config import logger load_dotenv() @@ -20,9 +20,6 @@ app.add_middleware( ) -logger = setup_logger("server:main") - - if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=5000, threads=10, workers=10) diff --git a/src/models/__init__.py b/src/models/__init__.py index c3d0a280..e4ec33e4 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -18,6 +18,7 @@ def select_model(config, model_provider=None, model_name=None): "siliconflow", "zhipu", "lingyiwanwu", + "together.ai", ]: return OpenAIBase( api_key=os.getenv(model_info["env"][0]), diff --git a/src/models/chat_model.py b/src/models/chat_model.py index cfbcf4bb..49f1df9b 100644 --- a/src/models/chat_model.py +++ b/src/models/chat_model.py @@ -1,9 +1,6 @@ import os from openai import OpenAI -from src.utils.logging_config import setup_logger - - -logger = setup_logger(__name__) +from src.utils import logger class OpenAIBase(): def __init__(self, api_key, base_url, model_name): diff --git a/src/models/embedding.py b/src/models/embedding.py index 2bcd8906..136b130b 100644 --- a/src/models/embedding.py +++ b/src/models/embedding.py @@ -4,12 +4,9 @@ import requests from FlagEmbedding import FlagModel from src.config import EMBED_MODEL_INFO -from src.utils.logging_config import setup_logger -from src.utils import hashstr +from src.utils import hashstr, logger -logger = setup_logger("EmbeddingModel") - class LocalEmbeddingModel(FlagModel): def __init__(self, config, **kwargs): info = EMBED_MODEL_INFO[config.embed_model] diff --git a/src/models/ollama_embedding.py b/src/models/ollama_embedding.py index 0819a46b..1aefc16c 100644 --- a/src/models/ollama_embedding.py +++ b/src/models/ollama_embedding.py @@ -4,9 +4,7 @@ import numpy as np from typing import List, Union, Dict from src.models.embedding import RemoteEmbeddingModel -from src.utils.logging_config import setup_logger - -logger = setup_logger("OllamaEmbedding") +from src.utils.logging_config import logger class OllamaEmbedding(RemoteEmbeddingModel): """ diff --git a/src/models/rerank_model.py b/src/models/rerank_model.py index dbd24b0d..5f02a75e 100644 --- a/src/models/rerank_model.py +++ b/src/models/rerank_model.py @@ -5,10 +5,7 @@ import numpy as np from FlagEmbedding import FlagReranker from src.config import RERANKER_LIST -from src.utils.logging_config import setup_logger - - -logger = setup_logger("RerankModel") +from src.utils.logging_config import logger class LocalReranker(FlagReranker): diff --git a/src/plugins/oneke.py b/src/plugins/oneke.py index 3ca3d60b..03bceacf 100644 --- a/src/plugins/oneke.py +++ b/src/plugins/oneke.py @@ -11,8 +11,7 @@ from transformers import ( BitsAndBytesConfig ) -from src.utils import setup_logger -logger = setup_logger("OneKE") +from src.utils import logger dotenv.load_dotenv() diff --git a/src/routers/base_router.py b/src/routers/base_router.py index 7d9b0f21..b89eff59 100644 --- a/src/routers/base_router.py +++ b/src/routers/base_router.py @@ -5,11 +5,11 @@ base = APIRouter() from fastapi import FastAPI, HTTPException from fastapi.responses import JSONResponse from fastapi import Request, Body -from src.core import HistoryManager -from src.utils.logging_config import setup_logger -from src.core.startup import startup -logger = setup_logger("server-base") +from src.core import HistoryManager +from src.core.startup import startup +from src.utils import logger + @base.get("/") async def route_index(): diff --git a/src/routers/chat_router.py b/src/routers/chat_router.py index f199c791..c15a63a5 100644 --- a/src/routers/chat_router.py +++ b/src/routers/chat_router.py @@ -5,10 +5,9 @@ from fastapi.responses import StreamingResponse, Response from concurrent.futures import ThreadPoolExecutor from src.core import HistoryManager from src.core.startup import startup -from src.utils.logging_config import setup_logger +from src.utils.logging_config import logger chat = APIRouter(prefix="/chat") -logger = setup_logger("server-chat") # 创建线程池 executor = ThreadPoolExecutor() diff --git a/src/routers/data_router.py b/src/routers/data_router.py index 9d39d01e..8bc9054a 100644 --- a/src/routers/data_router.py +++ b/src/routers/data_router.py @@ -2,12 +2,11 @@ import os from typing import List, Optional from fastapi import APIRouter, File, UploadFile, HTTPException, Depends, Body -from src.utils import setup_logger, hashstr +from src.utils import logger, hashstr from src.core.startup import startup data = APIRouter(prefix="/data") -logger = setup_logger("server-database") @data.get("/") def get_databases(): diff --git a/src/routers/tool_router.py b/src/routers/tool_router.py index 6732fed9..5690c086 100644 --- a/src/routers/tool_router.py +++ b/src/routers/tool_router.py @@ -3,11 +3,10 @@ from fastapi import APIRouter, Body from pydantic import BaseModel from typing import List, Dict, Any, Optional -from src.utils import setup_logger +from src.utils import logger tool = APIRouter(prefix="/tool") -logger = setup_logger("server-tools") class Tool(BaseModel): name: str diff --git a/src/utils/__init__.py b/src/utils/__init__.py index 5b613e7a..37ff7c4a 100644 --- a/src/utils/__init__.py +++ b/src/utils/__init__.py @@ -1,6 +1,6 @@ import time import random -from src.utils.logging_config import setup_logger, logger +from src.utils.logging_config import logger def is_text_pdf(pdf_path): import fitz diff --git a/src/utils/web_search.py b/src/utils/web_search.py index e66ecf86..cb343ddd 100644 --- a/src/utils/web_search.py +++ b/src/utils/web_search.py @@ -1,9 +1,7 @@ import os from typing import List, Dict from tavily import TavilyClient -from src.utils.logging_config import setup_logger - -logger = setup_logger("web-search") +from src.utils.logging_config import logger class WebSearcher: def __init__(self):