Graph chat (#525)

* added graph mode

* added doc names parameter

* added the graph mode

* chat modes configurable and added the tabs for graph only mode

* changed the example env

* cypher cleaning

* fixed the default tab mode for graph mode

---------

Co-authored-by: vasanthasaikalluri <[email protected]>

Author: kartikpersistent

backend/score.py

@@ -286,8 +286,10 @@ async def chat_bot(uri=Form(None),model=Form(None),userName=Form(None), password
     logging.info(f"QA_RAG called at {datetime.now()}")
     qa_rag_start_time = time.time()
     try:
-        # database = "neo4j"
-        graph = create_graph_database_connection(uri, userName, password, database)
+        if mode == "graph":
+            graph = Neo4jGraph( url=uri,username=userName,password=password,database=database,sanitize = True, refresh_schema=True)
+        else:
+            graph = create_graph_database_connection(uri, userName, password, database)
         result = await asyncio.to_thread(QA_RAG,graph=graph,model=model,question=question,document_names=document_names,session_id=session_id,mode=mode)
 
         total_call_time = time.time() - qa_rag_start_time

backend/src/QA_integration_new.py

@@ -21,6 +21,7 @@
 from langchain_core.messages import HumanMessage,AIMessage
 from src.shared.constants import *
 from src.llm import get_llm
+from langchain.chains import GraphCypherQAChain
 import json
 
 ## Chat models
@@ -226,11 +227,10 @@ def setup_chat(model, graph, session_id, document_names,retrieval_query):
     logging.info(f"Model called in chat {model} and model version is {model_name}")
     retriever = get_neo4j_retriever(graph=graph,retrieval_query=retrieval_query,document_names=document_names)
     doc_retriever = create_document_retriever_chain(llm, retriever)
-    history = create_neo4j_chat_message_history(graph, session_id)
     chat_setup_time = time.time() - start_time
     logging.info(f"Chat setup completed in {chat_setup_time:.2f} seconds")
 
-    return llm, doc_retriever, history, model_name
+    return llm, doc_retriever, model_name
 
 def retrieve_documents(doc_retriever, messages):
     start_time = time.time()
@@ -264,25 +264,87 @@ def summarize_and_log(history, messages, llm):
     history_summarized_time = time.time() - start_time
     logging.info(f"Chat History summarized in {history_summarized_time:.2f} seconds")
 
+
+def create_graph_chain(model, graph):
+    try:
+        logging.info(f"Graph QA Chain using LLM model: {model}")
+
+        cypher_llm,model_name = get_llm(model)
+        qa_llm,model_name = get_llm(model)
+        graph_chain = GraphCypherQAChain.from_llm(
+            cypher_llm=cypher_llm,
+            qa_llm=qa_llm,
+            validate_cypher= True,
+            graph=graph,
+            # verbose=True, 
+            return_intermediate_steps = True,
+            top_k=3
+        )
+
+        logging.info("GraphCypherQAChain instance created successfully.")
+        return graph_chain,qa_llm,model_name
+
+    except Exception as e:
+        logging.error(f"An error occurred while creating the GraphCypherQAChain instance. : {e}") 
+
+
+def get_graph_response(graph_chain, question):
+    try:
+        cypher_res = graph_chain.invoke({"query": question})
+        
+        response = cypher_res.get("result")
+        cypher_query = ""
+        context = []
+
+        for step in cypher_res.get("intermediate_steps", []):
+            if "query" in step:
+                cypher_string = step["query"]
+                cypher_query = cypher_string.replace("cypher\n", "").replace("\n", " ").strip() 
+            elif "context" in step:
+                context = step["context"]
+        return {
+            "response": response,
+            "cypher_query": cypher_query,
+            "context": context
+        }
+    
+    except Exception as e:
+        logging.error("An error occurred while getting the graph response : {e}")
+
 def QA_RAG(graph, model, question, document_names,session_id, mode):
     try:
         logging.info(f"Chat Mode : {mode}")
-        if mode == "vector":
-            retrieval_query = VECTOR_SEARCH_QUERY
-        elif mode == "graph":
-            #WIP
-            result =  {
-                "session_id": session_id,
+        history = create_neo4j_chat_message_history(graph, session_id)
+        messages = history.messages
+        user_question = HumanMessage(content=question)
+        messages.append(user_question)
+
+        if mode == "graph":
+            graph_chain, qa_llm,model_version = create_graph_chain(model,graph)
+            graph_response = get_graph_response(graph_chain,question)
+            ai_response = AIMessage(content=graph_response["response"])
+            messages.append(ai_response)
+            summarize_and_log(history, messages, qa_llm)
+
+            result = {
+                "session_id": session_id, 
+                "message": graph_response["response"], 
+                "info": {
+                    "model": model_version,
+                    'cypher_query':graph_response["cypher_query"],
+                    "context" : graph_response["context"],
+                    "mode" : mode,
+                    "response_time": 0
+                },
                 "user": "chatbot"
-                }
+            } 
             return result
+        elif mode == "vector":
+            retrieval_query = VECTOR_SEARCH_QUERY
         else:
             retrieval_query = VECTOR_GRAPH_SEARCH_QUERY
 
-        llm, doc_retriever, history, model_version = setup_chat(model, graph, session_id, document_names,retrieval_query)
-        messages = history.messages
-        user_question = HumanMessage(content=question)
-        messages.append(user_question)
+        llm, doc_retriever, model_version = setup_chat(model, graph, session_id, document_names,retrieval_query)
 
         docs = retrieve_documents(doc_retriever, messages)
 

frontend/example.env

@@ -7,4 +7,5 @@ TIME_PER_CHUNK=4
 TIME_PER_PAGE=50
 CHUNK_SIZE=5242880
 LARGE_FILE_SIZE=5242880
-GOOGLE_CLIENT_ID=""
+GOOGLE_CLIENT_ID=""
+CHAT_MODES=""

frontend/src/components/ChatBot/ChatInfoModal.tsx

@@ -1,13 +1,23 @@
-import { Box, Typography, TextLink, Flex, Tabs, LoadingSpinner } from '@neo4j-ndl/react';
-import { DocumentTextIconOutline } from '@neo4j-ndl/react/icons';
+import {
+  Box,
+  Typography,
+  TextLink,
+  Flex,
+  Tabs,
+  LoadingSpinner,
+  CypherCodeBlock,
+  CypherCodeBlockProps,
+  useCopyToClipboard,
+} from '@neo4j-ndl/react';
+import { DocumentDuplicateIconOutline, DocumentTextIconOutline } from '@neo4j-ndl/react/icons';
 import '../../styling/info.css';
 import Neo4jRetrievalLogo from '../../assets/images/Neo4jRetrievalLogo.png';
 import wikipedialogo from '../../assets/images/wikipedia.svg';
 import youtubelogo from '../../assets/images/youtube.svg';
 import gcslogo from '../../assets/images/gcs.webp';
 import s3logo from '../../assets/images/s3logo.png';
 import { Chunk, Entity, GroupedEntity, UserCredentials, chatInfoMessage } from '../../types';
-import { useEffect, useMemo, useState } from 'react';
+import { useContext, useEffect, useMemo, useState } from 'react';
 import HoverableLink from '../UI/HoverableLink';
 import GraphViewButton from '../Graph/GraphViewButton';
 import { chunkEntitiesAPI } from '../../services/ChunkEntitiesInfo';
@@ -17,44 +27,88 @@ import { calcWordColor } from '@neo4j-devtools/word-color';
 import ReactMarkdown from 'react-markdown';
 import { GlobeAltIconOutline } from '@neo4j-ndl/react/icons';
 import { youtubeLinkValidation } from '../../utils/Utils';
+import { ThemeWrapperContext } from '../../context/ThemeWrapper';
+import { ClipboardDocumentCheckIconOutline } from '@neo4j-ndl/react/icons';
 
-const ChatInfoModal: React.FC<chatInfoMessage> = ({ sources, model, total_tokens, response_time, chunk_ids, mode }) => {
-  const [activeTab, setActiveTab] = useState<number>(3);
+const ChatInfoModal: React.FC<chatInfoMessage> = ({
+  sources,
+  model,
+  total_tokens,
+  response_time,
+  chunk_ids,
+  mode,
+  cypher_query,
+  graphonly_entities,
+}) => {
+  const [activeTab, setActiveTab] = useState<number>(mode === 'graph' ? 4 : 3);
   const [infoEntities, setInfoEntities] = useState<Entity[]>([]);
   const [loading, setLoading] = useState<boolean>(false);
   const { userCredentials } = useCredentials();
   const [nodes, setNodes] = useState<Node[]>([]);
   const [relationships, setRelationships] = useState<Relationship[]>([]);
   const [chunks, setChunks] = useState<Chunk[]>([]);
+  const themeUtils = useContext(ThemeWrapperContext);
+  const [, copy] = useCopyToClipboard();
+  const [copiedText, setcopiedText] = useState<boolean>(false);
+
   const parseEntity = (entity: Entity) => {
     const { labels, properties } = entity;
     const label = labels[0];
     const text = properties.id;
     return { label, text };
   };
+  const actions: CypherCodeBlockProps['actions'] = useMemo(
+    () => [
+      {
+        title: 'copy',
+        'aria-label': 'copy',
+        children: (
+          <>
+            {copiedText ? (
+              <ClipboardDocumentCheckIconOutline className='n-size-token-7' />
+            ) : (
+              <DocumentDuplicateIconOutline className='text-palette-neutral-text-icon' />
+            )}
+          </>
+        ),
+        onClick: () => {
+          void copy(cypher_query as string);
+          setcopiedText(true);
+        },
+      },
+    ],
+    [copiedText, cypher_query]
+  );
   useEffect(() => {
-    setLoading(true);
-    chunkEntitiesAPI(userCredentials as UserCredentials, chunk_ids.map((c) => c.id).join(','))
-      .then((response) => {
-        setInfoEntities(response.data.data.nodes);
-        setNodes(response.data.data.nodes);
-        setRelationships(response.data.data.relationships);
-        const chunks = response.data.data.chunk_data.map((chunk: any) => {
-          const chunkScore = chunk_ids.find((chunkdetail) => chunkdetail.id === chunk.id);
-          return {
-            ...chunk,
-            score: chunkScore?.score,
-          };
+    if (mode != 'graph') {
+      setLoading(true);
+      chunkEntitiesAPI(userCredentials as UserCredentials, chunk_ids.map((c) => c.id).join(','))
+        .then((response) => {
+          console.log({ response });
+          setInfoEntities(response.data.data.nodes);
+          setNodes(response.data.data.nodes);
+          setRelationships(response.data.data.relationships);
+          const chunks = response.data.data.chunk_data.map((chunk: any) => {
+            const chunkScore = chunk_ids.find((chunkdetail) => chunkdetail.id === chunk.id);
+            return {
+              ...chunk,
+              score: chunkScore?.score,
+            };
+          });
+          const sortedchunks = chunks.sort((a: any, b: any) => b.score - a.score);
+          setChunks(sortedchunks);
+          setLoading(false);
+        })
+        .catch((error) => {
+          console.error('Error fetching entities:', error);
+          setLoading(false);
         });
-        const sortedchunks = chunks.sort((a: any, b: any) => b.score - a.score);
-        setChunks(sortedchunks);
-        setLoading(false);
-      })
-      .catch((error) => {
-        console.error('Error fetching entities:', error);
-        setLoading(false);
-      });
-  }, [chunk_ids]);
+    }
+
+    () => {
+      setcopiedText(false);
+    };
+  }, [chunk_ids, mode]);
   const groupedEntities = useMemo<{ [key: string]: GroupedEntity }>(() => {
     return infoEntities.reduce((acc, entity) => {
       const { label, text } = parseEntity(entity);
@@ -107,9 +161,10 @@ const ChatInfoModal: React.FC<chatInfoMessage> = ({ sources, model, total_tokens
         </Box>
       </Box>
       <Tabs size='large' fill='underline' onChange={onChangeTabs} value={activeTab}>
-        <Tabs.Tab tabId={3}>Sources used</Tabs.Tab>
-        {mode === 'graph+vector' && <Tabs.Tab tabId={4}>Top Entities used</Tabs.Tab>}
-        <Tabs.Tab tabId={5}>Chunks</Tabs.Tab>
+        {mode != 'graph' && <Tabs.Tab tabId={3}>Sources used</Tabs.Tab>}
+        {(mode === 'graph+vector' || mode === 'graph') && <Tabs.Tab tabId={4}>Top Entities used</Tabs.Tab>}
+        {mode === 'graph' && cypher_query?.trim().length && <Tabs.Tab tabId={6}>Generated Cypher Query</Tabs.Tab>}
+        {mode != 'graph' && <Tabs.Tab tabId={5}>Chunks</Tabs.Tab>}
       </Tabs>
       <Flex className='p-4'>
         <Tabs.TabPanel className='n-flex n-flex-col n-gap-token-4 n-p-token-6' value={activeTab} tabId={3}>
@@ -225,28 +280,39 @@ const ChatInfoModal: React.FC<chatInfoMessage> = ({ sources, model, total_tokens
             <Box className='flex justify-center items-center'>
               <LoadingSpinner size='small' />
             </Box>
-          ) : Object.keys(groupedEntities).length > 0 ? (
+          ) : Object.keys(groupedEntities).length > 0 || Object.keys(graphonly_entities).length > 0 ? (
             <ul className='list-none p-4 max-h-80 overflow-auto'>
-              {sortedLabels.map((label, index) => (
-                <li
-                  key={index}
-                  className='flex items-center mb-2 text-ellipsis whitespace-nowrap max-w-[100%)] overflow-hidden'
-                >
-                  <div
-                    key={index}
-                    style={{ backgroundColor: `${groupedEntities[label].color}` }}
-                    className='legend mr-2'
-                  >
-                    {label} ({labelCounts[label]})
-                  </div>
-                  <Typography
-                    className='entity-text text-ellipsis whitespace-nowrap max-w-[calc(100%-120px)] overflow-hidden'
-                    variant='body-medium'
-                  >
-                    {Array.from(groupedEntities[label].texts).slice(0, 3).join(', ')}
-                  </Typography>
-                </li>
-              ))}
+              {mode == 'graph'
+                ? graphonly_entities.map((label, index) => (
+                    <li
+                      key={index}
+                      className='flex items-center mb-2 text-ellipsis whitespace-nowrap max-w-[100%)] overflow-hidden'
+                    >
+                      <div style={{ backgroundColor: calcWordColor(Object.keys(label)[0]) }} className='legend mr-2'>
+                        {Object.keys(label)[0]}
+                      </div>
+                    </li>
+                  ))
+                : sortedLabels.map((label, index) => (
+                    <li
+                      key={index}
+                      className='flex items-center mb-2 text-ellipsis whitespace-nowrap max-w-[100%)] overflow-hidden'
+                    >
+                      <div
+                        key={index}
+                        style={{ backgroundColor: `${groupedEntities[label].color}` }}
+                        className='legend mr-2'
+                      >
+                        {label} ({labelCounts[label]})
+                      </div>
+                      <Typography
+                        className='entity-text text-ellipsis whitespace-nowrap max-w-[calc(100%-120px)] overflow-hidden'
+                        variant='body-medium'
+                      >
+                        {Array.from(groupedEntities[label].texts).slice(0, 3).join(', ')}
+                      </Typography>
+                    </li>
+                  ))}
             </ul>
           ) : (
             <span className='h6 text-center'>No Entities Found</span>
@@ -340,6 +406,15 @@ const ChatInfoModal: React.FC<chatInfoMessage> = ({ sources, model, total_tokens
             <span className='h6 text-center'>No Chunks Found</span>
           )}
         </Tabs.TabPanel>
+        <Tabs.TabPanel value={activeTab} tabId={6}>
+          <CypherCodeBlock
+            code={cypher_query as string}
+            actions={actions}
+            headerTitle=''
+            theme={themeUtils.colorMode}
+            className='min-h-40'
+          />
+        </Tabs.TabPanel>
       </Flex>
       {activeTab == 4 && nodes.length && relationships.length ? (
         <Box className='button-container flex mt-2 justify-center'>