import os from dotenv import load_dotenv from langchain_community.retrievers import WikipediaRetriever from langchain_google_genai import ChatGoogleGenerativeAI from langchain.schema import HumanMessage, SystemMessage import gradio as gr import re # Load environment variables load_dotenv() # Get the API key from the environment variable api_key = os.getenv("GOOGLE_API_KEY") if not api_key: raise ValueError("GOOGLE_API_KEY environment variable is not set") os.environ["GOOGLE_API_KEY"] = api_key # Initiate WikipediaRetriever retriever = WikipediaRetriever() # Initiate chat model chat_model = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0.2) # Function to get book info from Wikipedia def get_book_info(title, author): query = f"{title} book by {author}" docs = retriever.get_relevant_documents(query) if docs: return docs[0].page_content return None # Function to extract book info def extract_book_info(book_content): messages = [ SystemMessage(content="You are an AI assistant who is an expert in analyzing and recommending books."), HumanMessage(content=f""" Based on the following information about this book: {book_content} Please provide these details using the following format, only using information already available to you. Do not use asterisks or any other markdown formatting: 1. Genre/Subject: [genre or subject of this book] 2. Synopsis (within 100 words maximum): [synopsis of this book] 3. Book Recommendations: - [Title of Recommended Book 1]: [brief explanation within 20 words maximum] - [Title of Recommended Book 2]: [brief explanation within 20 words maximum] - [Title of Recommended Book 3]: [brief explanation within 20 words maximum] - [Title of Recommended Book 4]: [brief explanation within 20 words maximum] - [Title of Recommended Book 5]: [brief explanation within 20 words maximum] If the information you have is not enough to provide recommendations, give the available information and state that you cannot give recommendations due to lack of information. """) ] response = chat_model(messages) return response.content # Function to ask other questions def chat_about_book(book_content, question): messages = [ SystemMessage(content=f"""You are an AI assistant who is an expert on this book:\n{book_content} Answer the questions about this book using only informations already available to you. Also, do not bold any of the text"""), HumanMessage(content=question) ] response = chat_model(messages) return response.content def parse_extracted_info(extracted_info): genre = "Genre not found" synopsis = "Synopsis not found" recommendations = "Recommendations not found" # Use regex to find each section genre_match = re.search(r'1\.\s*Genre/Subject:\s*(.*?)(?=\n2\.|\Z)', extracted_info, re.DOTALL) synopsis_match = re.search(r'2\.\s*Synopsis.*?:\s*(.*?)(?=\n3\.|\Z)', extracted_info, re.DOTALL) recommendations_match = re.search(r'3\.\s*(?:5\s*)?Book Recommendations?:?(.*)', extracted_info, re.DOTALL) if genre_match: genre = genre_match.group(1).strip() if synopsis_match: synopsis = synopsis_match.group(1).strip() if recommendations_match: recommendations = recommendations_match.group(1).strip() # Remove asterisks from recommendations recommendations = re.sub(r'\*+', '', recommendations) return genre, synopsis, recommendations def process_book(title, author): if not title or not author: return "Book not found", "No synopsis found", "No book recommendations available" try: book_info = get_book_info(title, author) if book_info: extracted_info = extract_book_info(book_info) genre, synopsis, recommendations = parse_extracted_info(extracted_info) return genre, synopsis, recommendations return "Book not found", "No synopsis found", "No book recommendations available" except Exception as e: print(f"Error in process_book: {str(e)}") return f"Error: {str(e)}", "", "" def chat(title, author, question): if not title or not author: return "You have not entered the book's title and author yet." book_info = get_book_info(title, author) if book_info: response = chat_about_book(book_info, question) # Remove stars from the response response = re.sub(r'\*+', '', response) return response return "Book not found. Please check the title and author." with gr.Blocks(title="Simple Book AI (Wikipedia-based)") as demo: gr.Markdown( """ # Simple Book AI (Wikipedia-based) Input the title and author(s) of the book to get the relevant information and book recommendations similar to your book. You could also ask other questions regarding your book. """ ) with gr.Row(): with gr.Column(scale=2): title_input = gr.Textbox(label="Title", placeholder="Enter the book title here...") author_input = gr.Textbox(label="Author(s)", placeholder="Enter the author's name here...") with gr.Row(): submit_book = gr.Button("Submit") clear_book = gr.Button("Clear") with gr.Column(scale=3): question_input = gr.Textbox(label="Any other questions about this book?", placeholder="Enter your question here...") with gr.Row(): submit_question = gr.Button("Submit") clear_question = gr.Button("Clear") with gr.Row(): with gr.Column(scale=2): genre_output = gr.Textbox(label="Genre") synopsis_output = gr.Textbox(label="Synopsis", lines=5) recommendations_output = gr.Textbox(label="Book Recommendations", lines=5) with gr.Column(scale=3): chat_output = gr.Textbox(label="Answer to your question", lines=10) with gr.Row(): retry_btn = gr.Button("Retry") clear_chat_btn = gr.Button("Clear") def clear_book_inputs(): return "", "" def clear_question_input(): return "" def clear_chat(): return "" def retry_last_question(title, author, question): return chat(title, author, question) submit_book.click(process_book, inputs=[title_input, author_input], outputs=[genre_output, synopsis_output, recommendations_output]) clear_book.click(clear_book_inputs, outputs=[title_input, author_input]) submit_question.click(chat, inputs=[title_input, author_input, question_input], outputs=[chat_output]) clear_question.click(clear_question_input, outputs=[question_input]) clear_chat_btn.click(clear_chat, outputs=[chat_output]) retry_btn.click(retry_last_question, inputs=[title_input, author_input, question_input], outputs=[chat_output]) demo.launch()