diff --git a/examples/awel/simple_rag_example.py b/examples/awel/simple_rag_example.py new file mode 100644 index 000000000..78c08ac2f --- /dev/null +++ b/examples/awel/simple_rag_example.py @@ -0,0 +1,70 @@ +"""AWEL: Simple rag example + + Example: + + .. code-block:: shell + + curl -X POST http://127.0.0.1:5000/api/v1/awel/trigger/examples/simple_rag \ + -H "Content-Type: application/json" -d '{ + "conv_uid": "36f0e992-8825-11ee-8638-0242ac150003", + "model_name": "proxyllm", + "chat_mode": "chat_knowledge", + "user_input": "What is DB-GPT?", + "select_param": "default" + }' + +""" + +from pilot.awel import HttpTrigger, DAG, MapOperator +from pilot.scene.operator._experimental import ( + ChatContext, + PromptManagerOperator, + ChatHistoryStorageOperator, + ChatHistoryOperator, + EmbeddingEngingOperator, + BaseChatOperator, +) +from pilot.scene.base import ChatScene +from pilot.openapi.api_view_model import ConversationVo +from pilot.model.base import ModelOutput +from pilot.model.operator.model_operator import ModelOperator + + +class RequestParseOperator(MapOperator[ConversationVo, ChatContext]): + def __init__(self, **kwargs): + super().__init__(**kwargs) + + async def map(self, input_value: ConversationVo) -> ChatContext: + return ChatContext( + current_user_input=input_value.user_input, + model_name=input_value.model_name, + chat_session_id=input_value.conv_uid, + select_param=input_value.select_param, + chat_scene=ChatScene.ChatKnowledge, + ) + + +with DAG("simple_rag_example") as dag: + trigger_task = HttpTrigger( + "/examples/simple_rag", methods="POST", request_body=ConversationVo + ) + req_parse_task = RequestParseOperator() + prompt_task = PromptManagerOperator() + history_storage_task = ChatHistoryStorageOperator() + history_task = ChatHistoryOperator() + embedding_task = EmbeddingEngingOperator() + chat_task = BaseChatOperator() + model_task = ModelOperator() + output_parser_task = MapOperator(lambda out: out.to_dict()["text"]) + + ( + trigger_task + >> req_parse_task + >> prompt_task + >> history_storage_task + >> history_task + >> embedding_task + >> chat_task + >> model_task + >> output_parser_task + ) diff --git a/pilot/awel/dag/base.py b/pilot/awel/dag/base.py index 2673cb280..ceb13c8ad 100644 --- a/pilot/awel/dag/base.py +++ b/pilot/awel/dag/base.py @@ -7,6 +7,7 @@ import logging from collections import deque from functools import cache +from concurrent.futures import Executor from pilot.component import SystemApp from ..resource.base import ResourceGroup @@ -102,6 +103,7 @@ class DAGVar: _thread_local = threading.local() _async_local = contextvars.ContextVar("current_dag_stack", default=deque()) _system_app: SystemApp = None + _executor: Executor = None @classmethod def enter_dag(cls, dag) -> None: @@ -157,6 +159,14 @@ def set_current_system_app(cls, system_app: SystemApp) -> None: else: cls._system_app = system_app + @classmethod + def get_executor(cls) -> Executor: + return cls._executor + + @classmethod + def set_executor(cls, executor: Executor) -> None: + cls._executor = executor + class DAGNode(DependencyMixin, ABC): resource_group: Optional[ResourceGroup] = None @@ -165,9 +175,10 @@ class DAGNode(DependencyMixin, ABC): def __init__( self, dag: Optional["DAG"] = None, - node_id: str = None, - node_name: str = None, - system_app: SystemApp = None, + node_id: Optional[str] = None, + node_name: Optional[str] = None, + system_app: Optional[SystemApp] = None, + executor: Optional[Executor] = None, ) -> None: super().__init__() self._upstream: List["DAGNode"] = [] @@ -176,6 +187,7 @@ def __init__( self._system_app: Optional[SystemApp] = ( system_app or DAGVar.get_current_system_app() ) + self._executor: Optional[Executor] = executor or DAGVar.get_executor() if not node_id and self._dag: node_id = self._dag._new_node_id() self._node_id: str = node_id diff --git a/pilot/awel/operator/base.py b/pilot/awel/operator/base.py index 1089420e1..09aa87141 100644 --- a/pilot/awel/operator/base.py +++ b/pilot/awel/operator/base.py @@ -14,7 +14,13 @@ ) import functools from inspect import signature -from pilot.component import SystemApp +from pilot.component import SystemApp, ComponentType +from pilot.utils.executor_utils import ( + ExecutorFactory, + DefaultExecutorFactory, + blocking_func_to_async, + BlockingFunction, +) from ..dag.base import DAGNode, DAGContext, DAGVar, DAG from ..task.base import ( @@ -71,6 +77,16 @@ def apply_defaults(self: "BaseOperator", *args: Any, **kwargs: Any) -> Any: system_app: Optional[SystemApp] = ( kwargs.get("system_app") or DAGVar.get_current_system_app() ) + executor = kwargs.get("executor") or DAGVar.get_executor() + if not executor: + if system_app: + executor = system_app.get_component( + ComponentType.EXECUTOR_DEFAULT, ExecutorFactory + ).create() + else: + executor = DefaultExecutorFactory().create() + DAGVar.set_executor(executor) + if not task_id and dag: task_id = dag._new_node_id() runner: Optional[WorkflowRunner] = kwargs.get("runner") or default_runner @@ -86,6 +102,8 @@ def apply_defaults(self: "BaseOperator", *args: Any, **kwargs: Any) -> Any: kwargs["runner"] = runner if not kwargs.get("system_app"): kwargs["system_app"] = system_app + if not kwargs.get("executor"): + kwargs["executor"] = executor real_obj = func(self, *args, **kwargs) return real_obj @@ -177,6 +195,11 @@ async def call_stream( out_ctx = await self._runner.execute_workflow(self, call_data) return out_ctx.current_task_context.task_output.output_stream + async def blocking_func_to_async( + self, func: BlockingFunction, *args, **kwargs + ) -> Any: + return await blocking_func_to_async(self._executor, func, *args, **kwargs) + def initialize_runner(runner: WorkflowRunner): global default_runner diff --git a/pilot/awel/runner/local_runner.py b/pilot/awel/runner/local_runner.py index 282e6a4e2..6f8a0a484 100644 --- a/pilot/awel/runner/local_runner.py +++ b/pilot/awel/runner/local_runner.py @@ -67,7 +67,7 @@ async def _execute_node( node_outputs[node.node_id] = task_ctx return try: - logger.info( + logger.debug( f"Begin run operator, node id: {node.node_id}, node name: {node.node_name}, cls: {node}" ) await node._run(dag_ctx) @@ -76,7 +76,7 @@ async def _execute_node( if isinstance(node, BranchOperator): skip_nodes = task_ctx.metadata.get("skip_node_names", []) - logger.info( + logger.debug( f"Current is branch operator, skip node names: {skip_nodes}" ) _skip_current_downstream_by_node_name(node, skip_nodes, skip_node_ids) 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 8afbaf06b..f1c25d633 100644 --- a/pilot/memory/chat_history/store_type/meta_db_history.py +++ b/pilot/memory/chat_history/store_type/meta_db_history.py @@ -47,7 +47,7 @@ def create(self, chat_mode, summary: str, user_name: str) -> None: logger.error("init create conversation log error！" + str(e)) def append(self, once_message: OnceConversation) -> None: - logger.info(f"db history append: {once_message}") + logger.debug(f"db history append: {once_message}") chat_history: ChatHistoryEntity = self.chat_history_dao.get_by_uid( self.chat_seesion_id ) diff --git a/pilot/model/proxy/llms/chatgpt.py b/pilot/model/proxy/llms/chatgpt.py index 9e6d1a20a..1da815bfa 100644 --- a/pilot/model/proxy/llms/chatgpt.py +++ b/pilot/model/proxy/llms/chatgpt.py @@ -143,9 +143,7 @@ def _build_request(model: ProxyModel, params): proxyllm_backend = proxyllm_backend or "gpt-3.5-turbo" payloads["model"] = proxyllm_backend - logger.info( - f"Send request to real model {proxyllm_backend}, openai_params: {openai_params}" - ) + logger.info(f"Send request to real model {proxyllm_backend}") return history, payloads diff --git a/pilot/scene/base_chat.py b/pilot/scene/base_chat.py index e5ec0e8b8..0e263f7e5 100644 --- a/pilot/scene/base_chat.py +++ b/pilot/scene/base_chat.py @@ -68,7 +68,7 @@ def __init__(self, chat_param: Dict): CFG.prompt_template_registry.get_prompt_template( self.chat_mode.value(), language=CFG.LANGUAGE, - model_name=CFG.LLM_MODEL, + model_name=self.llm_model, proxyllm_backend=CFG.PROXYLLM_BACKEND, ) ) @@ -141,13 +141,7 @@ def get_llm_speak(self, prompt_define_response): return speak_to_user async def __call_base(self): - import inspect - - input_values = ( - await self.generate_input_values() - if inspect.isawaitable(self.generate_input_values()) - else self.generate_input_values() - ) + input_values = await self.generate_input_values() ### Chat sequence advance self.current_message.chat_order = len(self.history_message) + 1 self.current_message.add_user_message(self.current_user_input) @@ -379,16 +373,18 @@ def generate_llm_text(self) -> str: if self.prompt_template.template_define: text += self.prompt_template.template_define + self.prompt_template.sep ### Load prompt - text += self.__load_system_message() + text += _load_system_message(self.current_message, self.prompt_template) ### Load examples - text += self.__load_example_messages() + text += _load_example_messages(self.prompt_template) ### Load History - text += self.__load_history_messages() + text += _load_history_messages( + self.prompt_template, self.history_message, self.chat_retention_rounds + ) ### Load User Input - text += self.__load_user_message() + text += _load_user_message(self.current_message, self.prompt_template) return text def generate_llm_messages(self) -> List[ModelMessage]: @@ -406,137 +402,26 @@ def generate_llm_messages(self) -> List[ModelMessage]: ) ) ### Load prompt - messages += self.__load_system_message(str_message=False) + messages += _load_system_message( + self.current_message, self.prompt_template, str_message=False + ) ### Load examples - messages += self.__load_example_messages(str_message=False) + messages += _load_example_messages(self.prompt_template, str_message=False) ### Load History - messages += self.__load_history_messages(str_message=False) + messages += _load_history_messages( + self.prompt_template, + self.history_message, + self.chat_retention_rounds, + str_message=False, + ) ### Load User Input - messages += self.__load_user_message(str_message=False) + messages += _load_user_message( + self.current_message, self.prompt_template, str_message=False + ) return messages - def __load_system_message(self, str_message: bool = True): - system_convs = self.current_message.get_system_conv() - system_text = "" - system_messages = [] - for system_conv in system_convs: - system_text += ( - system_conv.type + ":" + system_conv.content + self.prompt_template.sep - ) - system_messages.append( - ModelMessage(role=system_conv.type, content=system_conv.content) - ) - return system_text if str_message else system_messages - - def __load_user_message(self, str_message: bool = True): - user_conv = self.current_message.get_user_conv() - user_messages = [] - if user_conv: - user_text = ( - user_conv.type + ":" + user_conv.content + self.prompt_template.sep - ) - user_messages.append( - ModelMessage(role=user_conv.type, content=user_conv.content) - ) - return user_text if str_message else user_messages - else: - raise ValueError("Hi! What do you want to talk about？") - - def __load_example_messages(self, str_message: bool = True): - example_text = "" - example_messages = [] - if self.prompt_template.example_selector: - for round_conv in self.prompt_template.example_selector.examples(): - for round_message in round_conv["messages"]: - if not round_message["type"] in [ - ModelMessageRoleType.VIEW, - ModelMessageRoleType.SYSTEM, - ]: - message_type = round_message["type"] - message_content = round_message["data"]["content"] - example_text += ( - message_type - + ":" - + message_content - + self.prompt_template.sep - ) - example_messages.append( - ModelMessage(role=message_type, content=message_content) - ) - return example_text if str_message else example_messages - - def __load_history_messages(self, str_message: bool = True): - history_text = "" - history_messages = [] - if self.prompt_template.need_historical_messages: - if self.history_message: - logger.info( - f"There are already {len(self.history_message)} rounds of conversations! Will use {self.chat_retention_rounds} rounds of content as history!" - ) - if len(self.history_message) > self.chat_retention_rounds: - for first_message in self.history_message[0]["messages"]: - if not first_message["type"] in [ - ModelMessageRoleType.VIEW, - ModelMessageRoleType.SYSTEM, - ]: - message_type = first_message["type"] - message_content = first_message["data"]["content"] - history_text += ( - message_type - + ":" - + message_content - + self.prompt_template.sep - ) - history_messages.append( - ModelMessage(role=message_type, content=message_content) - ) - if self.chat_retention_rounds > 1: - index = self.chat_retention_rounds - 1 - for round_conv in self.history_message[-index:]: - for round_message in round_conv["messages"]: - if not round_message["type"] in [ - ModelMessageRoleType.VIEW, - ModelMessageRoleType.SYSTEM, - ]: - message_type = round_message["type"] - message_content = round_message["data"]["content"] - history_text += ( - message_type - + ":" - + message_content - + self.prompt_template.sep - ) - history_messages.append( - ModelMessage( - role=message_type, content=message_content - ) - ) - - else: - ### user all history - for conversation in self.history_message: - for message in conversation["messages"]: - ### histroy message not have promot and view info - if not message["type"] in [ - ModelMessageRoleType.VIEW, - ModelMessageRoleType.SYSTEM, - ]: - message_type = message["type"] - message_content = message["data"]["content"] - history_text += ( - message_type - + ":" - + message_content - + self.prompt_template.sep - ) - history_messages.append( - ModelMessage(role=message_type, content=message_content) - ) - - return history_text if str_message else history_messages - def current_ai_response(self) -> str: for message in self.current_message.messages: if message.type == "view": @@ -656,3 +541,127 @@ def _build_model_operator( cache_check_branch_node >> cached_node >> join_node return join_node + + +def _load_system_message( + current_message: OnceConversation, + prompt_template: PromptTemplate, + str_message: bool = True, +): + system_convs = current_message.get_system_conv() + system_text = "" + system_messages = [] + for system_conv in system_convs: + system_text += ( + system_conv.type + ":" + system_conv.content + prompt_template.sep + ) + system_messages.append( + ModelMessage(role=system_conv.type, content=system_conv.content) + ) + return system_text if str_message else system_messages + + +def _load_user_message( + current_message: OnceConversation, + prompt_template: PromptTemplate, + str_message: bool = True, +): + user_conv = current_message.get_user_conv() + user_messages = [] + if user_conv: + user_text = user_conv.type + ":" + user_conv.content + prompt_template.sep + user_messages.append( + ModelMessage(role=user_conv.type, content=user_conv.content) + ) + return user_text if str_message else user_messages + else: + raise ValueError("Hi! What do you want to talk about？") + + +def _load_example_messages(prompt_template: PromptTemplate, str_message: bool = True): + example_text = "" + example_messages = [] + if prompt_template.example_selector: + for round_conv in prompt_template.example_selector.examples(): + for round_message in round_conv["messages"]: + if not round_message["type"] in [ + ModelMessageRoleType.VIEW, + ModelMessageRoleType.SYSTEM, + ]: + message_type = round_message["type"] + message_content = round_message["data"]["content"] + example_text += ( + message_type + ":" + message_content + prompt_template.sep + ) + example_messages.append( + ModelMessage(role=message_type, content=message_content) + ) + return example_text if str_message else example_messages + + +def _load_history_messages( + prompt_template: PromptTemplate, + history_message: List[OnceConversation], + chat_retention_rounds: int, + str_message: bool = True, +): + history_text = "" + history_messages = [] + if prompt_template.need_historical_messages: + if history_message: + logger.info( + f"There are already {len(history_message)} rounds of conversations! Will use {chat_retention_rounds} rounds of content as history!" + ) + if len(history_message) > chat_retention_rounds: + for first_message in history_message[0]["messages"]: + if not first_message["type"] in [ + ModelMessageRoleType.VIEW, + ModelMessageRoleType.SYSTEM, + ]: + message_type = first_message["type"] + message_content = first_message["data"]["content"] + history_text += ( + message_type + ":" + message_content + prompt_template.sep + ) + history_messages.append( + ModelMessage(role=message_type, content=message_content) + ) + if chat_retention_rounds > 1: + index = chat_retention_rounds - 1 + for round_conv in history_message[-index:]: + for round_message in round_conv["messages"]: + if not round_message["type"] in [ + ModelMessageRoleType.VIEW, + ModelMessageRoleType.SYSTEM, + ]: + message_type = round_message["type"] + message_content = round_message["data"]["content"] + history_text += ( + message_type + + ":" + + message_content + + prompt_template.sep + ) + history_messages.append( + ModelMessage(role=message_type, content=message_content) + ) + + else: + ### user all history + for conversation in history_message: + for message in conversation["messages"]: + ### histroy message not have promot and view info + if not message["type"] in [ + ModelMessageRoleType.VIEW, + ModelMessageRoleType.SYSTEM, + ]: + message_type = message["type"] + message_content = message["data"]["content"] + history_text += ( + message_type + ":" + message_content + prompt_template.sep + ) + history_messages.append( + ModelMessage(role=message_type, content=message_content) + ) + + return history_text if str_message else history_messages diff --git a/pilot/scene/chat_data/chat_excel/excel_reader.py b/pilot/scene/chat_data/chat_excel/excel_reader.py index 6aa1d3d91..00cb27a2b 100644 --- a/pilot/scene/chat_data/chat_excel/excel_reader.py +++ b/pilot/scene/chat_data/chat_excel/excel_reader.py @@ -6,7 +6,6 @@ import sqlparse import pandas as pd import chardet -import pandas as pd import numpy as np from pyparsing import ( CaselessKeyword, @@ -27,6 +26,8 @@ from pilot.common.pd_utils import csv_colunm_foramt from pilot.common.string_utils import is_chinese_include_number +logger = logging.getLogger(__name__) + def excel_colunm_format(old_name: str) -> str: new_column = old_name.strip() @@ -263,7 +264,7 @@ def __init__(self, file_path): file_name = os.path.basename(file_path) self.file_name_without_extension = os.path.splitext(file_name)[0] encoding, confidence = detect_encoding(file_path) - logging.error(f"Detected Encoding: {encoding} (Confidence: {confidence})") + logger.error(f"Detected Encoding: {encoding} (Confidence: {confidence})") self.excel_file_name = file_name self.extension = os.path.splitext(file_name)[1] # read excel file @@ -323,7 +324,7 @@ def run(self, sql): colunms.append(descrip[0]) return colunms, results.fetchall() except Exception as e: - logging.error("excel sql run error!", e) + logger.error(f"excel sql run error!, {str(e)}") raise ValueError(f"Data Query Exception!\\nSQL[{sql}].\\nError:{str(e)}") def get_df_by_sql_ex(self, sql): diff --git a/pilot/scene/chat_db/auto_execute/out_parser.py b/pilot/scene/chat_db/auto_execute/out_parser.py index 1cd5765da..bd1dd9de8 100644 --- a/pilot/scene/chat_db/auto_execute/out_parser.py +++ b/pilot/scene/chat_db/auto_execute/out_parser.py @@ -37,7 +37,7 @@ def is_sql_statement(self, statement): def parse_prompt_response(self, model_out_text): clean_str = super().parse_prompt_response(model_out_text) - logging.info("clean prompt response:", clean_str) + logger.info(f"clean prompt response: {clean_str}") # Compatible with community pure sql output model if self.is_sql_statement(clean_str): return SqlAction(clean_str, "") @@ -51,7 +51,7 @@ def parse_prompt_response(self, model_out_text): thoughts = response[key] return SqlAction(sql, thoughts) except Exception as e: - logging.error("json load faild") + logger.error("json load faild") return SqlAction("", clean_str) def parse_view_response(self, speak, data, prompt_response) -> str: diff --git a/pilot/scene/chat_knowledge/extract_entity/chat.py b/pilot/scene/chat_knowledge/extract_entity/chat.py index bb52961b5..373bb4e5d 100644 --- a/pilot/scene/chat_knowledge/extract_entity/chat.py +++ b/pilot/scene/chat_knowledge/extract_entity/chat.py @@ -24,7 +24,7 @@ def __init__(self, chat_param: Dict): self.user_input = chat_param["current_user_input"] self.extract_mode = chat_param["select_param"] - def generate_input_values(self): + async def generate_input_values(self): input_values = { "text": self.user_input, } diff --git a/pilot/scene/chat_knowledge/extract_triplet/chat.py b/pilot/scene/chat_knowledge/extract_triplet/chat.py index 11fe871ab..28152b92e 100644 --- a/pilot/scene/chat_knowledge/extract_triplet/chat.py +++ b/pilot/scene/chat_knowledge/extract_triplet/chat.py @@ -24,7 +24,7 @@ def __init__(self, chat_param: Dict): self.user_input = chat_param["current_user_input"] self.extract_mode = chat_param["select_param"] - def generate_input_values(self): + async def generate_input_values(self): input_values = { "text": self.user_input, } diff --git a/pilot/scene/chat_knowledge/refine_summary/chat.py b/pilot/scene/chat_knowledge/refine_summary/chat.py index a257332ae..2f3181d5e 100644 --- a/pilot/scene/chat_knowledge/refine_summary/chat.py +++ b/pilot/scene/chat_knowledge/refine_summary/chat.py @@ -23,7 +23,7 @@ def __init__(self, chat_param: Dict): self.existing_answer = chat_param["select_param"] - def generate_input_values(self): + async def generate_input_values(self): input_values = { # "context": self.user_input, "existing_answer": self.existing_answer, diff --git a/pilot/scene/chat_knowledge/summary/chat.py b/pilot/scene/chat_knowledge/summary/chat.py index 7327b7a5b..be4ee00c3 100644 --- a/pilot/scene/chat_knowledge/summary/chat.py +++ b/pilot/scene/chat_knowledge/summary/chat.py @@ -23,7 +23,7 @@ def __init__(self, chat_param: Dict): self.user_input = chat_param["select_param"] - def generate_input_values(self): + async def generate_input_values(self): input_values = { "context": self.user_input, } diff --git a/pilot/scene/chat_knowledge/v1/chat.py b/pilot/scene/chat_knowledge/v1/chat.py index 551fe2d36..a0c15e658 100644 --- a/pilot/scene/chat_knowledge/v1/chat.py +++ b/pilot/scene/chat_knowledge/v1/chat.py @@ -104,7 +104,7 @@ async def generate_input_values(self) -> Dict: self.current_user_input, self.top_k, ) - self.sources = self.merge_by_key( + self.sources = _merge_by_key( list(map(lambda doc: doc.metadata, docs)), "source" ) @@ -149,29 +149,6 @@ def parse_source_view(self, sources: List): ) return html - def merge_by_key(self, data, key): - result = {} - for item in data: - if item.get(key): - item_key = os.path.basename(item.get(key)) - if item_key in result: - if "pages" in result[item_key] and "page" in item: - result[item_key]["pages"].append(str(item["page"])) - elif "page" in item: - result[item_key]["pages"] = [ - result[item_key]["pages"], - str(item["page"]), - ] - else: - if "page" in item: - result[item_key] = { - "source": item_key, - "pages": [str(item["page"])], - } - else: - result[item_key] = {"source": item_key} - return list(result.values()) - @property def chat_type(self) -> str: return ChatScene.ChatKnowledge.value() @@ -179,3 +156,27 @@ def chat_type(self) -> str: def get_space_context(self, space_name): service = KnowledgeService() return service.get_space_context(space_name) + + +def _merge_by_key(data, key): + result = {} + for item in data: + if item.get(key): + item_key = os.path.basename(item.get(key)) + if item_key in result: + if "pages" in result[item_key] and "page" in item: + result[item_key]["pages"].append(str(item["page"])) + elif "page" in item: + result[item_key]["pages"] = [ + result[item_key]["pages"], + str(item["page"]), + ] + else: + if "page" in item: + result[item_key] = { + "source": item_key, + "pages": [str(item["page"])], + } + else: + result[item_key] = {"source": item_key} + return list(result.values()) diff --git a/pilot/server/componet_configs.py b/pilot/scene/operator/__init__.py similarity index 100% rename from pilot/server/componet_configs.py rename to pilot/scene/operator/__init__.py diff --git a/pilot/scene/operator/_experimental.py b/pilot/scene/operator/_experimental.py new file mode 100644 index 000000000..f0ee06179 --- /dev/null +++ b/pilot/scene/operator/_experimental.py @@ -0,0 +1,255 @@ +from typing import Dict, Optional, List, Any +from dataclasses import dataclass +import datetime +import os +from pilot.awel import MapOperator +from pilot.prompts.prompt_new import PromptTemplate +from pilot.configs.config import Config +from pilot.scene.base import ChatScene +from pilot.scene.message import OnceConversation +from pilot.scene.base_message import ModelMessage, ModelMessageRoleType + + +from pilot.memory.chat_history.base import BaseChatHistoryMemory +from pilot.memory.chat_history.chat_hisotry_factory import ChatHistory + +# TODO move global config +CFG = Config() + + +@dataclass +class ChatContext: + current_user_input: str + model_name: Optional[str] + chat_session_id: Optional[str] = None + select_param: Optional[str] = None + chat_scene: Optional[ChatScene] = ChatScene.ChatNormal + prompt_template: Optional[PromptTemplate] = None + chat_retention_rounds: Optional[int] = 0 + history_storage: Optional[BaseChatHistoryMemory] = None + history_manager: Optional["ChatHistoryManager"] = None + # The input values for prompt template + input_values: Optional[Dict] = None + echo: Optional[bool] = False + + def build_model_payload(self) -> Dict: + if not self.input_values: + raise ValueError("The input value can't be empty") + llm_messages = self.history_manager._new_chat(self.input_values) + return { + "model": self.model_name, + "prompt": "", + "messages": llm_messages, + "temperature": float(self.prompt_template.temperature), + "max_new_tokens": int(self.prompt_template.max_new_tokens), + "echo": self.echo, + } + + +class ChatHistoryManager: + def __init__( + self, + chat_ctx: ChatContext, + prompt_template: PromptTemplate, + history_storage: BaseChatHistoryMemory, + chat_retention_rounds: Optional[int] = 0, + ) -> None: + self._chat_ctx = chat_ctx + self.chat_retention_rounds = chat_retention_rounds + self.current_message: OnceConversation = OnceConversation( + chat_ctx.chat_scene.value() + ) + self.prompt_template = prompt_template + self.history_storage: BaseChatHistoryMemory = history_storage + self.history_message: List[OnceConversation] = history_storage.messages() + self.current_message.model_name = chat_ctx.model_name + if chat_ctx.select_param: + if len(chat_ctx.chat_scene.param_types()) > 0: + self.current_message.param_type = chat_ctx.chat_scene.param_types()[0] + self.current_message.param_value = chat_ctx.select_param + + def _new_chat(self, input_values: Dict) -> List[ModelMessage]: + self.current_message.chat_order = len(self.history_message) + 1 + self.current_message.add_user_message(self._chat_ctx.current_user_input) + self.current_message.start_date = datetime.datetime.now().strftime( + "%Y-%m-%d %H:%M:%S" + ) + self.current_message.tokens = 0 + if self.prompt_template.template: + current_prompt = self.prompt_template.format(**input_values) + self.current_message.add_system_message(current_prompt) + return self._generate_llm_messages() + + def _generate_llm_messages(self) -> List[ModelMessage]: + from pilot.scene.base_chat import ( + _load_system_message, + _load_example_messages, + _load_history_messages, + _load_user_message, + ) + + messages = [] + ### Load scene setting or character definition as system message + if self.prompt_template.template_define: + messages.append( + ModelMessage( + role=ModelMessageRoleType.SYSTEM, + content=self.prompt_template.template_define, + ) + ) + ### Load prompt + messages += _load_system_message( + self.current_message, self.prompt_template, str_message=False + ) + ### Load examples + messages += _load_example_messages(self.prompt_template, str_message=False) + + ### Load History + messages += _load_history_messages( + self.prompt_template, + self.history_message, + self.chat_retention_rounds, + str_message=False, + ) + + ### Load User Input + messages += _load_user_message( + self.current_message, self.prompt_template, str_message=False + ) + return messages + + +class PromptManagerOperator(MapOperator[ChatContext, ChatContext]): + def __init__(self, prompt_template: PromptTemplate = None, **kwargs): + super().__init__(**kwargs) + self._prompt_template = prompt_template + + async def map(self, input_value: ChatContext) -> ChatContext: + if not self._prompt_template: + self._prompt_template: PromptTemplate = ( + CFG.prompt_template_registry.get_prompt_template( + input_value.chat_scene.value(), + language=CFG.LANGUAGE, + model_name=input_value.model_name, + proxyllm_backend=CFG.PROXYLLM_BACKEND, + ) + ) + input_value.prompt_template = self._prompt_template + return input_value + + +class ChatHistoryStorageOperator(MapOperator[ChatContext, ChatContext]): + def __init__(self, history: BaseChatHistoryMemory = None, **kwargs): + super().__init__(**kwargs) + self._history = history + + async def map(self, input_value: ChatContext) -> ChatContext: + if self._history: + return self._history + chat_history_fac = ChatHistory() + input_value.history_storage = chat_history_fac.get_store_instance( + input_value.chat_session_id + ) + return input_value + + +class ChatHistoryOperator(MapOperator[ChatContext, ChatContext]): + def __init__(self, history: BaseChatHistoryMemory = None, **kwargs): + super().__init__(**kwargs) + self._history = history + + async def map(self, input_value: ChatContext) -> ChatContext: + history_storage = self._history or input_value.history_storage + if not history_storage: + from pilot.memory.chat_history.store_type.mem_history import ( + MemHistoryMemory, + ) + + history_storage = MemHistoryMemory(input_value.chat_session_id) + input_value.history_storage = history_storage + input_value.history_manager = ChatHistoryManager( + input_value, + input_value.prompt_template, + history_storage, + input_value.chat_retention_rounds, + ) + return input_value + + +class EmbeddingEngingOperator(MapOperator[ChatContext, ChatContext]): + def __init__(self, **kwargs): + super().__init__(**kwargs) + + async def map(self, input_value: ChatContext) -> ChatContext: + from pilot.configs.model_config import EMBEDDING_MODEL_CONFIG + from pilot.embedding_engine.embedding_engine import EmbeddingEngine + from pilot.embedding_engine.embedding_factory import EmbeddingFactory + from pilot.scene.chat_knowledge.v1.chat import _merge_by_key + + # TODO, decompose the current operator into some atomic operators + knowledge_space = input_value.select_param + vector_store_config = { + "vector_store_name": knowledge_space, + "vector_store_type": CFG.VECTOR_STORE_TYPE, + } + embedding_factory = self.system_app.get_component( + "embedding_factory", EmbeddingFactory + ) + knowledge_embedding_client = EmbeddingEngine( + model_name=EMBEDDING_MODEL_CONFIG[CFG.EMBEDDING_MODEL], + vector_store_config=vector_store_config, + embedding_factory=embedding_factory, + ) + space_context = await self._get_space_context(knowledge_space) + top_k = ( + CFG.KNOWLEDGE_SEARCH_TOP_SIZE + if space_context is None + else int(space_context["embedding"]["topk"]) + ) + max_token = ( + CFG.KNOWLEDGE_SEARCH_MAX_TOKEN + if space_context is None or space_context.get("prompt") is None + else int(space_context["prompt"]["max_token"]) + ) + input_value.prompt_template.template_is_strict = False + if space_context and space_context.get("prompt"): + input_value.prompt_template.template_define = space_context["prompt"][ + "scene" + ] + input_value.prompt_template.template = space_context["prompt"]["template"] + + docs = await self.blocking_func_to_async( + knowledge_embedding_client.similar_search, + input_value.current_user_input, + top_k, + ) + sources = _merge_by_key(list(map(lambda doc: doc.metadata, docs)), "source") + if not docs or len(docs) == 0: + print("no relevant docs to retrieve") + context = "no relevant docs to retrieve" + else: + context = [d.page_content for d in docs] + context = context[:max_token] + relations = list( + set([os.path.basename(str(d.metadata.get("source", ""))) for d in docs]) + ) + input_value.input_values = { + "context": context, + "question": input_value.current_user_input, + "relations": relations, + } + return input_value + + async def _get_space_context(self, space_name): + from pilot.server.knowledge.service import KnowledgeService + + service = KnowledgeService() + return await self.blocking_func_to_async(service.get_space_context, space_name) + + +class BaseChatOperator(MapOperator[ChatContext, Dict]): + def __init__(self, **kwargs): + super().__init__(**kwargs) + + async def map(self, input_value: ChatContext) -> Dict: + return input_value.build_model_payload()