from collections.abc import Generator, Sequence from typing import Any, Union from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity from core.app.llm import deduct_llm_quota from core.model_manager import ModelInstance, ModelManager from core.prompt.advanced_prompt_transform import AdvancedPromptTransform from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate from core.rag.retrieval.output_parser.react_output import ReactAction from core.rag.retrieval.output_parser.structured_chat import StructuredChatOutputParser from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage from graphon.model_runtime.entities.message_entities import PromptMessage, PromptMessageRole, PromptMessageTool from graphon.model_runtime.entities.model_entities import ModelType PREFIX = """Respond to the human as helpfully and accurately as possible. You have access to the following tools:""" SUFFIX = """Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation:. Thought:""" # noqa: E501 FORMAT_INSTRUCTIONS = """Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input). The nouns in the format of "Thought", "Action", "Action Input", "Final Answer" must be expressed in English. Valid "action" values: "Final Answer" or {tool_names} Provide only ONE action per $JSON_BLOB, as shown: ``` {{ "action": $TOOL_NAME, "action_input": $INPUT }} ``` Follow this format: Question: input question to answer Thought: consider previous and subsequent steps Action: ``` $JSON_BLOB ``` Observation: action result ... (repeat Thought/Action/Observation N times) Thought: I know what to respond Action: ``` {{ "action": "Final Answer", "action_input": "Final response to human" }} ```""" # noqa: E501 class ReactMultiDatasetRouter: def invoke( self, query: str, dataset_tools: list[PromptMessageTool], model_config: ModelConfigWithCredentialsEntity, model_instance: ModelInstance, user_id: str, tenant_id: str, ) -> tuple[Union[str, None], LLMUsage]: """Given input, decided what to do. Returns: Action specifying what tool to use. """ if len(dataset_tools) == 0: return None, LLMUsage.empty_usage() elif len(dataset_tools) == 1: return dataset_tools[0].name, LLMUsage.empty_usage() try: return self._react_invoke( query=query, model_config=model_config, model_instance=model_instance, tools=dataset_tools, user_id=user_id, tenant_id=tenant_id, ) except Exception: return None, LLMUsage.empty_usage() def _react_invoke( self, query: str, model_config: ModelConfigWithCredentialsEntity, model_instance: ModelInstance, tools: Sequence[PromptMessageTool], user_id: str, tenant_id: str, prefix: str = PREFIX, suffix: str = SUFFIX, format_instructions: str = FORMAT_INSTRUCTIONS, ) -> tuple[Union[str, None], LLMUsage]: prompt: Union[list[ChatModelMessage], CompletionModelPromptTemplate] if model_config.mode == "chat": prompt = self.create_chat_prompt( query=query, tools=tools, prefix=prefix, suffix=suffix, format_instructions=format_instructions, ) else: prompt = self.create_completion_prompt( tools=tools, prefix=prefix, format_instructions=format_instructions, ) stop = ["Observation:"] # handle invoke result prompt_transform = AdvancedPromptTransform() prompt_messages = prompt_transform.get_prompt( prompt_template=prompt, inputs={}, query="", files=[], context="", memory_config=None, memory=None, model_config=model_config, model_instance=model_instance, ) result_text, usage = self._invoke_llm( completion_param=model_config.parameters, model_instance=model_instance, prompt_messages=prompt_messages, stop=stop, user_id=user_id, tenant_id=tenant_id, ) output_parser = StructuredChatOutputParser() react_decision = output_parser.parse(result_text) if isinstance(react_decision, ReactAction): return react_decision.tool, usage return None, usage def _invoke_llm( self, completion_param: dict[str, Any], model_instance: ModelInstance, prompt_messages: list[PromptMessage], stop: list[str], user_id: str, tenant_id: str, ) -> tuple[str, LLMUsage]: """ Invoke large language model :param model_instance: model instance :param prompt_messages: prompt messages :param stop: stop :return: """ bound_model_instance = ModelManager.for_tenant(tenant_id=tenant_id, user_id=user_id).get_model_instance( tenant_id=tenant_id, provider=model_instance.provider, model_type=ModelType.LLM, model=model_instance.model_name, ) invoke_result: Generator[LLMResult, None, None] = bound_model_instance.invoke_llm( prompt_messages=prompt_messages, model_parameters=completion_param, stop=stop, stream=True, ) # handle invoke result text, usage = self._handle_invoke_result(invoke_result=invoke_result) # deduct quota deduct_llm_quota(tenant_id=tenant_id, model_instance=bound_model_instance, usage=usage) return text, usage def _handle_invoke_result(self, invoke_result: Generator) -> tuple[str, LLMUsage]: """ Handle invoke result :param invoke_result: invoke result :return: """ model = None prompt_messages: list[PromptMessage] = [] full_text = "" usage = None for result in invoke_result: text = result.delta.message.content full_text += text if not model: model = result.model if not prompt_messages: prompt_messages = result.prompt_messages if not usage and result.delta.usage: usage = result.delta.usage if not usage: usage = LLMUsage.empty_usage() return full_text, usage def create_chat_prompt( self, query: str, tools: Sequence[PromptMessageTool], prefix: str = PREFIX, suffix: str = SUFFIX, format_instructions: str = FORMAT_INSTRUCTIONS, ) -> list[ChatModelMessage]: tool_strings = [] for tool in tools: tool_strings.append( f"{tool.name}: {tool.description}, args: {{'query': {{'title': 'Query'," f" 'description': 'Query for the dataset to be used to retrieve the dataset.', 'type': 'string'}}}}" ) formatted_tools = "\n".join(tool_strings) unique_tool_names = {tool.name for tool in tools} tool_names = ", ".join('"' + name + '"' for name in unique_tool_names) format_instructions = format_instructions.format(tool_names=tool_names) template = "\n\n".join([prefix, formatted_tools, format_instructions, suffix]) prompt_messages = [] system_prompt_messages = ChatModelMessage(role=PromptMessageRole.SYSTEM, text=template) prompt_messages.append(system_prompt_messages) user_prompt_message = ChatModelMessage(role=PromptMessageRole.USER, text=query) prompt_messages.append(user_prompt_message) return prompt_messages def create_completion_prompt( self, tools: Sequence[PromptMessageTool], prefix: str = PREFIX, format_instructions: str = FORMAT_INSTRUCTIONS, ) -> CompletionModelPromptTemplate: """Create prompt in the style of the zero shot agent. Args: tools: List of tools the agent will have access to, used to format the prompt. prefix: String to put before the list of tools. format_instructions: The format instruction prompt. Returns: A PromptTemplate with the template assembled from the pieces here. """ suffix = """Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation:. Question: {input} Thought: {agent_scratchpad} """ # noqa: E501 tool_strings = "\n".join([f"{tool.name}: {tool.description}" for tool in tools]) tool_names = ", ".join([tool.name for tool in tools]) format_instructions = format_instructions.format(tool_names=tool_names) template = "\n\n".join([prefix, tool_strings, format_instructions, suffix]) return CompletionModelPromptTemplate(text=template)