import json import logging import re from collections.abc import Sequence from typing import Any, NotRequired, Protocol, TypedDict, cast import json_repair from sqlalchemy import select from core.app.app_config.entities import ModelConfig from core.llm_generator.entities import RuleCodeGeneratePayload, RuleGeneratePayload, RuleStructuredOutputPayload from core.llm_generator.output_parser.rule_config_generator import RuleConfigGeneratorOutputParser from core.llm_generator.output_parser.suggested_questions_after_answer import SuggestedQuestionsAfterAnswerOutputParser from core.llm_generator.prompts import ( CONVERSATION_TITLE_PROMPT, GENERATOR_QA_PROMPT, JAVASCRIPT_CODE_GENERATOR_PROMPT_TEMPLATE, LLM_MODIFY_CODE_SYSTEM, LLM_MODIFY_PROMPT_SYSTEM, PYTHON_CODE_GENERATOR_PROMPT_TEMPLATE, SYSTEM_STRUCTURED_OUTPUT_GENERATE, WORKFLOW_RULE_CONFIG_PROMPT_GENERATE_TEMPLATE, ) from core.model_manager import ModelManager from core.ops.entities.trace_entity import TraceTaskName from core.ops.ops_trace_manager import TraceQueueManager, TraceTask from core.ops.utils import measure_time from core.prompt.utils.prompt_template_parser import PromptTemplateParser from extensions.ext_database import db from extensions.ext_storage import storage from graphon.enums import WorkflowNodeExecutionMetadataKey from graphon.model_runtime.entities.llm_entities import LLMResult from graphon.model_runtime.entities.message_entities import PromptMessage, SystemPromptMessage, UserPromptMessage from graphon.model_runtime.entities.model_entities import ModelType from graphon.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError from models import App, Message, WorkflowNodeExecutionModel from models.workflow import Workflow logger = logging.getLogger(__name__) class SuggestedQuestionsModelConfig(TypedDict): provider: str name: str completion_params: NotRequired[dict[str, object]] def _normalize_completion_params(completion_params: dict[str, object]) -> tuple[dict[str, object], list[str]]: """ Normalize raw completion params into invocation parameters and stop sequences. This mirrors the app-model access path by separating ``stop`` from provider parameters before invocation, then drops non-positive token limits because some plugin-backed models reject ``0`` after mapping ``max_tokens`` to their provider-specific output-token field. """ normalized_parameters = dict(completion_params) stop_value = normalized_parameters.pop("stop", []) if isinstance(stop_value, list) and all(isinstance(item, str) for item in stop_value): stop = stop_value else: stop = [] for token_limit_key in ("max_tokens", "max_output_tokens"): token_limit = normalized_parameters.get(token_limit_key) if isinstance(token_limit, int | float) and token_limit <= 0: normalized_parameters.pop(token_limit_key, None) return normalized_parameters, stop class WorkflowServiceInterface(Protocol): def get_draft_workflow(self, app_model: App, workflow_id: str | None = None) -> Workflow | None: pass def get_node_last_run(self, app_model: App, workflow: Workflow, node_id: str) -> WorkflowNodeExecutionModel | None: pass class CodeGenerateResultDict(TypedDict): code: str language: str error: str class StructuredOutputResultDict(TypedDict): output: str error: str class LLMGenerator: @classmethod def generate_conversation_name( cls, tenant_id: str, query, conversation_id: str | None = None, app_id: str | None = None ): prompt = CONVERSATION_TITLE_PROMPT if len(query) > 2000: query = query[:300] + "...[TRUNCATED]..." + query[-300:] query = query.replace("\n", " ") prompt += query + "\n" model_manager = ModelManager.for_tenant(tenant_id=tenant_id) model_instance = model_manager.get_default_model_instance( tenant_id=tenant_id, model_type=ModelType.LLM, ) prompts = [UserPromptMessage(content=prompt)] with measure_time() as timer: response: LLMResult = model_instance.invoke_llm( prompt_messages=list(prompts), model_parameters={"max_tokens": 500, "temperature": 1}, stream=False ) answer = response.message.get_text_content() if answer == "": return "" try: result_dict = json.loads(answer) except json.JSONDecodeError: result_dict = json_repair.loads(answer) if not isinstance(result_dict, dict): answer = query else: output = result_dict.get("Your Output") if isinstance(output, str) and output.strip(): answer = output.strip() else: answer = query name = answer.strip() if len(name) > 75: name = name[:75] + "..." # get tracing instance trace_manager = TraceQueueManager(app_id=app_id) trace_manager.add_trace_task( TraceTask( TraceTaskName.GENERATE_NAME_TRACE, conversation_id=conversation_id, generate_conversation_name=name, inputs=prompt, timer=timer, tenant_id=tenant_id, ) ) return name @classmethod def generate_suggested_questions_after_answer( cls, tenant_id: str, histories: str, *, instruction_prompt: str | None = None, model_config: object | None = None, ) -> Sequence[str]: output_parser = SuggestedQuestionsAfterAnswerOutputParser(instruction_prompt=instruction_prompt) format_instructions = output_parser.get_format_instructions() prompt_template = PromptTemplateParser(template="{{histories}}\n{{format_instructions}}\nquestions:\n") prompt = prompt_template.format({"histories": histories, "format_instructions": format_instructions}) try: model_manager = ModelManager.for_tenant(tenant_id=tenant_id) configured_model = cast(dict[str, object], model_config) if isinstance(model_config, dict) else {} provider = configured_model.get("provider") model_name = configured_model.get("name") use_configured_model = False if isinstance(provider, str) and provider and isinstance(model_name, str) and model_name: try: model_instance = model_manager.get_model_instance( tenant_id=tenant_id, model_type=ModelType.LLM, provider=provider, model=model_name, ) use_configured_model = True except Exception: logger.warning( "Failed to use configured suggested-questions model %s/%s, fallback to default model", provider, model_name, exc_info=True, ) model_instance = model_manager.get_default_model_instance( tenant_id=tenant_id, model_type=ModelType.LLM, ) else: model_instance = model_manager.get_default_model_instance( tenant_id=tenant_id, model_type=ModelType.LLM, ) except InvokeAuthorizationError: return [] prompt_messages = [UserPromptMessage(content=prompt)] questions: Sequence[str] = [] try: configured_completion_params = configured_model.get("completion_params") if use_configured_model and isinstance(configured_completion_params, dict): model_parameters, stop = _normalize_completion_params(configured_completion_params) elif use_configured_model: model_parameters = {} stop = [] else: # Default-model generation keeps the built-in suggested-questions tuning. model_parameters = { "max_tokens": 2560, "temperature": 0.0, } stop = [] response: LLMResult = model_instance.invoke_llm( prompt_messages=list(prompt_messages), model_parameters=model_parameters, stop=stop, stream=False, ) text_content = response.message.get_text_content() questions = output_parser.parse(text_content) if text_content else [] except Exception: logger.exception("Failed to generate suggested questions after answer") questions = [] return questions @classmethod def generate_rule_config(cls, tenant_id: str, args: RuleGeneratePayload): output_parser = RuleConfigGeneratorOutputParser() error = "" error_step = "" rule_config = {"prompt": "", "variables": [], "opening_statement": "", "error": ""} model_parameters = args.model_config_data.completion_params if args.no_variable: prompt_template = PromptTemplateParser(WORKFLOW_RULE_CONFIG_PROMPT_GENERATE_TEMPLATE) prompt_generate = prompt_template.format( inputs={ "TASK_DESCRIPTION": args.instruction, }, remove_template_variables=False, ) prompt_messages = [UserPromptMessage(content=prompt_generate)] model_manager = ModelManager.for_tenant(tenant_id=tenant_id) model_instance = model_manager.get_model_instance( tenant_id=tenant_id, model_type=ModelType.LLM, provider=args.model_config_data.provider, model=args.model_config_data.name, ) try: response: LLMResult = model_instance.invoke_llm( prompt_messages=list(prompt_messages), model_parameters=model_parameters, stream=False ) rule_config["prompt"] = response.message.get_text_content() except InvokeError as e: error = str(e) error_step = "generate rule config" except Exception as e: logger.exception("Failed to generate rule config, model: %s", args.model_config_data.name) error = str(e) error_step = "generate rule config" rule_config["error"] = f"Failed to {error_step}. Error: {error}" if error else "" return rule_config # get rule config prompt, parameter and statement prompt_generate, parameter_generate, statement_generate = output_parser.get_format_instructions() prompt_template = PromptTemplateParser(prompt_generate) parameter_template = PromptTemplateParser(parameter_generate) statement_template = PromptTemplateParser(statement_generate) # format the prompt_generate_prompt prompt_generate_prompt = prompt_template.format( inputs={ "TASK_DESCRIPTION": args.instruction, }, remove_template_variables=False, ) prompt_messages = [UserPromptMessage(content=prompt_generate_prompt)] # get model instance model_manager = ModelManager.for_tenant(tenant_id=tenant_id) model_instance = model_manager.get_model_instance( tenant_id=tenant_id, model_type=ModelType.LLM, provider=args.model_config_data.provider, model=args.model_config_data.name, ) try: try: # the first step to generate the task prompt prompt_content: LLMResult = model_instance.invoke_llm( prompt_messages=list(prompt_messages), model_parameters=model_parameters, stream=False ) except InvokeError as e: error = str(e) error_step = "generate prefix prompt" rule_config["error"] = f"Failed to {error_step}. Error: {error}" if error else "" return rule_config rule_config["prompt"] = prompt_content.message.get_text_content() parameter_generate_prompt = parameter_template.format( inputs={ "INPUT_TEXT": prompt_content.message.get_text_content(), }, remove_template_variables=False, ) parameter_messages = [UserPromptMessage(content=parameter_generate_prompt)] # the second step to generate the task_parameter and task_statement statement_generate_prompt = statement_template.format( inputs={ "TASK_DESCRIPTION": args.instruction, "INPUT_TEXT": prompt_content.message.get_text_content(), }, remove_template_variables=False, ) statement_messages = [UserPromptMessage(content=statement_generate_prompt)] try: parameter_content: LLMResult = model_instance.invoke_llm( prompt_messages=list(parameter_messages), model_parameters=model_parameters, stream=False ) rule_config["variables"] = re.findall(r'"\s*([^"]+)\s*"', parameter_content.message.get_text_content()) except InvokeError as e: error = str(e) error_step = "generate variables" try: statement_content: LLMResult = model_instance.invoke_llm( prompt_messages=list(statement_messages), model_parameters=model_parameters, stream=False ) rule_config["opening_statement"] = statement_content.message.get_text_content() except InvokeError as e: error = str(e) error_step = "generate conversation opener" except Exception as e: logger.exception("Failed to generate rule config, model: %s", args.model_config_data.name) error = str(e) error_step = "handle unexpected exception" rule_config["error"] = f"Failed to {error_step}. Error: {error}" if error else "" return rule_config @classmethod def generate_code( cls, tenant_id: str, args: RuleCodeGeneratePayload, ) -> CodeGenerateResultDict: if args.code_language == "python": prompt_template = PromptTemplateParser(PYTHON_CODE_GENERATOR_PROMPT_TEMPLATE) else: prompt_template = PromptTemplateParser(JAVASCRIPT_CODE_GENERATOR_PROMPT_TEMPLATE) prompt = prompt_template.format( inputs={ "INSTRUCTION": args.instruction, "CODE_LANGUAGE": args.code_language, }, remove_template_variables=False, ) model_manager = ModelManager.for_tenant(tenant_id=tenant_id) model_instance = model_manager.get_model_instance( tenant_id=tenant_id, model_type=ModelType.LLM, provider=args.model_config_data.provider, model=args.model_config_data.name, ) prompt_messages = [UserPromptMessage(content=prompt)] model_parameters = args.model_config_data.completion_params try: response: LLMResult = model_instance.invoke_llm( prompt_messages=list(prompt_messages), model_parameters=model_parameters, stream=False ) generated_code = response.message.get_text_content() return {"code": generated_code, "language": args.code_language, "error": ""} except InvokeError as e: error = str(e) return {"code": "", "language": args.code_language, "error": f"Failed to generate code. Error: {error}"} except Exception as e: logger.exception( "Failed to invoke LLM model, model: %s, language: %s", args.model_config_data.name, args.code_language ) return {"code": "", "language": args.code_language, "error": f"An unexpected error occurred: {str(e)}"} @classmethod def generate_qa_document(cls, tenant_id: str, query, document_language: str): prompt = GENERATOR_QA_PROMPT.format(language=document_language) model_manager = ModelManager.for_tenant(tenant_id=tenant_id) model_instance = model_manager.get_default_model_instance( tenant_id=tenant_id, model_type=ModelType.LLM, ) prompt_messages: list[PromptMessage] = [SystemPromptMessage(content=prompt), UserPromptMessage(content=query)] # Explicitly use the non-streaming overload result = model_instance.invoke_llm( prompt_messages=prompt_messages, model_parameters={"temperature": 0.01, "max_tokens": 2000}, stream=False, ) # Runtime type check for overload narrowing. if not isinstance(result, LLMResult): raise TypeError("Expected LLMResult when stream=False") response = result answer = response.message.get_text_content() return answer.strip() @classmethod def generate_structured_output( cls, tenant_id: str, args: RuleStructuredOutputPayload ) -> StructuredOutputResultDict: model_manager = ModelManager.for_tenant(tenant_id=tenant_id) model_instance = model_manager.get_model_instance( tenant_id=tenant_id, model_type=ModelType.LLM, provider=args.model_config_data.provider, model=args.model_config_data.name, ) prompt_messages = [ SystemPromptMessage(content=SYSTEM_STRUCTURED_OUTPUT_GENERATE), UserPromptMessage(content=args.instruction), ] model_parameters = args.model_config_data.completion_params try: response: LLMResult = model_instance.invoke_llm( prompt_messages=list(prompt_messages), model_parameters=model_parameters, stream=False ) raw_content = response.message.get_text_content() try: parsed_content = json.loads(raw_content) except json.JSONDecodeError: parsed_content = json_repair.loads(raw_content) if not isinstance(parsed_content, dict | list): raise ValueError(f"Failed to parse structured output from llm: {raw_content}") generated_json_schema = json.dumps(parsed_content, indent=2, ensure_ascii=False) return {"output": generated_json_schema, "error": ""} except InvokeError as e: error = str(e) return {"output": "", "error": f"Failed to generate JSON Schema. Error: {error}"} except Exception as e: logger.exception("Failed to invoke LLM model, model: %s", args.model_config_data.name) return {"output": "", "error": f"An unexpected error occurred: {str(e)}"} @staticmethod def instruction_modify_legacy( tenant_id: str, flow_id: str, current: str, instruction: str, model_config: ModelConfig, ideal_output: str | None, ): last_run: Message | None = db.session.scalar( select(Message).where(Message.app_id == flow_id).order_by(Message.created_at.desc()).limit(1) ) if not last_run: return LLMGenerator.__instruction_modify_common( tenant_id=tenant_id, model_config=model_config, last_run=None, current=current, error_message="", instruction=instruction, node_type="llm", ideal_output=ideal_output, ) last_run_dict = { "query": last_run.query, "answer": last_run.answer, "error": last_run.error, } return LLMGenerator.__instruction_modify_common( tenant_id=tenant_id, model_config=model_config, last_run=last_run_dict, current=current, error_message=str(last_run.error), instruction=instruction, node_type="llm", ideal_output=ideal_output, ) @staticmethod def instruction_modify_workflow( tenant_id: str, flow_id: str, node_id: str, current: str, instruction: str, model_config: ModelConfig, ideal_output: str | None, workflow_service: WorkflowServiceInterface, ): session = db.session() app: App | None = session.scalar(select(App).where(App.id == flow_id).limit(1)) if not app: raise ValueError("App not found.") workflow = workflow_service.get_draft_workflow(app_model=app) if not workflow: raise ValueError("Workflow not found for the given app model.") last_run = workflow_service.get_node_last_run(app_model=app, workflow=workflow, node_id=node_id) try: node_type = cast(WorkflowNodeExecutionModel, last_run).node_type except Exception: try: node_type = [it for it in workflow.graph_dict["graph"]["nodes"] if it["id"] == node_id][0]["data"][ "type" ] except Exception: node_type = "llm" if not last_run: # Node is not executed yet return LLMGenerator.__instruction_modify_common( tenant_id=tenant_id, model_config=model_config, last_run=None, current=current, error_message="", instruction=instruction, node_type=node_type, ideal_output=ideal_output, ) def agent_log_of(node_execution: WorkflowNodeExecutionModel) -> Sequence: raw_agent_log = node_execution.execution_metadata_dict.get(WorkflowNodeExecutionMetadataKey.AGENT_LOG, []) if not raw_agent_log: return [] return [ { "status": event["status"], "error": event["error"], "data": event["data"], } for event in raw_agent_log ] inputs = last_run.load_full_inputs(session, storage) last_run_dict = { "inputs": inputs, "status": last_run.status, "error": last_run.error, "agent_log": agent_log_of(last_run), } return LLMGenerator.__instruction_modify_common( tenant_id=tenant_id, model_config=model_config, last_run=last_run_dict, current=current, error_message=last_run.error, instruction=instruction, node_type=last_run.node_type, ideal_output=ideal_output, ) @staticmethod def __instruction_modify_common( tenant_id: str, model_config: ModelConfig, last_run: dict[str, Any] | None, current: str | None, error_message: str | None, instruction: str, node_type: str, ideal_output: str | None, ): LAST_RUN = "{{#last_run#}}" CURRENT = "{{#current#}}" ERROR_MESSAGE = "{{#error_message#}}" injected_instruction = instruction if LAST_RUN in injected_instruction: injected_instruction = injected_instruction.replace(LAST_RUN, json.dumps(last_run)) if CURRENT in injected_instruction: injected_instruction = injected_instruction.replace(CURRENT, current or "null") if ERROR_MESSAGE in injected_instruction: injected_instruction = injected_instruction.replace(ERROR_MESSAGE, error_message or "null") model_instance = ModelManager.for_tenant(tenant_id=tenant_id).get_model_instance( tenant_id=tenant_id, model_type=ModelType.LLM, provider=model_config.provider, model=model_config.name, ) match node_type: case "llm" | "agent": system_prompt = LLM_MODIFY_PROMPT_SYSTEM case "code": system_prompt = LLM_MODIFY_CODE_SYSTEM case _: system_prompt = LLM_MODIFY_PROMPT_SYSTEM prompt_messages = [ SystemPromptMessage(content=system_prompt), UserPromptMessage( content=json.dumps( { "current": current, "last_run": last_run, "instruction": injected_instruction, "ideal_output": ideal_output, } ) ), ] model_parameters = {"temperature": 0.4} try: response: LLMResult = model_instance.invoke_llm( prompt_messages=list(prompt_messages), model_parameters=model_parameters, stream=False ) generated_raw = response.message.get_text_content() first_brace = generated_raw.find("{") last_brace = generated_raw.rfind("}") if first_brace == -1 or last_brace == -1 or last_brace < first_brace: raise ValueError(f"Could not find a valid JSON object in response: {generated_raw}") json_str = generated_raw[first_brace : last_brace + 1] data = json_repair.loads(json_str) if not isinstance(data, dict): raise TypeError(f"Expected a JSON object, but got {type(data).__name__}") return data except InvokeError as e: error = str(e) return {"error": f"Failed to generate code. Error: {error}"} except Exception as e: logger.exception("Failed to invoke LLM model, model: %s", json.dumps(model_config.name), exc_info=True) return {"error": f"An unexpected error occurred: {str(e)}"}