module interface
    

export agent, addNewMessage, clearMessage, removeLatestMsg, generatePrompt_tokenPrefix,
        generatePrompt_tokenSuffix

using JSON3, DataStructures, Dates

# ---------------------------------------------------------------------------- #
#                              pythoncall setting                              #
# ---------------------------------------------------------------------------- #
# Ref: https://github.com/JuliaPy/PythonCall.jl/issues/252
# by setting the following variables, PythonCall will use system python or conda python and 
# packages installed by system or conda
# if these setting are not set (comment out), PythonCall will use its own python and package that
# installed by CondaPkg
ENV["JULIA_CONDAPKG_BACKEND"] = "Null"
systemPython = split(read(`which python`, String), "\n")[1]
ENV["JULIA_PYTHONCALL_EXE"] = systemPython # find python location with $> which python  ex. raw"/root/conda/bin/python"

using PythonCall
const py_agents = PythonCall.pynew()
const py_llms = PythonCall.pynew()
function __init__()
    # PythonCall.pycopy!(py_cv2, pyimport("cv2"))

    # equivalent to     from urllib.request import urlopen    in python
    PythonCall.pycopy!(py_agents, pyimport("langchain.agents"))
    PythonCall.pycopy!(py_llms, pyimport("langchain.llms"))
end

#------------------------------------------------------------------------------------------------100



@kwdef mutable struct agent
    availableRole=["system", "user", "assistant"]
    maxUserMsg::Int= 10
    llmAIRequestTopic_openblas = "llm/openblas/request"
    llmAIRequestTopic_gpu = "llm/api/v0.0.1/gpu/request"
    self_llmReceiveTopic = "chatbothub/llm/respond"

    """ Dict(Role=> Content)  ;   Role can be system, user, assistant
      Example:
        messages=[
              Dict(:role=>"system", :content=> "You are a helpful assistant."),
              Dict(:role=>"assistant", :content=> "How may I help you"),
              Dict(:role=>"user", :content=> "Hello, how are you"),
              ]
    """
    # Ref: Chat prompt format https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGML/discussions/3
    messages=[Dict(:role=>"system", :content=> "You are a helpful assistant.", :timestamp=> Dates.now()),]
end

@kwdef mutable struct agentLangchain
  availableRole=["system", "user", "assistant"]
  maxUserMsg::Int= 10
  llmAIRequestTopic_openblas = "llm/openblas/request"
  llmAIRequestTopic_gpu = "llm/api/v0.0.1/gpu/request"
  self_llmReceiveTopic = "chatbothub/llm/respond"

  """ Dict(Role=> Content)  ;   Role can be system, user, assistant
    Example:
      messages=[
            Dict(:role=>"system", :content=> "You are a helpful assistant."),
            Dict(:role=>"assistant", :content=> "How may I help you"),
            Dict(:role=>"user", :content=> "Hello, how are you"),
            ]
  """
  # Ref: https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGML/discussions/3
  # 
  messages=[Dict(:role=>"system", :content=> "You are a helpful assistant.", :timestamp=> Dates.now()),]
end

""" 
  add new message to agent
# Example
```jldoctest
julia> addNewMessage(agent1, "user", "Where should I go to buy snacks")
````
"""
function addNewMessage(a::agent, role::String, content::String)
  if role ∉ a.availableRole   # guard against typo
    error("role is not in agent.availableRole")
  end

  # check whether user messages exceed limit
  userMsg = 0
  for i in a.messages
    if i[:role] == "user"
      userMsg += 1
    end
  end
  messageleft = 0

  if userMsg > a.maxUserMsg # delete all conversation
    clearMessage(a)
    messageleft = a.maxUserMsg
  else
    userMsg += 1
    d = Dict(:role=> role, :content=> content, :timestamp=> Dates.now())
    push!(a.messages, d)
    messageleft = a.maxUserMsg - userMsg
  end

  return messageleft
end


function clearMessage(a::agent)
  for i in eachindex(a.messages)
    if length(a.messages) > 1 # system instruction will NOT be deleted
      pop!(a.messages)
    else
      break
    end
  end
end

function removeLatestMsg(a::agent)
  if length(a.messages) > 1
    pop!(a.messages)
  end
end

function generatePrompt_tokenSuffix(a::agent; 
                      userToken::String="[/INST]", assistantToken="[INST]", 
                      systemToken="[INST]<<SYS>> content <</SYS>>")
  prompt = nothing
  for msg in a.messages
    role = msg[:role]
    content = msg[:content]

    if role == "system"
      prompt = replace(systemToken, "content" => content) * " "
    elseif role == "user"
      prompt *= " " * content * " " * userToken
    elseif role == "assistant"
      prompt *= " " * content * " " * assistantToken
    else
      error("undefied condition role = $role")
    end
  end

  return prompt
end

function generatePrompt_tokenPrefix(a::agent; 
                      userToken::String="Q:", assistantToken="A:", 
                      systemToken="[INST]<<SYS>> content <</SYS>>")
  prompt = nothing
  for msg in a.messages
    role = msg[:role]
    content = msg[:content]

    if role == "system"
      prompt = replace(systemToken, "content" => content) * " "
    elseif role == "user"
      prompt *= userToken * " " * content * " "
    elseif role == "assistant"
      prompt *= assistantToken * " " * content * " "
    else
      error("undefied condition role = $role")
    end
  end

  return prompt
end

























































end # module