From e5e4f54734b45571e706715a1f692eda180cfd95 Mon Sep 17 00:00:00 2001 From: FangYin Cheng Date: Tue, 24 Oct 2023 16:07:58 +0800 Subject: [PATCH] fix(core): Fix chat chat tracer no spans bug --- .../agent/commands/disply_type/show_chart_gen.py | 1 - pilot/base_modules/agent/hub/agent_hub.py | 2 +- pilot/base_modules/agent/plugins_util.py | 4 +++- pilot/base_modules/meta_data/meta_data.py | 2 +- pilot/logs.py | 3 ++- .../chat_history/store_type/meta_db_history.py | 2 +- pilot/openapi/api_v1/api_v1.py | 13 ++++++++++--- pilot/scene/chat_agent/out_parser.py | 5 ----- pilot/scene/chat_factory.py | 7 ++++++- 9 files changed, 24 insertions(+), 15 deletions(-) diff --git a/pilot/base_modules/agent/commands/disply_type/show_chart_gen.py b/pilot/base_modules/agent/commands/disply_type/show_chart_gen.py index 166992822..956a064c2 100644 --- a/pilot/base_modules/agent/commands/disply_type/show_chart_gen.py +++ b/pilot/base_modules/agent/commands/disply_type/show_chart_gen.py @@ -1,7 +1,6 @@ from pandas import DataFrame from pilot.base_modules.agent.commands.command_mange import command -from pilot.configs.config import Config import pandas as pd import uuid import os diff --git a/pilot/base_modules/agent/hub/agent_hub.py b/pilot/base_modules/agent/hub/agent_hub.py index 063a2e03a..470b4ab71 100644 --- a/pilot/base_modules/agent/hub/agent_hub.py +++ b/pilot/base_modules/agent/hub/agent_hub.py @@ -12,7 +12,7 @@ from ..common.schema import PluginStorageType from ..plugins_util import scan_plugins, update_from_git -logger = logging.getLogger("agent_hub") +logger = logging.getLogger(__name__) Default_User = "default" DEFAULT_PLUGIN_REPO = "https://github.com/eosphoros-ai/DB-GPT-Plugins.git" TEMP_PLUGIN_PATH = "" diff --git a/pilot/base_modules/agent/plugins_util.py b/pilot/base_modules/agent/plugins_util.py index facc1472d..cf8fb8df3 100644 --- a/pilot/base_modules/agent/plugins_util.py +++ b/pilot/base_modules/agent/plugins_util.py @@ -9,6 +9,7 @@ import git import threading import datetime +import logging from pathlib import Path from typing import List from urllib.parse import urlparse @@ -19,7 +20,8 @@ from pilot.configs.config import Config from pilot.configs.model_config import PLUGINS_DIR -from pilot.logs import logger + +logger = logging.getLogger(__name__) def inspect_zip_for_modules(zip_path: str, debug: bool = False) -> list[str]: diff --git a/pilot/base_modules/meta_data/meta_data.py b/pilot/base_modules/meta_data/meta_data.py index 11cd29e38..687138e77 100644 --- a/pilot/base_modules/meta_data/meta_data.py +++ b/pilot/base_modules/meta_data/meta_data.py @@ -20,7 +20,7 @@ from pilot.configs.config import Config -logger = logging.getLogger("meta_data") +logger = logging.getLogger(__name__) CFG = Config() default_db_path = os.path.join(os.getcwd(), "meta_data") diff --git a/pilot/logs.py b/pilot/logs.py index 673799307..32528bdaa 100644 --- a/pilot/logs.py +++ b/pilot/logs.py @@ -249,7 +249,8 @@ def remove_color_codes(s: str) -> str: return ansi_escape.sub("", s) -logger: Logger = Logger() +# Remove current logger +# logger: Logger = Logger() def print_assistant_thoughts( diff --git a/pilot/memory/chat_history/store_type/meta_db_history.py b/pilot/memory/chat_history/store_type/meta_db_history.py index ed8187d99..cc51804a1 100644 --- a/pilot/memory/chat_history/store_type/meta_db_history.py +++ b/pilot/memory/chat_history/store_type/meta_db_history.py @@ -14,7 +14,7 @@ from pilot.memory.chat_history.base import MemoryStoreType CFG = Config() -logger = logging.getLogger("db_chat_history") +logger = logging.getLogger(__name__) class DbHistoryMemory(BaseChatHistoryMemory): diff --git a/pilot/openapi/api_v1/api_v1.py b/pilot/openapi/api_v1/api_v1.py index ae8e74a2b..23cba7948 100644 --- a/pilot/openapi/api_v1/api_v1.py +++ b/pilot/openapi/api_v1/api_v1.py @@ -47,6 +47,7 @@ from pilot.memory.chat_history.chat_hisotry_factory import ChatHistory from pilot.model.cluster import BaseModelController, WorkerManager, WorkerManagerFactory from pilot.model.base import FlatSupportedModel +from pilot.utils.tracer import root_tracer, SpanType from pilot.utils.executor_utils import ExecutorFactory, blocking_func_to_async router = APIRouter() @@ -389,7 +390,10 @@ async def chat_completions(dialogue: ConversationVo = Body()): print( f"chat_completions:{dialogue.chat_mode},{dialogue.select_param},{dialogue.model_name}" ) - chat: BaseChat = await get_chat_instance(dialogue) + with root_tracer.start_span( + "get_chat_instance", span_type=SpanType.CHAT, metadata=dialogue.dict() + ): + chat: BaseChat = await get_chat_instance(dialogue) # background_tasks = BackgroundTasks() # background_tasks.add_task(release_model_semaphore) headers = { @@ -440,8 +444,9 @@ async def model_supports(worker_manager: WorkerManager = Depends(get_worker_mana async def no_stream_generator(chat): - msg = await chat.nostream_call() - yield f"data: {msg}\n\n" + with root_tracer.start_span("no_stream_generator"): + msg = await chat.nostream_call() + yield f"data: {msg}\n\n" async def stream_generator(chat, incremental: bool, model_name: str): @@ -458,6 +463,7 @@ async def stream_generator(chat, incremental: bool, model_name: str): Yields: _type_: streaming responses """ + span = root_tracer.start_span("stream_generator") msg = "[LLM_ERROR]: llm server has no output, maybe your prompt template is wrong." stream_id = f"chatcmpl-{str(uuid.uuid1())}" @@ -483,6 +489,7 @@ async def stream_generator(chat, incremental: bool, model_name: str): await asyncio.sleep(0.02) if incremental: yield "data: [DONE]\n\n" + span.end() def message2Vo(message: dict, order, model_name) -> MessageVo: diff --git a/pilot/scene/chat_agent/out_parser.py b/pilot/scene/chat_agent/out_parser.py index ecea328bb..2684ae0c9 100644 --- a/pilot/scene/chat_agent/out_parser.py +++ b/pilot/scene/chat_agent/out_parser.py @@ -1,11 +1,6 @@ import json from typing import Dict, NamedTuple -from pilot.utils import build_logger from pilot.out_parser.base import BaseOutputParser, T -from pilot.configs.model_config import LOGDIR - - -logger = build_logger("webserver", LOGDIR + "DbChatOutputParser.log") class PluginAction(NamedTuple): diff --git a/pilot/scene/chat_factory.py b/pilot/scene/chat_factory.py index d33f89e76..e50e56efd 100644 --- a/pilot/scene/chat_factory.py +++ b/pilot/scene/chat_factory.py @@ -1,5 +1,6 @@ from pilot.scene.base_chat import BaseChat from pilot.singleton import Singleton +from pilot.utils.tracer import root_tracer class ChatFactory(metaclass=Singleton): @@ -20,7 +21,11 @@ def get_implementation(chat_mode, **kwargs): implementation = None for cls in chat_classes: if cls.chat_scene == chat_mode: - implementation = cls(**kwargs) + metadata = {"cls": str(cls), "params": kwargs} + with root_tracer.start_span( + "get_implementation_of_chat", metadata=metadata + ): + implementation = cls(**kwargs) if implementation == None: raise Exception(f"Invalid implementation name:{chat_mode}") return implementation