import time from collections.abc import Generator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Optional, Union from core.app.app_config.entities import ExternalDataVariableEntity, PromptTemplateEntity from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom from core.app.entities.app_invoke_entities import ( AppGenerateEntity, EasyUIBasedAppGenerateEntity, InvokeFrom, ModelConfigWithCredentialsEntity, ) from core.app.entities.queue_entities import QueueAgentMessageEvent, QueueLLMChunkEvent, QueueMessageEndEvent from core.app.features.annotation_reply.annotation_reply import AnnotationReplyFeature from core.app.features.hosting_moderation.hosting_moderation import HostingModerationFeature from core.external_data_tool.external_data_fetch import ExternalDataFetch from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage from core.model_runtime.entities.message_entities import AssistantPromptMessage, PromptMessage from core.model_runtime.entities.model_entities import ModelPropertyKey from core.model_runtime.errors.invoke import InvokeBadRequestError from core.moderation.input_moderation import InputModeration from core.prompt.advanced_prompt_transform import AdvancedPromptTransform from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig from core.prompt.simple_prompt_transform import ModelMode, SimplePromptTransform from models.model import App, AppMode, Message, MessageAnnotation if TYPE_CHECKING: from core.file.models import File class AppRunner: def get_pre_calculate_rest_tokens( self, app_record: App, model_config: ModelConfigWithCredentialsEntity, prompt_template_entity: PromptTemplateEntity, inputs: Mapping[str, str], files: Sequence["File"], query: Optional[str] = None, ) -> int: """ Get pre calculate rest tokens :param app_record: app record :param model_config: model config entity :param prompt_template_entity: prompt template entity :param inputs: inputs :param files: files :param query: query :return: """ # Invoke model model_instance = ModelInstance( provider_model_bundle=model_config.provider_model_bundle, model=model_config.model ) model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE) max_tokens = 0 for parameter_rule in model_config.model_schema.parameter_rules: if parameter_rule.name == "max_tokens" or ( parameter_rule.use_template and parameter_rule.use_template == "max_tokens" ): max_tokens = ( model_config.parameters.get(parameter_rule.name) or model_config.parameters.get(parameter_rule.use_template or "") ) or 0 if model_context_tokens is None: return -1 if max_tokens is None: max_tokens = 0 # get prompt messages without memory and context prompt_messages, stop = self.organize_prompt_messages( app_record=app_record, model_config=model_config, prompt_template_entity=prompt_template_entity, inputs=inputs, files=files, query=query, ) prompt_tokens = model_instance.get_llm_num_tokens(prompt_messages) rest_tokens: int = model_context_tokens - max_tokens - prompt_tokens if rest_tokens < 0: raise InvokeBadRequestError( "Query or prefix prompt is too long, you can reduce the prefix prompt, " "or shrink the max token, or switch to a llm with a larger token limit size." ) return rest_tokens def recalc_llm_max_tokens( self, model_config: ModelConfigWithCredentialsEntity, prompt_messages: list[PromptMessage] ): # recalc max_tokens if sum(prompt_token + max_tokens) over model token limit model_instance = ModelInstance( provider_model_bundle=model_config.provider_model_bundle, model=model_config.model ) model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE) max_tokens = 0 for parameter_rule in model_config.model_schema.parameter_rules: if parameter_rule.name == "max_tokens" or ( parameter_rule.use_template and parameter_rule.use_template == "max_tokens" ): max_tokens = ( model_config.parameters.get(parameter_rule.name) or model_config.parameters.get(parameter_rule.use_template or "") ) or 0 if model_context_tokens is None: return -1 if max_tokens is None: max_tokens = 0 prompt_tokens = model_instance.get_llm_num_tokens(prompt_messages) if prompt_tokens + max_tokens > model_context_tokens: max_tokens = max(model_context_tokens - prompt_tokens, 16) for parameter_rule in model_config.model_schema.parameter_rules: if parameter_rule.name == "max_tokens" or ( parameter_rule.use_template and parameter_rule.use_template == "max_tokens" ): model_config.parameters[parameter_rule.name] = max_tokens def organize_prompt_messages( self, app_record: App, model_config: ModelConfigWithCredentialsEntity, prompt_template_entity: PromptTemplateEntity, inputs: Mapping[str, str], files: Sequence["File"], query: Optional[str] = None, context: Optional[str] = None, memory: Optional[TokenBufferMemory] = None, ) -> tuple[list[PromptMessage], Optional[list[str]]]: """ Organize prompt messages :param context: :param app_record: app record :param model_config: model config entity :param prompt_template_entity: prompt template entity :param inputs: inputs :param files: files :param query: query :param memory: memory :return: """ # get prompt without memory and context if prompt_template_entity.prompt_type == PromptTemplateEntity.PromptType.SIMPLE: prompt_transform: Union[SimplePromptTransform, AdvancedPromptTransform] prompt_transform = SimplePromptTransform() prompt_messages, stop = prompt_transform.get_prompt( app_mode=AppMode.value_of(app_record.mode), prompt_template_entity=prompt_template_entity, inputs=inputs, query=query or "", files=files, context=context, memory=memory, model_config=model_config, ) else: memory_config = MemoryConfig(window=MemoryConfig.WindowConfig(enabled=False)) model_mode = ModelMode.value_of(model_config.mode) prompt_template: Union[CompletionModelPromptTemplate, list[ChatModelMessage]] if model_mode == ModelMode.COMPLETION: advanced_completion_prompt_template = prompt_template_entity.advanced_completion_prompt_template if not advanced_completion_prompt_template: raise InvokeBadRequestError("Advanced completion prompt template is required.") prompt_template = CompletionModelPromptTemplate(text=advanced_completion_prompt_template.prompt) if advanced_completion_prompt_template.role_prefix: memory_config.role_prefix = MemoryConfig.RolePrefix( user=advanced_completion_prompt_template.role_prefix.user, assistant=advanced_completion_prompt_template.role_prefix.assistant, ) else: if not prompt_template_entity.advanced_chat_prompt_template: raise InvokeBadRequestError("Advanced chat prompt template is required.") prompt_template = [] for message in prompt_template_entity.advanced_chat_prompt_template.messages: prompt_template.append(ChatModelMessage(text=message.text, role=message.role)) prompt_transform = AdvancedPromptTransform() prompt_messages = prompt_transform.get_prompt( prompt_template=prompt_template, inputs=inputs, query=query or "", files=files, context=context, memory_config=memory_config, memory=memory, model_config=model_config, ) stop = model_config.stop return prompt_messages, stop def direct_output( self, queue_manager: AppQueueManager, app_generate_entity: EasyUIBasedAppGenerateEntity, prompt_messages: list, text: str, stream: bool, usage: Optional[LLMUsage] = None, ) -> None: """ Direct output :param queue_manager: application queue manager :param app_generate_entity: app generate entity :param prompt_messages: prompt messages :param text: text :param stream: stream :param usage: usage :return: """ if stream: index = 0 for token in text: chunk = LLMResultChunk( model=app_generate_entity.model_conf.model, prompt_messages=prompt_messages, delta=LLMResultChunkDelta(index=index, message=AssistantPromptMessage(content=token)), ) queue_manager.publish(QueueLLMChunkEvent(chunk=chunk), PublishFrom.APPLICATION_MANAGER) index += 1 time.sleep(0.01) queue_manager.publish( QueueMessageEndEvent( llm_result=LLMResult( model=app_generate_entity.model_conf.model, prompt_messages=prompt_messages, message=AssistantPromptMessage(content=text), usage=usage or LLMUsage.empty_usage(), ), ), PublishFrom.APPLICATION_MANAGER, ) def _handle_invoke_result( self, invoke_result: Union[LLMResult, Generator[Any, None, None]], queue_manager: AppQueueManager, stream: bool, agent: bool = False, ) -> None: """ Handle invoke result :param invoke_result: invoke result :param queue_manager: application queue manager :param stream: stream :param agent: agent :return: """ if not stream and isinstance(invoke_result, LLMResult): self._handle_invoke_result_direct(invoke_result=invoke_result, queue_manager=queue_manager, agent=agent) elif stream and isinstance(invoke_result, Generator): self._handle_invoke_result_stream(invoke_result=invoke_result, queue_manager=queue_manager, agent=agent) else: raise NotImplementedError(f"unsupported invoke result type: {type(invoke_result)}") def _handle_invoke_result_direct( self, invoke_result: LLMResult, queue_manager: AppQueueManager, agent: bool ) -> None: """ Handle invoke result direct :param invoke_result: invoke result :param queue_manager: application queue manager :param agent: agent :return: """ queue_manager.publish( QueueMessageEndEvent( llm_result=invoke_result, ), PublishFrom.APPLICATION_MANAGER, ) def _handle_invoke_result_stream( self, invoke_result: Generator, queue_manager: AppQueueManager, agent: bool ) -> None: """ Handle invoke result :param invoke_result: invoke result :param queue_manager: application queue manager :param agent: agent :return: """ model: str = "" prompt_messages: list[PromptMessage] = [] text = "" usage = None for result in invoke_result: if not agent: queue_manager.publish(QueueLLMChunkEvent(chunk=result), PublishFrom.APPLICATION_MANAGER) else: queue_manager.publish(QueueAgentMessageEvent(chunk=result), PublishFrom.APPLICATION_MANAGER) text += result.delta.message.content if not model: model = result.model if not prompt_messages: prompt_messages = result.prompt_messages if result.delta.usage: usage = result.delta.usage if not usage: usage = LLMUsage.empty_usage() llm_result = LLMResult( model=model, prompt_messages=prompt_messages, message=AssistantPromptMessage(content=text), usage=usage ) queue_manager.publish( QueueMessageEndEvent( llm_result=llm_result, ), PublishFrom.APPLICATION_MANAGER, ) def moderation_for_inputs( self, *, app_id: str, tenant_id: str, app_generate_entity: AppGenerateEntity, inputs: Mapping[str, Any], query: str | None = None, message_id: str, ) -> tuple[bool, Mapping[str, Any], str]: """ Process sensitive_word_avoidance. :param app_id: app id :param tenant_id: tenant id :param app_generate_entity: app generate entity :param inputs: inputs :param query: query :param message_id: message id :return: """ moderation_feature = InputModeration() return moderation_feature.check( app_id=app_id, tenant_id=tenant_id, app_config=app_generate_entity.app_config, inputs=dict(inputs), query=query or "", message_id=message_id, trace_manager=app_generate_entity.trace_manager, ) def check_hosting_moderation( self, application_generate_entity: EasyUIBasedAppGenerateEntity, queue_manager: AppQueueManager, prompt_messages: list[PromptMessage], ) -> bool: """ Check hosting moderation :param application_generate_entity: application generate entity :param queue_manager: queue manager :param prompt_messages: prompt messages :return: """ hosting_moderation_feature = HostingModerationFeature() moderation_result = hosting_moderation_feature.check( application_generate_entity=application_generate_entity, prompt_messages=prompt_messages ) if moderation_result: self.direct_output( queue_manager=queue_manager, app_generate_entity=application_generate_entity, prompt_messages=prompt_messages, text="I apologize for any confusion, but I'm an AI assistant to be helpful, harmless, and honest.", stream=application_generate_entity.stream, ) return moderation_result def fill_in_inputs_from_external_data_tools( self, tenant_id: str, app_id: str, external_data_tools: list[ExternalDataVariableEntity], inputs: Mapping[str, Any], query: str, ) -> Mapping[str, Any]: """ Fill in variable inputs from external data tools if exists. :param tenant_id: workspace id :param app_id: app id :param external_data_tools: external data tools configs :param inputs: the inputs :param query: the query :return: the filled inputs """ external_data_fetch_feature = ExternalDataFetch() return external_data_fetch_feature.fetch( tenant_id=tenant_id, app_id=app_id, external_data_tools=external_data_tools, inputs=inputs, query=query ) def query_app_annotations_to_reply( self, app_record: App, message: Message, query: str, user_id: str, invoke_from: InvokeFrom ) -> Optional[MessageAnnotation]: """ Query app annotations to reply :param app_record: app record :param message: message :param query: query :param user_id: user id :param invoke_from: invoke from :return: """ annotation_reply_feature = AnnotationReplyFeature() return annotation_reply_feature.query( app_record=app_record, message=message, query=query, user_id=user_id, invoke_from=invoke_from )