Spaces:
Sleeping
Sleeping
| 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() |