Merge pull request #220 from neo4j-labs/configurable_chatbot

karanchellani · web-flow · commit 7c07230ccd76 · 2024-04-18T17:47:53.000+05:30
Configurable chatbot
diff --git a/backend/score.py b/backend/score.py
@@ -213,9 +213,11 @@ async def update_similarity_graph(uri=Form(None), userName=Form(None), password=
         return create_api_response(job_status,message=message,error=error_message)
         
 @app.post("/chat_bot")
-async def chat_bot(uri=Form(None), userName=Form(None), password=Form(None), question=Form(None), session_id=Form(None)):
+async def chat_bot(uri=Form(None),model=Form(None),userName=Form(None), password=Form(None), question=Form(None), session_id=Form(None)):
     try:
-        result = await asyncio.to_thread(QA_RAG,uri=uri,userName=userName,password=password,question=question,session_id=session_id)
+        # model=Form(None),
+        # model = "Gemini Pro"
+        result = await asyncio.to_thread(QA_RAG,uri=uri,model=model,userName=userName,password=password,question=question,session_id=session_id)
         return create_api_response('Success',data=result)
     except Exception as e:
         job_status = "Failed"
diff --git a/backend/src/QA_integration.py b/backend/src/QA_integration.py
@@ -7,12 +7,79 @@
 from langchain_openai import ChatOpenAI
 from langchain_openai import OpenAIEmbeddings
 from langchain_google_vertexai import VertexAIEmbeddings
+from langchain_google_vertexai import ChatVertexAI
+from langchain_google_vertexai import HarmBlockThreshold, HarmCategory
 import logging
 from langchain_community.chat_message_histories import Neo4jChatMessageHistory
+from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
+from src.shared.common_fn import load_embedding_model
+import re
+
 load_dotenv()
 
 openai_api_key = os.environ.get('OPENAI_API_KEY')
-model_version='gpt-4-0125-preview'
+
+
+# def get_embedding_function(embedding_model_name: str):
+#     if embedding_model_name == "openai":
+#         embedding_function = OpenAIEmbeddings()
+#         dimension = 1536
+#         logging.info(f"Embedding: Using OpenAI Embeddings , Dimension:{dimension}")
+#     elif embedding_model_name == "vertexai":        
+#         embedding_function = VertexAIEmbeddings(
+#             model_name="textembedding-gecko@003"
+#         )
+#         dimension = 768
+#         logging.info(f"Embedding: Using Vertex AI Embeddings , Dimension:{dimension}")
+#     else:
+#         embedding_function = SentenceTransformerEmbeddings(
+#             model_name="all-MiniLM-L6-v2"#, cache_folder="/embedding_model"
+#         )
+#         dimension = 384
+#         logging.info(f"Embedding: Using SentenceTransformer , Dimension:{dimension}")
+#     return embedding_function
+
+def get_llm(model : str):
+    if model == "OpenAI GPT 3.5":
+        model_version = "gpt-3.5-turbo-16k"
+        logging.info(f"Chat Model: GPT 3.5, Model Version : {model_version}")
+        llm = ChatOpenAI(model= model_version, temperature=0)
+    
+    elif model == "Gemini Pro" :
+        # model_version = "gemini-1.0-pro"
+        model_version = 'gemini-1.0-pro-001' 
+        logging.info(f"Chat Model: Gemini , Model Version : {model_version}")
+        llm = ChatVertexAI(model_name=model_version,
+                        #    max_output_tokens=100,
+                           convert_system_message_to_human=True,
+                           temperature=0,
+                           safety_settings={
+                               HarmCategory.HARM_CATEGORY_UNSPECIFIED: HarmBlockThreshold.BLOCK_NONE, 
+                               HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+                               HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE, 
+                               HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE, 
+                               HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+                               })
+    elif model == "Gemini 1.5 Pro" :
+        model_version = "gemini-1.5-pro-preview-0409"
+        logging.info(f"Chat Model: Gemini 1.5 , Model Version : {model_version}")
+        llm = ChatVertexAI(model_name=model_version,
+                        #    max_output_tokens=100,
+                           convert_system_message_to_human=True,
+                           temperature=0,
+                           safety_settings={
+                               HarmCategory.HARM_CATEGORY_UNSPECIFIED: HarmBlockThreshold.BLOCK_NONE, 
+                               HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+                               HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE, 
+                               HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE, 
+                               HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+                               })
+    else: 
+        ## for model == "OpenAI GPT 4" or model == "Diffbot" 
+        model_version = "gpt-4-0125-preview"
+        logging.info(f"Chat Model: GPT 4, Model Version : {model_version}")
+        llm = ChatOpenAI(model= model_version, temperature=0)
+    return llm
 
 def vector_embed_results(qa,question):
     vector_res={}
@@ -92,29 +159,52 @@ def get_chat_history(llm,uri,userName,password,session_id):
         error_message = str(e)
         logging.exception(f'Exception in retrieving chat history:{error_message}')
         # raise Exception(error_message)
-        return ''
+        return '' 
+
+def extract_and_remove_source(message):
+    pattern = r'\[Source: ([^\]]+)\]'
+    match = re.search(pattern, message)
+    if match:
+        sources_string = match.group(1)
+        sources = [source.strip().strip("'") for source in sources_string.split(',')]
+        new_message = re.sub(pattern, '', message).strip()
+        response = {
+            "message" : new_message,
+            "sources" : sources
+        }
+    else:
+        response = {
+            "message" : message,
+            "sources" : []
+        }
+    return response
 
-def QA_RAG(uri,userName,password,question,session_id):
+def QA_RAG(uri,model,userName,password,question,session_id):
     try:
         retrieval_query="""
         MATCH (node)-[:PART_OF]->(d:Document)
         WITH d, apoc.text.join(collect(node.text),"\n----\n") as text, avg(score) as score
         RETURN text, score, {source: COALESCE(CASE WHEN d.url CONTAINS "None" THEN d.fileName ELSE d.url END, d.fileName)} as metadata
         """
         embedding_model = os.getenv('EMBEDDING_MODEL')
+        embedding_function, _ = load_embedding_model(embedding_model)
         neo_db=Neo4jVector.from_existing_index(
-                embedding = VertexAIEmbeddings(model_name=embedding_model),
+                embedding = embedding_function,
                 url=uri,
                 username=userName,
                 password=password,
                 database="neo4j",
                 index_name="vector",
                 retrieval_query=retrieval_query,
             )
-        llm = ChatOpenAI(model= model_version, temperature=0)
+        # model = "Gemini Pro"
+        llm = get_llm(model = model)
 
         qa = RetrievalQA.from_chain_type(
-            llm=llm, chain_type="stuff", retriever=neo_db.as_retriever(search_kwargs={'k': 3,"score_threshold": 0.5}), return_source_documents=True
+            llm=llm, 
+            chain_type="stuff", 
+            retriever=neo_db.as_retriever(search_kwargs={'k': 3,"score_threshold": 0.5}),
+            return_source_documents=True
         )
 
         vector_res=vector_embed_results(qa,question)
@@ -133,32 +223,58 @@ def QA_RAG(uri,userName,password,question,session_id):
 
         chat_summary=get_chat_history(llm,uri,userName,password,session_id)
 
-        final_prompt = f"""You are a helpful question-answering agent. Your task is to analyze
-        and synthesize information from two sources: the top result from a similarity search
-        (unstructured information) and relevant data from a graph database (structured information). 
-        If structured information fails to find an answer then use the answer from unstructured information 
-        and vice versa. I only want a straightforward answer without mentioning from which source you got the answer. You are also receiving 
-        a chat history of the earlier conversation. You should be able to understand the context from the chat history and answer the question.
-        Given the user's query: {question}, provide a meaningful and efficient answer based
-        on the insights derived from the following data:
-        chat_summary:{chat_summary}
-        Structured information:  .
-        Unstructured information: {vector_res.get('result','')}.
 
+        # final_prompt = f"""You are a helpful question-answering agent. Your task is to analyze
+        # and synthesize information from two sources: the top result from a similarity search
+        # (unstructured information) and relevant data from a graph database (structured information). 
+        # If structured information fails to find an answer then use the answer from unstructured information 
+        # and vice versa. I only want a straightforward answer without mentioning from which source you got the answer. You are also receiving 
+        # a chat history of the earlier conversation. You should be able to understand the context from the chat history and answer the question.
+        # Given the user's query: {question}, provide a meaningful and efficient answer based
+        # on the insights derived from the following data:
+        # chat_summary:{chat_summary}
+        # Structured information:  .
+        # Unstructured information: {vector_res.get('result','')}.
+        # """ 
+
+        final_prompt = f"""
+            You are an AI-powered question-answering agent tasked with providing accurate and direct responses to user queries. Utilize information from the chat history, current user input, and relevant unstructured data effectively.
+
+            Response Requirements:
+            - Deliver concise and direct answers to the user's query without headers unless requested.
+            - Acknowledge and utilize relevant previous interactions based on the chat history summary.
+            - Respond to initial greetings appropriately, but avoid including a greeting in subsequent responses unless the chat is restarted or significantly paused.
+            - Clearly state if an answer is unknown; avoid speculating.
+
+            Instructions:
+            - Prioritize directly answering the User Input: {question}.
+            - Use the Chat History Summary: {chat_summary} to provide context-aware responses.
+            - Refer to Additional Unstructured Information: {vector_res.get('result', '')} only if it directly relates to the query.
+            - Cite sources clearly when using unstructured data in your response [Sources: {vector_res.get('source', '')}]. The Source must be printed only at the last in the format [Source: source1,source2]
+            Ensure that answers are straightforward and context-aware, focusing on being relevant and concise.
         """ 
 
         print(final_prompt)
+        llm = get_llm(model = model)
         response = llm.predict(final_prompt)
+        # print(response)
+
         ai_message=response
         user_message=question
         save_chat_history(uri,userName,password,session_id,user_message,ai_message)
 
-        res={"session_id":session_id,"message":response,"user":"chatbot"}
+        reponse = extract_and_remove_source(response)
+        message = reponse["message"]
+        sources = reponse["sources"]
+        # print(extract_and_remove_source(response))
+        print(response)
+        res={"session_id":session_id,"message":message,"sources":sources,"user":"chatbot"}
         return res
     except Exception as e:
       error_message = str(e)
       logging.exception(f'Exception in in QA component:{error_message}')
-    #   raise Exception(error_message)
-      return {"session_id":session_id,"message":"Something went wrong","user":"chatbot"}
+      message = "Something went wrong"
+      sources = []
+    #   raise Exception(error_message)  
+      return {"session_id":session_id,"message":message,"sources":sources,"user":"chatbot"}
 
- 
diff --git a/backend/src/shared/common_fn.py b/backend/src/shared/common_fn.py
@@ -1,5 +1,8 @@
 import logging
 from src.document_sources.youtube import create_youtube_url
+from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
+from langchain_google_vertexai import VertexAIEmbeddings
+from langchain_openai import OpenAIEmbeddings
 from langchain.docstore.document import Document
 import re
 import os
@@ -56,4 +59,24 @@ def get_chunk_and_graphDocument(graph_document_list, chunkId_chunkDoc_list):
             lst_chunk_chunkId_document.append({'graph_doc':graph_document,'chunk_id':chunk_id})
                   
   return lst_chunk_chunkId_document  
+
+
+def load_embedding_model(embedding_model_name: str):
+    if embedding_model_name == "openai":
+        embeddings = OpenAIEmbeddings()
+        dimension = 1536
+        logging.info(f"Embedding: Using OpenAI Embeddings , Dimension:{dimension}")
+    elif embedding_model_name == "vertexai":        
+        embeddings = VertexAIEmbeddings(
+            model="textembedding-gecko@003"
+        )
+        dimension = 768
+        logging.info(f"Embedding: Using Vertex AI Embeddings , Dimension:{dimension}")
+    else:
+        embeddings = SentenceTransformerEmbeddings(
+            model_name="all-MiniLM-L6-v2"#, cache_folder="/embedding_model"
+        )
+        dimension = 384
+        logging.info(f"Embedding: Using SentenceTransformer , Dimension:{dimension}")
+    return embeddings, dimension
                  
diff --git a/frontend/src/assets/ChatbotMessages.json b/frontend/src/assets/ChatbotMessages.json
@@ -10,7 +10,8 @@
       "id": 2,
       "message": " Welcome to the Neo4j Knowledge Graph Chat. You can ask questions related to documents which have been completely processed.",
       "user": "chatbot",
-      "datetime": "01/01/2024 00:00:00"
+      "datetime": "01/01/2024 00:00:00",
+      "sources":["https://neo4j.com/"]
     },
     {
       "id": 3,
diff --git a/frontend/src/components/Chatbot.tsx b/frontend/src/components/Chatbot.tsx
@@ -1,19 +1,21 @@
 /* eslint-disable no-confusing-arrow */
 import { useEffect, useRef, useState } from 'react';
-import { Button, Widget, Typography, Avatar, TextInput } from '@neo4j-ndl/react';
+import { Button, Widget, Typography, Avatar, TextInput, TextLink } from '@neo4j-ndl/react';
 import ChatBotUserAvatar from '../assets/images/chatbot-user.png';
 import ChatBotAvatar from '../assets/images/chatbot-ai.png';
 import { ChatbotProps, UserCredentials } from '../types';
 import { useCredentials } from '../context/UserCredentials';
 import chatBotAPI from '../services/QnaAPI';
 import { v4 as uuidv4 } from 'uuid';
+import { useFileContext } from '../context/UsersFiles';
 
 export default function Chatbot(props: ChatbotProps) {
   const { messages: listMessages, setMessages: setListMessages } = props;
   const [inputMessage, setInputMessage] = useState('');
   const formattedTextStyle = { color: 'rgb(var(--theme-palette-discovery-bg-strong))' };
   const [loading, setLoading] = useState<boolean>(false);
   const { userCredentials } = useCredentials();
+  const { model } = useFileContext();
   const messagesEndRef = useRef<HTMLDivElement>(null);
   const [sessionId, setSessionId] = useState<string>(sessionStorage.getItem('session_id') ?? '');
 
@@ -29,17 +31,24 @@ export default function Chatbot(props: ChatbotProps) {
     }
   }, []);
 
-  const simulateTypingEffect = (responseText: string, index = 0) => {
-    if (index < responseText.length) {
+  const simulateTypingEffect = (response: { reply: string; sources?: [string] }, index = 0) => {
+    if (index < response.reply.length) {
       const nextIndex = index + 1;
-      const currentTypedText = responseText.substring(0, nextIndex);
+      const currentTypedText = response.reply.substring(0, nextIndex);
       if (index === 0) {
         const date = new Date();
         const datetime = `${date.toLocaleDateString()} ${date.toLocaleTimeString()}`;
-        if (responseText.length <= 1) {
+        if (response.reply.length <= 1) {
           setListMessages((msgs) => [
             ...msgs,
-            { id: Date.now(), user: 'chatbot', message: currentTypedText, datetime: datetime, isTyping: true },
+            {
+              id: Date.now(),
+              user: 'chatbot',
+              message: currentTypedText,
+              datetime: datetime,
+              isTyping: true,
+              sources: response?.sources,
+            },
           ]);
         } else {
           setListMessages((msgs) => {
@@ -49,6 +58,7 @@ export default function Chatbot(props: ChatbotProps) {
             lastmsg.message = currentTypedText;
             lastmsg.datetime = datetime;
             lastmsg.isTyping = true;
+            lastmsg.sources = response?.sources;
             return msgs.map((msg, index) => {
               if (index === msgs.length - 1) {
                 return lastmsg;
@@ -60,7 +70,7 @@ export default function Chatbot(props: ChatbotProps) {
       } else {
         setListMessages((msgs) => msgs.map((msg) => (msg.isTyping ? { ...msg, message: currentTypedText } : msg)));
       }
-      setTimeout(() => simulateTypingEffect(responseText, nextIndex), 20);
+      setTimeout(() => simulateTypingEffect(response, nextIndex), 20);
     } else {
       setListMessages((msgs) => msgs.map((msg) => (msg.isTyping ? { ...msg, isTyping: false } : msg)));
     }
@@ -79,15 +89,15 @@ export default function Chatbot(props: ChatbotProps) {
     try {
       setLoading(true);
       setInputMessage('');
-      simulateTypingEffect(' ');
-      const chatresponse = await chatBotAPI(userCredentials as UserCredentials, inputMessage, sessionId);
-      chatbotReply = chatresponse?.data?.message;
-      simulateTypingEffect(chatbotReply);
+      simulateTypingEffect({ reply: ' ' });
+      const chatresponse = await chatBotAPI(userCredentials as UserCredentials, inputMessage, sessionId, model);
+      chatbotReply = chatresponse?.data?.data?.message;
+      simulateTypingEffect({ reply: chatbotReply, sources: chatresponse?.data?.data?.sources });
       setLoading(false);
     } catch (error) {
       chatbotReply = "Oops! It seems we couldn't retrieve the answer. Please try again later";
       setInputMessage('');
-      simulateTypingEffect(chatbotReply);
+      simulateTypingEffect({ reply: chatbotReply });
       setLoading(false);
     }
   };
@@ -160,6 +170,21 @@ export default function Chatbot(props: ChatbotProps) {
                   </div>
                   <div className='text-right align-bottom pt-3'>
                     <Typography variant='body-small'>{chat.datetime}</Typography>
+                    {chat?.sources?.length ? (
+                      <div className={`flex ${chat.sources?.length > 1 ? 'flex-col' : 'flex-row justify-end'} gap-1`}>
+                        {chat.sources.map((link) => (
+                          <div className='text-right'>
+                            {link.startsWith('http') || link.startsWith('https') ? (
+                              <TextLink href={link} externalLink={true}>
+                                Source
+                              </TextLink>
+                            ) : (
+                              <Typography variant='body-small'>{link}</Typography>
+                            )}
+                          </div>
+                        ))}
+                      </div>
+                    ) : null}
                   </div>
                 </Widget>
               </div>
diff --git a/frontend/src/services/QnaAPI.ts b/frontend/src/services/QnaAPI.ts
diff --git a/frontend/src/types.ts b/frontend/src/types.ts
