app.py

### 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",
    "logg inn",
    "booket",
    "reservation number",
    "bestilling",
    "order number",
    "booking ID",
    "identyfikacyjny płatności"
]


daysoff_assistant_booking_template = """
You are a Norwegian customer support AI assistant for Daysoff.no. By default, 
you respond in Norwegian to user question: {question}. In all other cases,
adapt to user's language and respond accordingly. 
You can retrieving booking information for a given booking ID and provide 
general helpful info about Daysoff's services."
Chat History: {chat_history}
Question: {question}
Answer:
"""
daysoff_assistant_booking_prompt= PromptTemplate(
    input_variables=["chat_history", "question"],
    template=daysoff_assistant_booking_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 = """
With the API Documentation for Daysoff's official API: {api_docs}
in mind, and IF {question} contains keywords from:{BOOKING_KEYWORDS},
and given this API URL: {api_url} for querying, here is the
response from Daysoff's API: {api_response}.
Please provide a summary 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 kundeservice is providing this information themselves.
Booking Info:
"""

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)
    
import re

def extract_booking_id(user_message):
# Extract booking ID from user message
    match = re.search(r'\b[A-Z]{6}\d{6}\b', user_message) 
    return match.group(0) if match else None
    
@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")

    # booking_id = extract_booking_id(user_message)

    
    #is_booking_query = any(
        #re.search(keyword, user_message, re.IGNORECASE) # re.search(keyword, r'\b[A-Z]{6}\d{6}\b', user_message, re.IGNORECASE)
        #for keyword in BOOKING_KEYWORDS + [BOOKING_ID]
    #)
    if any(keyword in user_message for keyword in ["BOOKING_KEYWORDS"]): # could be(?): if any(keyword in user_message for keyword in BOOKING_KEYWORDS + [BOOKING_ID]
            
        # <f any of the keywords are in user_message>,  use api_chain
        response = await api_chain.acall(user_message, 
                                         callbacks=[cl.AsyncLangchainCallbackHandler()])

    elif: 
        # <or if the booking_id regex format is discovered in user message>, use api_chain
        re.search(keyword, r'\b[A-Z]{6}\d{6}\b', user_message):  
            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()