### title: 010125-daysoff-assistant-api
### file: app.py

import asyncio
import os
import re
import time
import json
import torch

from api_docs_mck import api_docs_str

import chainlit as cl

from langchain import hub
from langchain.chains import LLMChain, APIChain
from langchain_core.prompts import PromptTemplate
#from langchain_community.llms import HuggingFaceHub
from langchain_huggingface import HuggingFaceEndpoint
from langchain.memory.buffer import ConversationBufferMemory


HUGGINGFACEHUB_API_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
BOOKING_ID = r'\b[A-Z]{6}\d{6}\b'
BOOKING_KEYWORDS = [
    "booking",
    "bestillingsnummer",
    "bookingen",
    "ordrenummer",
    "reservation",
    "rezerwacji",
    "bookingreferanse",
    "rezerwacja",
    "booket",
    "reservation number",
    "bestilling",
    "order number",
    "booking ID",
    "identyfikacyjny płatności"
]


daysoff_assistant_system_template = """
You are an AI customer support assistant for Daysoff. By default, 
you respond in Norwegian to {question}. In all other cases,
adapt to user's language and respond accordingly. 
You can retrieving booking information for a given booking ID. In addition,
you can inform on Daysoff's personvernspolicy and verticals."
Chat History: {chat_history}
Question: {question}
Answer:
"""
daysoff_assistant_system_prompt= PromptTemplate(
    input_variables=["chat_history", "question"],
    template=daysoff_assistant_system_template
)

api_url_template = """
Given the following API Documentation for Daysoff's official
booking information API: {api_docs}
Your task is to construct the most efficient API URL to answer
the user's question, ensuring the
call is optimized to include only the necessary information.
Question: {question}
API URL:
"""
api_url_prompt = PromptTemplate(input_variables=['api_docs', 'question'],
                                template=api_url_template)

api_response_template = """
IF and only IF {question} contains alphanumeric identifier:{BOOKING_ID},
with the API Documentation for Daysoff's official API: {api_docs}
in mind, and given this API URL: {api_url} for querying, here is the
response from Daysoff's API: {api_response}.
Please provide only information that directly addresses the user's question, 
omitting technical details like response format, and 
focusing on delivering the answer with clarity and conciseness, 
as if Daysoff's human customer service agent is providing this information themselves.
Her er informasjon om bestilligen:
"""

api_response_prompt = PromptTemplate(
    input_variables=['api_docs', 'question', 'api_url', 'api_response'],
    template=api_response_template
)

@cl.on_chat_start
def setup_multiple_chains():
    
    llm = HuggingFaceEndpoint(
    repo_id="google/gemma-2-2b-it", 
    huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN, 
    #max_new_tokens=512,  
    temperature=0.7,     
    task="text-generation"  
    )

    conversation_memory = ConversationBufferMemory(memory_key="chat_history",
                                                   max_len=200,
                                                   return_messages=True,
                                                   )
    llm_chain = LLMChain(llm=llm,
                         prompt=daysoff_assistant_booking_prompt,
                         memory=conversation_memory
                        )

    cl.user_session.set("llm_chain", llm_chain)

    api_chain = APIChain.from_llm_and_api_docs(
        llm=llm,
        api_docs=api_docs_str,
        api_url_prompt=api_url_prompt,
        api_response_prompt=api_response_prompt,
        verbose=True,
        limit_to_domains=None
    )


    cl.user_session.set("api_chain", api_chain)
    


@cl.on_message
async def handle_message(message: cl.Message):
    user_message = message.content.lower()
    llm_chain = cl.user_session.get("llm_chain")
    api_chain = cl.user_session.get("api_chain")

    def is_booking_query(user_message):
        match = re.search(r'\b[A-Z]{6}\d{6}\b', user_message)

        return match

    if match: 
        response = await api_chain.acall(user_message, 
                                         callbacks=[cl.AsyncLangchainCallbackHandler()])

    else: 
        response = await llm_chain.acall(user_message,
                                         callbacks=[cl.AsyncLangchainCallbackHandler()])

    response_key = "output" if "output" in response else "text"
    await cl.Message(response.get(response_key, "")).send()
    return message.content
    
# await cl.Message(response.get(response_key, "")).send()
# content=f"Her er informasjonen for bestillingsnummer {booking_id}:\n{booking_info}").send()