package com.smtaiserver.smtaiserver.javaai.jsonflow.node;

import java.util.List;

import com.smtaiserver.smtaiserver.core.SMTAIServerApp;
import com.smtaiserver.smtaiserver.javaai.SMTJavaAIError;
import com.smtaiserver.smtaiserver.javaai.jsonflow.core.SMTJsonFlowExecArg;
import com.smtaiserver.smtaiserver.javaai.jsonflow.core.SMTJsonFlowManager;
import com.smtaiserver.smtaiserver.javaai.jsonflow.core.SMTJsonFlowNodeOnlyOutput;
import com.smtaiserver.smtaiserver.javaai.llm.core.SMTLLMConnect;
import com.smtservlet.util.Json;
import com.smtservlet.util.SMTStatic;

public class SMTJsonFlowNodeLLM extends SMTJsonFlowNodeOnlyOutput
{
	protected String 	_llmId;
	protected String	_prompt;
	protected String	_argKey;
	
	@Override
	public void initInstane(SMTJsonFlowManager manager, Json jsonNode) throws Exception
	{
		super.initInstane(manager, jsonNode);
		
		List<Json> jsonGroupParams = jsonNode.getJsonPath("data|group_params|", false).asJsonList();
		for(Json jsonParams : jsonGroupParams)
		{
			String type = jsonParams.getJsonPath("params|0|type", false).asString();
			if("llm_model".equals(type))
			{
				_llmId = jsonParams.getJsonPath("params|0|value", false).asString();
			}
			else if("textarea".equals(type))
			{
				_prompt = jsonParams.getJsonPath("params|0|value", false).asString();
			}
			else if("input".equals(type))
			{
				_argKey = jsonParams.getJsonPath("params|0|value", false).asString();
			}
		}
		
		if(SMTStatic.isNullOrEmpty(_llmId))
			throw new Exception("can' find llmId");
		
		if(SMTStatic.isNullOrEmpty(_prompt))
			throw new Exception("can' find prompt");
		
		if(SMTStatic.isNullOrEmpty(_argKey))
			throw new Exception("can' find _argKey");
	}
	
	@Override
	public SMTJavaAIError executeFlowNode(SMTJsonFlowExecArg execArg) throws Exception 
	{
		SMTLLMConnect llm = SMTAIServerApp.getApp().allocLLMConnect(_llmId);
		try
		{
			String realPrompt = SMTAIServerApp.getApp().convGlobalMacroString(_prompt, execArg._jsonArgs);
			execArg._tranReq.traceLLMDebug("call llm : " +_llmId + "\n" + realPrompt);
			
			String result = llm.callWithMessage(null, realPrompt, execArg._tranReq);
			execArg._tranReq.traceLLMDebug("call llm result:\n" + result);
			execArg._jsonArgs.set(_argKey, result);
		}
		finally
		{
			llm.close();
		}
		
		
		return super.executeFlowNode(execArg);
	}
}