Merge pull request #5 from neo4j-contrib/retrieval

update loader & bot

Author: tomasonjo

bot.py

@@ -5,12 +5,10 @@
 from langchain.embeddings.openai import OpenAIEmbeddings
 from langchain.embeddings import OllamaEmbeddings
 from langchain.chat_models import ChatOpenAI, ChatOllama
-from langchain.chains import ConversationalRetrievalChain
-from langchain.memory import ConversationBufferMemory
+from langchain.chains import RetrievalQAWithSourcesChain
 from langchain.prompts.chat import (
     ChatPromptTemplate,
     SystemMessagePromptTemplate,
-    AIMessagePromptTemplate,
     HumanMessagePromptTemplate,
 )
 from dotenv import load_dotenv
@@ -28,7 +26,7 @@
 # llm = ChatOllama(temperature=0, base_url=ollama_base_url)
 
 embeddings = OpenAIEmbeddings()
-llm = ChatOpenAI(temperature=0)
+llm = ChatOpenAI(temperature=0, model_name="gpt-4")
 
 # LLM only response
 template = "You are a helpful assistant that helps with programming questions."
@@ -57,10 +55,21 @@ def generate_llm_output(user_input: str) -> str:
     database="neo4j",  # neo4j by default
     index_name="stackoverflow",  # vector by default
     text_node_property="body",  # text by default
+    retrieval_query="""
+CALL  { with node
+    MATCH (node)<-[:ANSWERS]-(a)
+    WITH a
+    ORDER BY a.is_accepted DESC, a.score DESC
+    WITH collect(a.body)[..1] as answers
+    RETURN reduce(str='', text IN answers | str +  text + '\n') as answerTexts
+} 
+RETURN node.body + '\n' + answerTexts AS text, score, {source:node.link} AS metadata
+""",
 )
 
-memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-qa = ConversationalRetrievalChain.from_llm(llm, neo4j_db.as_retriever(), memory=memory)
+qa = RetrievalQAWithSourcesChain.from_chain_type(
+    llm, chain_type="stuff", retriever=neo4j_db.as_retriever(search_kwargs={"k": 2})
+)
 
 # Rag + KG
 kg = Neo4jVector.from_existing_index(
@@ -74,8 +83,9 @@ def generate_llm_output(user_input: str) -> str:
     retrieval_query="RETURN 'fancy' AS text, 1 AS score, {} AS metadata",  # Fix this
 )
 
-kg_memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-kg_qa = ConversationalRetrievalChain.from_llm(llm, kg.as_retriever(), memory=kg_memory)
+kg_qa = RetrievalQAWithSourcesChain.from_chain_type(
+    llm, chain_type="stuff", retriever=kg.as_retriever(search_kwargs={"k": 2})
+)
 
 # Streamlit stuff
 styl = f"""
@@ -91,6 +101,7 @@ def generate_llm_output(user_input: str) -> str:
 """
 st.markdown(styl, unsafe_allow_html=True)
 
+
 def chat_input():
     # Session state
     if "generated" not in st.session_state:
@@ -105,11 +116,16 @@ def chat_input():
     user_input = st.chat_input("What coding issue can I help you resolve today?")
 
     if user_input:
-        output = output_function(user_input)
+        try:
+            data = output_function(user_input)
+            output = data["answer"] + "\n" + data["sources"]
+        except KeyError:
+            output = output_function(user_input)
         st.session_state[f"user_input"].append(user_input)
         st.session_state[f"generated"].append(output)
         st.session_state[f"rag_mode"].append(name)
 
+
 def display_chat():
     if st.session_state[f"generated"]:
         size = len(st.session_state[f"generated"])
@@ -122,19 +138,19 @@ def display_chat():
                 st.caption(f"Mode: {st.session_state[f'rag_mode'][i]}")
                 st.write(st.session_state[f"generated"][i])
 
+
 def mode_select() -> str:
     options = ["LLM only", "Vector", "Vector + Graph"]
     return st.radio("Select sophistication mode", options, horizontal=True)
 
+
 name = mode_select()
-if(name == "LLM only"):
+if name == "LLM only":
     output_function = generate_llm_output
-elif(name == "Vector"):
-    output_function = qa.run
-elif(name == "Vector + Graph"):
-    output_function = kg_qa.run
+elif name == "Vector":
+    output_function = qa
+elif name == "Vector + Graph":
+    output_function = kg_qa
 
 chat_input()
 display_chat()
-
-

loader.py

@@ -2,7 +2,6 @@
 import requests
 
 from dotenv import load_dotenv
-from bs4 import BeautifulSoup
 from langchain.embeddings import OllamaEmbeddings, OpenAIEmbeddings
 from langchain.graphs import Neo4jGraph
 
@@ -43,6 +42,7 @@ def create_constraints():
 
 create_constraints()
 
+
 def create_vector_index(dimension):
     # TODO use Neo4jVector Code from LangChain on the existing graph
     index_query = "CALL db.index.vector.createNodeIndex('stackoverflow', 'Question', 'embedding', $dimension, 'cosine')"
@@ -56,27 +56,22 @@ def create_vector_index(dimension):
 
 
 def load_so_data(tag: str = "neo4j", page: int = 1) -> None:
-    base_url = "https://api.stackexchange.com/2.2/questions"
+    base_url = "https://api.stackexchange.com/2.3/search/advanced"
     parameters = (
-        f"?pagesize=100&page={page}&order=desc&sort=creation&tagged={tag}"
-        "&site=stackoverflow&filter=!6WPIomnMNcVD9"
+        f"?pagesize=100&page={page}&order=desc&sort=creation&answers=1&tagged={tag}"
+        "&site=stackoverflow&filter=!51dU0b1n(WTdqj5MH1iGsNShY6BhXXwJ)xwV5b"
     )
     data = requests.get(base_url + parameters).json()
     # Convert html to text and calculate embedding values
     for q in data["items"]:
-        question_text = BeautifulSoup(q["body"], features="html.parser").text
-        q["body"] = question_text
-        q["embedding"] = embeddings.embed_query(q["title"] + " " + question_text)
-        if q.get("answers"):
-            for a in q.get("answers"):
-                a["body"] = BeautifulSoup(a["body"], features="html.parser").text
+        q["embedding"] = embeddings.embed_query(q["title"] + " " + q["body_markdown"])
 
     import_query = """
     UNWIND $data AS q
     MERGE (question:Question {id:q.question_id}) 
     ON CREATE SET question.title = q.title, question.link = q.link,
-        question.favorite_count = q.favorite_count, question.creation_date = q.creation_date,
-        question.body = q.body, question.embedding = q.embedding
+        question.favorite_count = q.favorite_count, question.creation_date = datetime({epochSeconds: q.creation_date}),
+        question.body = q.body_markdown, question.embedding = q.embedding
     FOREACH (tagName IN q.tags | 
         MERGE (tag:Tag {name:tagName}) 
         MERGE (question)-[:TAGGED]->(tag)
@@ -85,8 +80,8 @@ def load_so_data(tag: str = "neo4j", page: int = 1) -> None:
         MERGE (question)<-[:ANSWERS]-(answer:Answer {id:a.answer_id})
         SET answer.is_accepted = a.is_accepted,
             answer.score = a.score,
-            answer.creation_date = a.creation_date,
-            answer.body = a.body
+            answer.creation_date = datetime({epochSeconds:a.creation_date}),
+            answer.body = a.body_markdown
         MERGE (answerer:User {id:coalesce(a.owner.user_id, "deleted")}) 
         ON CREATE SET answerer.display_name = a.owner.display_name,
                       answerer.reputation= a.owner.reputation
@@ -103,19 +98,24 @@ def load_so_data(tag: str = "neo4j", page: int = 1) -> None:
 
 # Streamlit
 def get_tag() -> str:
-    input_text = st.text_input("Which tag questions do you want to import?", value="neo4j")
+    input_text = st.text_input(
+        "Which tag questions do you want to import?", value="neo4j"
+    )
     return input_text
 
 
 def get_pages():
     col1, col2 = st.columns(2)
     with col1:
-        num_pages = st.number_input("Number of pages (100 questions per page)", step=1, min_value=1)
+        num_pages = st.number_input(
+            "Number of pages (100 questions per page)", step=1, min_value=1
+        )
     with col2:
         start_page = st.number_input("Start page", step=1, min_value=1)
     st.caption("Only questions with answers will be imported.")
     return (int(num_pages), int(start_page))
 
+
 st.header("StackOverflow Loader")
 st.subheader("Choose StackOverflow tags to load into Neo4j")
 st.caption("Go to http://localhost:7474/browser/ to explore the graph.")
@@ -127,7 +127,7 @@ def get_pages():
     with st.spinner("Loading... This might take a minute or two."):
         try:
             for page in range(1, num_pages + 1):
-                load_so_data(user_input, start_page + (page-1))
+                load_so_data(user_input, start_page + (page - 1))
             st.success("Import successful", icon="✅")
         except Exception as e:
             st.error(f"Error: {e}", icon="🚨")