YiemAgent/test/runtest.jl

using Revise      # remove when this package is completed
using YiemAgent, GeneralUtils, JSON3, MQTTClient, Dates, UUIDs, LibPQ, Base64, DataFrames
using Base.Threads

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

config = copy(JSON3.read("config.json"))

# instanceInternalTopic = config[:serviceInternalTopic][:mqtttopic] * "/1" 

# client, connection = MakeConnection(config[:mqttServerInfo][:broker], 
#                                     config[:mqttServerInfo][:port])

receiveUserMsgChannel = Channel{Dict}(4)
# receiveInternalMsgChannel = Channel{Dict}(4)
# println(typeof(connection))
# msgMeta = GeneralUtils.generate_msgMeta(
#   "N/A", 
#   replyTopic = config[:servicetopic][:mqtttopic]  # ask frontend reply to this instance_chat_topic
#   )  

function executeSQL(sql::T) where {T<:AbstractString}
  DBconnection = LibPQ.Connection("host=192.168.88.12 port=5432 dbname=yiem_wine_assistant user=yiem password=yiem@Postgres_0.0")
  result = LibPQ.execute(DBconnection, sql)
  close(DBconnection)
  return result
end

function text2textInstructLLM(prompt::String)
  msgMeta = GeneralUtils.generate_msgMeta(
    config[:externalservice][:text2textinstruct][:mqtttopic];
    msgPurpose= "inference",
    senderName= "yiemagent",
    senderId= string(uuid4()),
    receiverName= "text2textinstruct",
    mqttBrokerAddress= config[:mqttServerInfo][:broker],
    mqttBrokerPort= config[:mqttServerInfo][:port],
  )

  outgoingMsg = Dict(
    :msgMeta=> msgMeta,
    :payload=> Dict(
      :text=> prompt,
      :kwargs=> Dict(
        :max_tokens=> 2048,
        :stop=> ["<|eot_id|>"],
        :temperature=> 0.2,
        )
      )
    )

  _response = GeneralUtils.sendReceiveMqttMsg(outgoingMsg; timeout=120)
  response = _response[:response][:text]

  return response
end

function executeSQLVectorDB(sql)
  DBconnection = LibPQ.Connection("host=192.168.88.12 port=5433 dbname=SQLVectorDB user=yiemtechnologies@gmail.com password=yiem@Postgres_0.0")
  result = LibPQ.execute(DBconnection, sql)
  close(DBconnection)
  return result
end

function addSQLVectorDB(state)
  # get embedding of the query
  query = [state[:thoughtHistory][:question]]
  msgMeta = GeneralUtils.generate_msgMeta(
    config[:externalservice][:text2textinstruct][:mqtttopic];
    msgPurpose= "embedding",
    senderName= "yiemagent",
    senderId= string(uuid4()),
    receiverName= "text2textinstruct",
    mqttBrokerAddress= config[:mqttServerInfo][:broker],
    mqttBrokerPort= config[:mqttServerInfo][:port],
  )

  outgoingMsg = Dict(
    :msgMeta=> msgMeta,
    :payload=> Dict(
      :text=> query
      )
    )
  response = GeneralUtils.sendReceiveMqttMsg(outgoingMsg)
  embedding = response[:response][:embeddings][1]

  # check whether there is close enough vector already store in vectorDB. if no, add, else skip 
  sql =
    """
    SELECT *, embedding <-> '$embedding' as distance
    FROM sql_statement_repository
    ORDER BY distance LIMIT 1;
    """
  response = executeSQLVectorDB(sql)
  df = DataFrame(response)
  row, col = size(df)
  distance = row == 0 ? Inf : df[1, :distance]
  if row == 0 || distance > 10 # no close enough SQL stored in the database
    latestKey, _ = GeneralUtils.findHighestIndexKey(state[:thoughtHistory], :action_input)
    _sqlStatement = state[:thoughtHistory][latestKey]
    if occursin("SELECT", _sqlStatement)  # make sure it is an SQL statement before adding into DB
      sqlStatementBase64 = base64encode(_sqlStatement)
      sqlStatement = replace(_sqlStatement, "'"=>"")
      sql =
      """
      INSERT INTO sql_statement_repository (question, sql_statement, sql_statement_base64, embedding) VALUES ('$query', '$sqlStatement', '$sqlStatementBase64', '$embedding');
      """
      _ = executeSQLVectorDB(sql)
      println("--> added new SQL statement to vectorDB ", @__FILE__, " ", @__LINE__)
      println(sqlStatement)
    end
  end
end

function querySQLVectorDB(state)

  # provide similarSQL at the first time thinking only
  latestKey, _ = GeneralUtils.findHighestIndexKey(state[:thoughtHistory], :action_input)
  if latestKey === nothing
    # get embedding of the query
    query = [state[:thoughtHistory][:question]]
    msgMeta = GeneralUtils.generate_msgMeta(
      config[:externalservice][:text2textinstruct][:mqtttopic];
      msgPurpose= "embedding",
      senderName= "yiemagent",
      senderId= string(uuid4()),
      receiverName= "text2textinstruct",
      mqttBrokerAddress= config[:mqttServerInfo][:broker],
      mqttBrokerPort= config[:mqttServerInfo][:port],
    )

    outgoingMsg = Dict(
      :msgMeta=> msgMeta,
      :payload=> Dict(
        :text=> query
        )
      )
    response = GeneralUtils.sendReceiveMqttMsg(outgoingMsg)
    embedding = response[:response][:embeddings][1]

    # check whether there is close enough vector already store in vectorDB. if no, add, else skip 
    sql =
    """
    SELECT *, embedding <-> '$embedding' as distance
    FROM sql_statement_repository
    ORDER BY distance LIMIT 1;
    """
    response = executeSQLVectorDB(sql)
    df = DataFrame(response)
    row, col = size(df)
    distance = row == 0 ? Inf : df[1, :distance]
    if row != 0 && distance < 100
      # if there is usable SQL, return it.
      sqlStatementBase64 = df[1, :sql_statement_base64]
      sqlStatement = String(base64decode(sqlStatementBase64))
      println("--> getting SQL statement from vectorDB ", @__FILE__, " ", @__LINE__)
      println(sqlStatement)
      return sqlStatement
    else
      return nothing
    end
  end
  return nothing
end


# Instantiate an agent
a = YiemAgent.sommelier(
  text2textInstructLLM,
  executeSQL,
  querySQLVectorDB,
  addSQLVectorDB;
  name= "Janie",
  id= sessionId,  # agent instance id
)

function main()
  for i in 1:10
    userinput = ""
    for i in 1:3
      if userinput == ""
        println("")
        println("--> user input:")
        userinput = readline()
      else
        break
      end
    end
    response = YiemAgent.conversation(a, Dict(:text=> userinput))
    println("")
    println("--> assistant response: \n", response)
  end
end

main()

"""
  I'm joining a graduation party this evening. I want a bottle of dry white wine from the US. I'm ok with any price range.
  Well, the party is small casual with close friends and no food serving.
  I'm open to suggestion since I have no specific idea.
  I'm ok with any region.




  The input is instructions on how you want the presentation to be conducted.
"""





# response = YiemAgent.conversation(a, Dict(:text=> "newtopic",) )

# response = YiemAgent.conversation(a, Dict(:text=> "Hello, I would like a get a bottle of wine."))
# println("---> YiemAgent: ", response)

# response = YiemAgent.conversation(a, Dict(:text=> "I'm having a graduation party this evening. I'll pay at most 30 bucks."))







# # input = "query=\"off dry, medium tannin, French Rosé\""
# input = "Search the database for wine type: white, country: France, sweetness level: 1"
# YiemAgent.winestock(a, input)




# input = "French dry white wines with medium body"
# input = "query=\"medium-bodied dry white wine\""
# # input = "the customer is looking for a medium-bodied, dry white wine."
# result = YiemAgent.checkinventory(a, input)





error("test done")




function run_with_timeout(f, args...; timeout=5)
  result = Ref{Any}()
  task = Threads.@spawn try
      result[] = f(args...)
  catch e
      println("Task interrupted: ", e)
  end

  Timer(timeout) do _
      if !istaskdone(task)
          schedule(task, InterruptException())
          println("Task did not complete in time. Aborting.")
      else
          println("Task completed within the timeout.")
      end
  end

  return result[]
end

# Example function that takes arguments and returns a value
function example_function(x, y)
  sleep(10)  # Simulate a long-running task
  return x + y
end

# Example usage
result = run_with_timeout(example_function, 3, 4; timeout=5)
println("Result: ", result)