Update app.py

app.py

@@ -8,16 +8,20 @@ import time
 import json
 import torch
 
-from api_docs_mck import daysoff_api_docs_str
+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_community.llms import HuggingFaceHub
+from langchain_huggingface import HuggingFaceEndpoint
 from langchain.memory.buffer import ConversationBufferMemory
 
+
+HUGGINGFACEHUB_API_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+
 daysoff_assistant_booking_template = """
 You are a customer support assistant for Daysoff.no. Your expertise is
 retrieving booking information for a given booking ID."
@@ -32,17 +36,17 @@ daysoff_assistant_booking_prompt= PromptTemplate(
 
 api_url_template = """
 Given the following API Documentation for Daysoff's official
-booking information API: {daysoff_api_docs}
+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_prompt = PromptTemplate(input_variables=['daysoff_api_docs', 'question'],
+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: {daysoff_api_docs}
+With the API Documentation for Daysoff's official API: {api_docs}
 and the specific user question: {question} in mind,
 and given this API URL: {base_url} for querying, here is the
 response from Daysoff's API: {response}.
@@ -52,16 +56,24 @@ as if a human customer service agent is providing this information.
 Adapt to user's language. By default, you speak Norwegian.
 Booking information:
 """
-api_response_prompt = PromptTemplate(input_variables=['daysoff_api_docs',
-                                                      'question',
+api_response_prompt = PromptTemplate(input_variables=['question',
+                                                      'api_docs',
                                                       'base_url',
                                                       'response'],
                                      template=api_response_template)
 
 @cl.on_chat_start
 def setup_multiple_chains():
-    llm = HuggingFaceHub(repo_id="google/gemma-2-2b-it") 
-                    
+    #llm = HuggingFaceEndpoint(repo_id="google/gemma-2-2b-it")
+
+    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,
@@ -75,7 +87,7 @@ def setup_multiple_chains():
 
     api_chain = APIChain.from_llm_and_api_docs(
         llm=llm,
-        api_docs=daysoff_api_docs_str,
+        api_docs=api_docs_str,
         api_url_prompt=api_url_prompt,
         api_response_prompt=api_response_prompt,
         verbose=True,
@@ -83,7 +95,7 @@ def setup_multiple_chains():
 
     cl.user_session.set("api_chain", api_chain)
 
-BOOKING_ID = r'\b[A-Z]{6}\d{6}\b' 
+BOOKING_ID = r'\b[A-Z]{6}\d{6}\b'
 
 BOOKING_KEYWORDS = [
     "booking",
@@ -91,16 +103,16 @@ BOOKING_KEYWORDS = [
     "bookingen",
     "ordrenummer",
     "reservation",
-    "rezerwacji", 
-    "bookingreferanse", 
-    "rezerwacja", 
-    "logg inn", 
-    "booket", 
+    "rezerwacji",
+    "bookingreferanse",
+    "rezerwacja",
+    "logg inn",
+    "booket",
     "reservation number",
-    "bestilling", 
+    "bestilling",
     "order number",
     "booking ID",
-    "identyfikacyjny płatności"  
+    "identyfikacyjny płatności"
 ]
 
 @cl.on_message
@@ -121,7 +133,7 @@ async def handle_message(message: cl.Message):
     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