# @Author : Kinhong Tsui
# @Date : 2024/6/4 9:49
import pickle
import sys,os
import yaml
sys.path.append(os.getcwd())

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from nlpbridge.persistent.redis import RedisDB
from nlpbridge.text.template_manager import TemplateManager

system = '''Respond to the human as helpfully and accurately as possible. You have access to the following tools:

{tools}

Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).

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"
}}

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'''

human = '''{inputs}{history_messages}'''


prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("placeholder", "{chat_history}"),
        ("human", human),
    ]
)

with open('../../config.yaml', 'r') as file:
    config = yaml.safe_load(file)
db = RedisDB(config)

template_manager = TemplateManager(db)

prompt_map = {
    'template1': prompt
}

save = pickle.dumps(prompt)
print(save)
print(type(save))

print("=======================")
res_save = template_manager.save_template('chat_prompt_template', prompt_map)
print(f"save: {res_save}")

print("=======================")
res_load = template_manager.load_template('chat_prompt_template', ['template1'])
print(len(res_load))
print(type(res_load[0]))
print(f"load: {res_load[0]}")

# print("=======================")
# res_delete = template_manager.delete_template('test_chat_prompt', ['template11', 'template21'])
# print(f"delete: {res_delete}")

