Spaces:
Sleeping
Sleeping
import os | |
import streamlit as st | |
from dotenv import load_dotenv | |
from langchain_core.messages import AIMessage, HumanMessage | |
from langchain_community.llms import HuggingFaceEndpoint | |
from langchain_core.output_parsers import StrOutputParser | |
from langchain_core.prompts import ChatPromptTemplate | |
import fitz # PyMuPDF | |
load_dotenv() | |
api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN") | |
# Define the repository ID and task | |
repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1" | |
task = "text-generation" | |
# App config | |
st.set_page_config(page_title="Yatra Sevak.AI", page_icon="🌍") | |
st.title("Yatra Sevak.AI ✈️") | |
# Define the template outside the function | |
template = """ | |
You are a travel assistant chatbot your name is Yatra Sevak.AI designed to help users plan their trips and provide travel-related information. Here are some scenarios you should be able to handle: | |
1. Destination Information: Provide information about popular travel destinations. Offer insights on attractions, local cuisine, cultural highlights, weather conditions, and best times to visit. | |
2. Travel Tips: Offer practical travel tips and advice. Topics may include packing essentials, visa requirements, currency exchange, local customs, and safety tips. | |
3. Local Attractions: Suggest local attractions and points of interest based on the user's destination. Highlight must-see landmarks, museums, parks, and recreational activities. | |
4. Customer Service: Address customer service inquiries and provide assistance with travel-related issues. | |
5. Online Marketing & Sales Associate: Act as an Online Marketing & Sales Associate and try to suggest our services to the user. If you think that the user can be converted to a client, then ask for their contact information and also give them our contact information. | |
Please ensure responses are informative, accurate, and tailored to the user's queries and preferences. Use natural language to engage users and provide a seamless experience throughout their travel planning journey. | |
Chat history: | |
{chat_history} | |
User question: | |
{user_question} | |
""" | |
pdf_path = "YatraAI_Travel_Plans.pdf" | |
def extract_text_from_pdf(pdf_path): | |
doc = fitz.open(pdf_path) | |
text = "" | |
for page in doc: | |
text += page.get_text() | |
return text | |
pdf_text = extract_text_from_pdf(pdf_path) | |
# Function to get a response from the model | |
def get_response(user_query, chat_history): | |
# Initialize the Hugging Face Endpoint | |
llm = HuggingFaceEndpoint( | |
huggingfacehub_api_token=api_token, | |
repo_id=repo_id, | |
task=task | |
) | |
# Include the PDF content in the prompt | |
prompt_with_pdf = f"{template}\n\nCompany Info and Travel Plans:\n{pdf_text}" | |
chain = ChatPromptTemplate.from_template(prompt_with_pdf) | llm | StrOutputParser() | |
response = chain.invoke({ | |
"chat_history": chat_history, | |
"user_question": user_query, | |
}) | |
return response | |
# Initialize session state | |
if "chat_history" not in st.session_state: | |
st.session_state.chat_history = [ | |
AIMessage(content="Hello, I am Yatra Sevak.AI. How can I help you?"), | |
] | |
# Display chat history | |
for message in st.session_state.chat_history: | |
if isinstance(message, AIMessage): | |
with st.chat_message("AI"): | |
st.write(message.content) | |
elif isinstance(message, HumanMessage): | |
with st.chat_message("Human"): | |
st.write(message.content) | |
# User input | |
user_query = st.chat_input("Type your message here...") | |
if user_query is not None and user_query != "": | |
st.session_state.chat_history.append(HumanMessage(content=user_query)) | |
with st.chat_message("Human"): | |
st.markdown(user_query) | |
response = get_response(user_query, st.session_state.chat_history) | |
# Remove any unwanted prefixes from the response | |
response = response.replace("AI response:", "").replace("chat response:", "").replace("bot response:", "").strip() | |
with st.chat_message("AI"): | |
st.write(response) | |
st.session_state.chat_history.append(AIMessage(content=response)) | |