Update app.py

app.py

@@ -1,137 +1,227 @@
-import gradio as gr
+import os
 import logging
+import gradio as gr
 from langchain.chains import ConversationalRetrievalChain
 from langchain_openai import ChatOpenAI
-from langchain.memory import ConversationBufferMemory  # Using the updated memory package
-from langchain_community.vectorstores import Chroma  # Corrected import for Chroma
-from langchain_openai import OpenAIEmbeddings  # Updated import for OpenAIEmbeddings
+from langchain.memory import ConversationBufferMemory
+from langchain_community.vectorstores import FAISS
+from langchain_openai import OpenAIEmbeddings
 from langchain_community.document_loaders import WikipediaLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.tools import StructuredTool
 from langchain.callbacks.base import BaseCallbackHandler
 
-# ================================
-# Step 1: Setup Logging for Debugging
-# ================================
+# Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# ================================
-# Step 2: Load Wikipedia Data
-# ================================
-def fetch_wikipedia_content():
-    """Fetches Wikipedia content using LangChain's WikipediaLoader."""
-    loader = WikipediaLoader(query="Generative artificial intelligence", lang="en")
-    documents = loader.load()
-    return documents[0].page_content if documents else "Page not found."
-
-wiki_text = fetch_wikipedia_content()
-
-# ================================
-# Step 3: Process Wikipedia Text for Retrieval
-# ================================
-def process_and_store_wikipedia(text):
-    """Splits Wikipedia content into chunks, embeds them, and stores in ChromaDB."""
-    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-    chunks = splitter.split_text(text)
-    
-    embeddings = OpenAIEmbeddings()  # Using updated OpenAI embeddings
-    vectorstore = Chroma.from_texts(chunks, embedding=embeddings, persist_directory="/home/user/chroma_db")  # Ensuring persistence
-    return vectorstore.as_retriever()
-
-retriever = process_and_store_wikipedia(wiki_text)
-
-# ================================
-# Step 4: Initialize Chat Model and Memory
-# ================================
-llm = ChatOpenAI(model_name="gpt-4o")
-memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)  # Initialize memory for conversation history
-
-# ================================
-# Step 5: Create Q/A Retrieval Chain
-# ================================
-qa_chain = ConversationalRetrievalChain.from_llm(
-    llm, retriever=retriever, memory=memory
-)
-
-# ================================
-# Step 6: Implement Chatbot Response Function with Caching
-# ================================
-def ask_with_memory(query):
-    """Retrieves the answer from memory if available, otherwise fetches it using LangChain's Q/A chain."""
-
-    # Load chat history
-    chat_history = memory.load_memory_variables({})["chat_history"]
-
-    # Check if the exact query has been answered before
-    for i in range(len(chat_history) - 1):
-        if chat_history[i].content == query:
-            return chat_history[i + 1].content  # Return cached answer
-
-    # If not cached, process the query
-    response = qa_chain.invoke({"question": query})["answer"]
-
-    # Save query-response pair in memory
-    memory.save_context({"question": query}, {"answer": response})
+# Memory cache for storing answers
+class MemoryCache:
+    def __init__(self):
+        self.cache = {}
 
-    return response
+    def get(self, query: str):
+        if query in self.cache:
+            logger.info(f"Cache hit: {query}")
+            return self.cache.get(query)
+        return None
 
+    def set(self, query: str, response: str):
+        logger.info(f"Saving to cache: {query}")
+        self.cache[query] = response
 
-# ================================
-# Step 7: Implement Structured Function Calling for Section Extraction
-# ================================
-def extract_section_by_query(query: str) -> str:
-    """Finds and returns the most relevant section based on a user query using embeddings."""
-    vector_store = retriever  # Use the existing retriever
-
-    # Retrieve the most relevant section
-    retrieved_docs = vector_store.get_relevant_documents(query)
-
-    if not retrieved_docs:
-        return "Section not found."
-
-    return f"Section: {retrieved_docs[0].metadata.get('title', 'Unknown')}\n\n{retrieved_docs[0].page_content}"
-
-section_extraction_tool = StructuredTool.from_function(
-    extract_section_by_query,
-    name="extract_section_by_query",
-    description="Finds the most relevant Wikipedia section based on a user query using embeddings."
-)
-
-# ================================
-# Step 8: Implement Callback Logging for Debugging
-# ================================
+# Callback handler for logging
 class LoggingCallbackHandler(BaseCallbackHandler):
     def on_chain_start(self, serialized, inputs, **kwargs):
-        logger.info(f"Starting chain execution with input: {inputs}")
+        logger.info(f"Chain start. Inputs: {inputs}")
 
     def on_chain_end(self, outputs, **kwargs):
-        logger.info(f"Chain execution finished. Output: {outputs}")
-
-callback_handler = LoggingCallbackHandler()
-qa_chain.callbacks = [callback_handler]
-
-# ================================
-# Step 9: Define Gradio Interface
-# ================================
-def respond(message, history, system_message, max_tokens, temperature, top_p):
-    """
-    Processes user query and retrieves answers from Wikipedia-based Q/A system with caching.
-    """
-    return ask_with_memory(message)
-
-# ================================
-# Step 10: Create Gradio Interface
-# ================================
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are an AI expert answering questions about Generative AI.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
-    ],
-)
+        logger.info(f"Chain end. Outputs: {outputs}")
+
+    def on_retriever_start(self, *args, **kwargs):
+        logger.info("Retrieval start.")
+
+    def on_retriever_end(self, *args, **kwargs):
+        logger.info("Retrieval end.")
+
+    def on_llm_start(self, *args, **kwargs):
+        logger.info("LLM start.")
+
+    def on_llm_end(self, result, *args, **kwargs):
+        try:
+            final_text = result.generations[0][0].text
+            logger.info(f"LLM end. Text: {final_text}")
+        except Exception as e:
+            logger.error(f"LLM error: {e}")
+
+class GenAIQASystem:
+    def __init__(self):
+        self.cache = MemoryCache()
+        self.callback_handler = LoggingCallbackHandler()
+        self.content = None
+        self.qa_chain = None
+        self.memory = None
+        self.wiki_loaded = False
+        self.api_key_set = False
+
+    def set_api_key(self, api_key):
+        if not api_key:
+            return "Please provide a valid API key."
+            
+        try:
+            os.environ["OPENAI_API_KEY"] = api_key
+            # Test if API key works
+            embeddings = OpenAIEmbeddings()
+            embeddings.embed_query("Test")
+            self.api_key_set = True
+            return "API key set successfully!"
+        except Exception as e:
+            logger.error(f"API key error: {e}")
+            return f"Error setting API key: {str(e)}"
+    
+    def load_wikipedia(self):
+        if not self.api_key_set:
+            return "Please set your OpenAI API key first."
+            
+        if self.wiki_loaded:
+            return "Wikipedia content already loaded."
+            
+        try:
+            logger.info("Loading Wikipedia content for Generative artificial intelligence")
+            
+            # Load Wikipedia content
+            loader = WikipediaLoader(query="Generative artificial intelligence", lang="en")
+            documents = loader.load()
+            self.content = documents[0].page_content
+            
+            # Split content into chunks
+            text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+            chunks = text_splitter.split_text(self.content)
+            
+            # Create vector store
+            embeddings = OpenAIEmbeddings()
+            vectorstore = FAISS.from_texts(chunks, embeddings)
+            
+            # Initialize memory
+            self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+            
+            # Create QA Chain
+            llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
+            self.qa_chain = ConversationalRetrievalChain.from_llm(
+                llm,
+                retriever=vectorstore.as_retriever(),
+                memory=self.memory,
+                callbacks=[self.callback_handler]
+            )
+            
+            self.wiki_loaded = True
+            return "Wikipedia content loaded successfully!"
+        except Exception as e:
+            logger.error(f"Error loading Wikipedia: {e}")
+            return f"Error loading Wikipedia: {str(e)}"
+            
+    def extract_section(self, query: str):
+        """Extracts a specific section from the Wikipedia content."""
+        if not self.content:
+            return None
+            
+        query_lower = query.lower()
+        content_lower = self.content.lower()
+        
+        # Dictionary of section headers to look for
+        sections = {
+            "early history": "== early history ==",
+            "generative models": "== generative models ==",
+            "academic artificial intelligence": "== academic artificial intelligence =="
+        }
+        
+        # Check if query matches any section
+        for key, header in sections.items():
+            if key in query_lower:
+                start_index = content_lower.find(header)
+                if start_index != -1:
+                    logger.info(f"Found header: {header}")
+                    end_index = self.content.find("\n==", start_index + len(header))
+                    section_text = self.content[start_index:end_index].strip() if end_index != -1 else self.content[start_index:].strip()
+                    return section_text
+                    
+        return None
+    
+    def process_query(self, query):
+        if not self.api_key_set:
+            return "Please set your OpenAI API key in the Settings tab first."
+            
+        if not self.wiki_loaded:
+            return "Please load Wikipedia content in the Settings tab first."
+            
+        # Check cache first
+        cached_answer = self.cache.get(query)
+        if cached_answer:
+            return cached_answer
+            
+        # Try to extract a specific section
+        extracted_section = self.extract_section(query)
+        if extracted_section:
+            self.cache.set(query, extracted_section)
+            return f"[Section Found] {extracted_section}"
+            
+        # Use the QA chain
+        try:
+            logger.info(f"Processing query: {query}")
+            result = self.qa_chain.invoke({"question": query})
+            answer = result.get("answer", "No answer found")
+            self.cache.set(query, answer)
+            return answer
+        except Exception as e:
+            logger.error(f"Error in QA chain: {e}")
+            return f"Error processing query: {str(e)}"
+
+# Initialize system
+qa_system = GenAIQASystem()
+
+# Define Gradio interface
+with gr.Blocks(title="Generative AI Q/A System") as demo:
+    gr.Markdown("# Generative AI Q/A System")
+    gr.Markdown("Ask questions about Generative AI using this LangChain-based Q/A system")
+    
+    with gr.Tab("Chat"):
+        chatbot = gr.Chatbot()
+        msg = gr.Textbox(label="Your Question")
+        clear = gr.Button("Clear")
+        
+        def respond(message, history):
+            response = qa_system.process_query(message)
+            history.append((message, response))
+            return "", history
+        
+        msg.submit(respond, [msg, chatbot], [msg, chatbot])
+        clear.click(lambda: [], None, chatbot, queue=False)
+        
+    with gr.Tab("Settings"):
+        with gr.Group():
+            gr.Markdown("### Step 1: Set OpenAI API Key")
+            api_key_input = gr.Textbox(type="password", label="OpenAI API Key")
+            api_submit = gr.Button("Set API Key")
+            api_status = gr.Textbox(label="API Status", interactive=False)
+            
+        with gr.Group():
+            gr.Markdown("### Step 2: Load Wikipedia Content")
+            load_wiki_button = gr.Button("Load Wikipedia Content")
+            wiki_status = gr.Textbox(label="Loading Status", interactive=False)
+        
+        api_submit.click(qa_system.set_api_key, [api_key_input], [api_status])
+        load_wiki_button.click(qa_system.load_wikipedia, [], [wiki_status])
+        
+    gr.Markdown("## About")
+    gr.Markdown("""
+    This Q/A system uses LangChain and OpenAI to answer questions based on the Wikipedia page about Generative AI.
+    
+    Features:
+    - Caching mechanism to avoid repeating work
+    - Function calls to extract specific sections
+    - Logging to track processing
+    
+    Created by Anjali Haryani
+    """)
 
 if __name__ == "__main__":
     demo.launch()