Update app.py

app.py

@@ -1,227 +1,137 @@
-import os
-import logging
 import gradio as gr
+import logging
 from langchain.chains import ConversationalRetrievalChain
 from langchain_openai import ChatOpenAI
-from langchain.memory import ConversationBufferMemory
-from langchain_community.vectorstores import FAISS
-from langchain_openai import OpenAIEmbeddings
+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_community.document_loaders import WikipediaLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.tools import StructuredTool
 from langchain.callbacks.base import BaseCallbackHandler
 
-# Setup logging
+# ================================
+# Step 1: Setup Logging for Debugging
+# ================================
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Memory cache for storing answers
-class MemoryCache:
-    def __init__(self):
-        self.cache = {}
+# ================================
+# 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})
 
-    def get(self, query: str):
-        if query in self.cache:
-            logger.info(f"Cache hit: {query}")
-            return self.cache.get(query)
-        return None
+    return response
 
-    def set(self, query: str, response: str):
-        logger.info(f"Saving to cache: {query}")
-        self.cache[query] = response
 
-# Callback handler for logging
+# ================================
+# 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
+# ================================
 class LoggingCallbackHandler(BaseCallbackHandler):
     def on_chain_start(self, serialized, inputs, **kwargs):
-        logger.info(f"Chain start. Inputs: {inputs}")
+        logger.info(f"Starting chain execution with input: {inputs}")
 
     def on_chain_end(self, outputs, **kwargs):
-        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
-    """)
+        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)"),
+    ],
+)
 
 if __name__ == "__main__":
     demo.launch()