Underground-Digital/Workflow-Engine

1from typing import Optional2 3from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity4from core.memory.token_buffer_memory import TokenBufferMemory5from core.model_manager import ModelInstance6from core.model_runtime.entities.message_entities import PromptMessage7from core.model_runtime.entities.model_entities import ModelPropertyKey8from core.prompt.entities.advanced_prompt_entities import MemoryConfig9 10 11class PromptTransform:12    def _append_chat_histories(13        self,14        memory: TokenBufferMemory,15        memory_config: MemoryConfig,16        prompt_messages: list[PromptMessage],17        model_config: ModelConfigWithCredentialsEntity,18    ) -> list[PromptMessage]:19        rest_tokens = self._calculate_rest_token(prompt_messages, model_config)20        histories = self._get_history_messages_list_from_memory(memory, memory_config, rest_tokens)21        prompt_messages.extend(histories)22 23        return prompt_messages24 25    def _calculate_rest_token(26        self, prompt_messages: list[PromptMessage], model_config: ModelConfigWithCredentialsEntity27    ) -> int:28        rest_tokens = 200029 30        model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)31        if model_context_tokens:32            model_instance = ModelInstance(33                provider_model_bundle=model_config.provider_model_bundle, model=model_config.model34            )35 36            curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages)37 38            max_tokens = 039            for parameter_rule in model_config.model_schema.parameter_rules:40                if parameter_rule.name == "max_tokens" or (41                    parameter_rule.use_template and parameter_rule.use_template == "max_tokens"42                ):43                    max_tokens = (44                        model_config.parameters.get(parameter_rule.name)45                        or model_config.parameters.get(parameter_rule.use_template)46                    ) or 047 48            rest_tokens = model_context_tokens - max_tokens - curr_message_tokens49            rest_tokens = max(rest_tokens, 0)50 51        return rest_tokens52 53    def _get_history_messages_from_memory(54        self,55        memory: TokenBufferMemory,56        memory_config: MemoryConfig,57        max_token_limit: int,58        human_prefix: Optional[str] = None,59        ai_prefix: Optional[str] = None,60    ) -> str:61        """Get memory messages."""62        kwargs = {"max_token_limit": max_token_limit}63 64        if human_prefix:65            kwargs["human_prefix"] = human_prefix66 67        if ai_prefix:68            kwargs["ai_prefix"] = ai_prefix69 70        if memory_config.window.enabled and memory_config.window.size is not None and memory_config.window.size > 0:71            kwargs["message_limit"] = memory_config.window.size72 73        return memory.get_history_prompt_text(**kwargs)74 75    def _get_history_messages_list_from_memory(76        self, memory: TokenBufferMemory, memory_config: MemoryConfig, max_token_limit: int77    ) -> list[PromptMessage]:78        """Get memory messages."""79        return memory.get_history_prompt_messages(80            max_token_limit=max_token_limit,81            message_limit=memory_config.window.size82            if (83                memory_config.window.enabled and memory_config.window.size is not None and memory_config.window.size > 084            )85            else None,86        )87