Update app.py

app.py

@@ -2,13 +2,13 @@ import gradio as gr
 from langchain_groq import ChatGroq
 from langchain_community.graphs.networkx_graph import NetworkxEntityGraph
 from langchain.chains import GraphQAChain
-from langchain_community.document_loaders import TextLoader  # Updated import
+from langchain_community.document_loaders import TextLoader
 from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.vectorstores import Pinecone  # Updated import
-from langchain_community.embeddings import HuggingFaceEmbeddings  # Correct import
+from langchain_community.vectorstores import Pinecone
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.schema.runnable import RunnablePassthrough
 from langchain.schema.output_parser import StrOutputParser
-from langchain_core.prompts import PromptTemplate  # Updated import
+from langchain import PromptTemplate
 from langchain_core.documents import Document
 from neo4j import GraphDatabase
 import networkx as nx
@@ -18,7 +18,6 @@ from datetime import datetime
 import gspread
 from oauth2client.service_account import ServiceAccountCredentials
 
-# Install the missing package
 os.system("pip install sentence-transformers")
 os.system("pip install gspread oauth2client")
 
@@ -35,27 +34,28 @@ def store_feedback_in_sheet(feedback, question, rag_response, graphrag_response)
     row = [timestamp, question, rag_response, graphrag_response, feedback]
     sheet.append_row(row)
 
-# Function to load data from Google Sheets
-def load_data():
-    data = sheet.get_all_records()
-    return data[-10:], len(data)  # return the last 10 rows and total count
-
-# Function to add review to Google Sheets
-def add_review(question, rag_response, graphrag_response, feedback):
-    store_feedback_in_sheet(feedback, question, rag_response, graphrag_response)
-    return None, None  # No output needed since we removed the data and count display
-
-# Initialize the chatbot and other necessary setups
+# RAG Setup
 text_path = r"./text_chunks.txt"
 loader = TextLoader(text_path, encoding='utf-8')
 documents = loader.load()
 text_splitter = CharacterTextSplitter(chunk_size=3000, chunk_overlap=4)
 docs = text_splitter.split_documents(documents)
 
+class CustomTextLoader(TextLoader):
+    def __init__(self, file_path: str, encoding: str = 'utf-8'):
+        super().__init__(file_path)
+        self.encoding = encoding
+
+    def load(self):
+        with open(self.file_path, encoding=self.encoding) as f:
+            text = f.read()
+        return [Document(page_content=text)]
+
 embeddings = HuggingFaceEmbeddings()
 
 from langchain.llms import HuggingFaceHub
 
+# Define the repo ID and connect to Mixtral model on Huggingface
 repo_id = "meta-llama/Meta-Llama-3-8B"
 llm = HuggingFaceHub(
   repo_id=repo_id, 
@@ -85,9 +85,9 @@ rag_llm = ChatGroq(
 )
 
 template = """
-You are a Thai rice assistants. These Human will ask you a questions about Thai Rice.
-Answer the question only in Thai languages.
-Use following piece of context to answer the question. 
+You are a Thai rice assistant. These humans will ask you questions about Thai rice.
+Answer the question only in Thai language.
+Use the following piece of context to answer the question. 
 If you don't know the answer, just say you don't know. 
 Keep the answer within 2 sentences and concise.
 Context: {context}
@@ -107,25 +107,118 @@ rag_chain = (
     | StrOutputParser()
 )
 
+class ChatBot():
+    loader = CustomTextLoader(r"./text_chunks.txt", encoding='utf-8')
+    documents = loader.load()
+
+    rag_chain = (
+        {"context": docsearch.as_retriever(),  "question": RunnablePassthrough()} 
+        | prompt 
+        | llm
+        | StrOutputParser() 
+    )
+
+graphrag_llm = ChatGroq(
+    model="Llama3-8b-8192",
+    temperature=0,
+    max_tokens=None,
+    timeout=None,
+    max_retries=5,
+    groq_api_key='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'
+)
+
+uri = "neo4j+s://46084f1a.databases.neo4j.io"
+user = "neo4j"
+password = "FwnX0ige_QYJk8eEYSXSF0l081mWWGIS7TFg6t8rLZc"
+driver = GraphDatabase.driver(uri, auth=(user, password))
+
+def fetch_nodes(tx):
+    query = "MATCH (n) RETURN id(n) AS id, labels(n) AS labels"
+    result = tx.run(query)
+    return result.data()
+
+def fetch_relationships(tx):
+    query = "MATCH (n)-[r]->(m) RETURN id(n) AS source, id(m) AS target, type(r) AS relation"
+    result = tx.run(query)
+    return result.data()
+
+def populate_networkx_graph():
+    G = nx.Graph()
+    with driver.session() as session:
+        nodes = session.read_transaction(fetch_nodes)
+        relationships = session.read_transaction(fetch_relationships)
+        for node in nodes:
+            G.add_node(node['id'], labels=node['labels'])
+        for relationship in relationships:
+            G.add_edge(
+                relationship['source'],
+                relationship['target'],
+                relation=relationship['relation']
+            )
+    return G
+
+networkx_graph = populate_networkx_graph()
+graph = NetworkxEntityGraph()
+graph._graph = networkx_graph
+
+graphrag_chain = GraphQAChain.from_llm(
+    llm=graphrag_llm,
+    graph=graph,
+    verbose=True
+)
+
 def get_rag_response(question):
     response = rag_chain.invoke(question)
     return response
 
-# Gradio Interface
+def get_graphrag_response(question):
+    system_prompt = "You are a Thai rice assistant that gives concise and direct answers. Do not explain the process, just provide the answer, provide the answer only in Thai."
+    formatted_question = f"System Prompt: {system_prompt}\n\nQuestion: {question}"
+    response = graphrag_chain.run(formatted_question)
+    return response
+
+def compare_models(question):
+    rag_response = get_rag_response(question)
+    graphrag_response = get_graphrag_response(question)
+    return rag_response, graphrag_response
+
+def handle_feedback(feedback, question, rag_response, graphrag_response):
+    try:
+        store_feedback_in_sheet(feedback, question, rag_response, graphrag_response)
+        return "ส่งสำเร็จ!"
+    except Exception as e:
+        return f"Error: {e}"
+
 with gr.Blocks() as demo:
+    gr.Markdown("## Thai Rice Assistant A/B Testing")
+
     with gr.Row():
         with gr.Column():
             question_input = gr.Textbox(label="ถามคำถามเกี่ยวกับข้าว:")
             submit_btn = gr.Button("ถาม")
+
+        with gr.Column():
             rag_output = gr.Textbox(label="Model A", interactive=False)
             graphrag_output = gr.Textbox(label="Model B", interactive=False)
 
+    with gr.Row():
         with gr.Column():
-            feedback_review = gr.Radio(label="How satisfied are you with the chatbot responses?", choices=["A ดีกว่า", "B ดีกว่า", "แย่ทั้งคู่", "เท่ากัน"])
-            feedback_comments = gr.Textbox(label="Comments", lines=10, placeholder="Any additional comments?")
-            feedback_submit = gr.Button(value="Submit Feedback")
+            choice = gr.Radio(["A ดีกว่า", "B ดีกว่า", "เท่ากัน", "แย่ทั้งคู่"], label="คำตอบไหนดีกว่ากัน?")
+            send_feedback_btn = gr.Button("ส่ง")
+
+    feedback_output = gr.Textbox(label="Feedback Status", interactive=False)
+
+    def on_submit(question):
+        rag_response, graphrag_response = compare_models(question)
+        return rag_response, graphrag_response
+
+    def on_feedback(feedback):
+        question = question_input.value
+        rag_response = rag_output.value
+        graphrag_response = graphrag_output.value
+        return handle_feedback(feedback, question, rag_response, graphrag_response)
 
-    submit_btn.click(lambda question: (get_rag_response(question), get_rag_response(question)), [question_input], [rag_output, graphrag_output])
-    feedback_submit.click(add_review, [question_input, rag_output, graphrag_output, feedback_review])
+    submit_btn.click(on_submit, inputs=[question_input], outputs=[rag_output, graphrag_output])
+    send_feedback_btn.click(on_feedback, inputs=[choice], outputs=[feedback_output])
 
 demo.launch(share=True)