Update app.py

app.py

@@ -3,11 +3,23 @@ warnings.simplefilter("ignore", category=FutureWarning)
 
 import os
 import streamlit as st
+from neo4j import GraphDatabase
 from huggingface_hub import InferenceClient
 from langchain_community.vectorstores import Neo4jVector
 from transformers import AutoTokenizer, AutoModel
 import torch
 
+# Hugging Face API Setup
+API_TOKEN = os.environ.get("HUGGINGFACE_API_TOKEN")
+MISTRAL_MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
+client = InferenceClient(api_key=API_TOKEN, )
+
+# Driver neo4j 
+driver = GraphDatabase.driver(
+        os.environ['NEO4J_URI'], 
+        auth=(os.environ['NEO4J_USERNAME'], os.environ['NEO4J_PASSWORD'])
+    )
+
 # Custom Embedding Class
 class CustomHuggingFaceEmbeddings:
     def __init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2"):
@@ -44,11 +56,6 @@ def setup_vector_index():
         embedding_node_property='embedding',
     )
 
-# Hugging Face API Setup
-API_TOKEN = os.environ.get("HUGGINGFACE_API_TOKEN")
-MISTRAL_MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
-client = InferenceClient(api_key=API_TOKEN, )
-
 # Query Mistral
 def query_from_mistral(context: str, user_input: str):
     messages = [
@@ -62,6 +69,17 @@ def query_from_mistral(context: str, user_input: str):
     )
     return completion.choices[0].message["content"]
 
+# Find keywords
+def query_article_keywords(name):
+    with driver.session() as session:
+        query = """
+        MATCH (a:Article)-[:CONTAIN]->(k:Keyword)
+        WHERE a.name = $name
+        RETURN k
+        """
+        result = session.run(query, name=name)
+        return [record["k"] for record in result]
+
 # extract data from retriever response
 def extract_data(documents):
     result = []
@@ -79,10 +97,14 @@ def extract_data(documents):
             elif line.lower().startswith("abstract:"):
                 abstract = line[len("abstract:"):].strip()
 
+        keywords = query_article_keywords(title)
+        keywords = [dict(node)['text'] for node in keywords]
+
         doc_data = {
             "Publication Date": publication_date,
             "Title": title,
             "Abstract": abstract,
+            "keywords": ','.join(keywords)
         }
         result.append(doc_data)
 
@@ -91,8 +113,25 @@ def extract_data(documents):
 # Main Streamlit Application
 def main():
     st.set_page_config(page_title="Vector Chat with Mistral", layout="centered")
-    st.title("🤖 Vector Chat with Mistral")
-    st.markdown("Chat with **Mistral-7B-Instruct** using context retrieved from a Neo4j vector index.")
+    
+    # App description and features
+    st.title("🤖 RAG with Mistral")
+    st.markdown("""
+        ## Description:
+        Chat with **Mistral-7B-Instruct** using context retrieved from a **Neo4j** vector index. This app allows you to ask questions, and the assistant will provide real-time, context-driven answers by querying relevant articles and their keywords from the database.
+    """)
+
+    st.image(image="image.jpg", caption="Neo4j")
+
+    st.markdown("""
+        ## Key Features:
+        - **Real-time context search** from a Neo4j vector index.
+        - **Integration with Mistral-7B-Instruct model** for natural language processing.
+        - **Keyword extraction** from relevant articles for enhanced context-based responses.
+
+        ## GitHub Repository:
+        You can find the source code and more information about this app on GitHub: [GitHub Repository Link](https://github.com/yourusername/your-repository-name)
+    """)
 
     # Initialize the vector index
     vector_index = setup_vector_index()
@@ -115,8 +154,17 @@ def main():
                     st.warning("No relevant context found. Please refine your query.")
                     response = "I'm sorry, I couldn't find any relevant information to answer your question."
                 else:
-                    context = "\n".join([f"Title: {doc['Title']}\nAbstract: {doc['Abstract']}\nPublication Date: {doc['Publication Date']}" 
-                                     for doc in extract_data(context_results)])
+                    data_dict = extract_data(context_results)
+
+                    # convert to string
+                    context = '\n'.join([ 
+                        f"Title: {doc['Title']}\n"
+                        f"Abstract: {doc['Abstract']}\n"
+                        f"Publication Date: {doc['Publication Date']}\n"
+                        f"Keywords: {doc['keywords']}"
+                        for doc in data_dict
+                    ])
+
                     response = query_from_mistral(context.strip(), user_input)
 
                 st.session_state.messages.append({"role": "bot", "content": response})