made changes to graph rag file and some ui elements

frontend/src/Components/AiComponents/ChatWindow.css

@@ -127,6 +127,7 @@
 .markdown {
   margin: -1rem 0 -0.8rem 0;
   line-height: 2rem;
+  white-space: pre-wrap;
 }
 
 /* Post-answer icons container: placed below the bot bubble */
@@ -275,4 +276,4 @@
   100% {
     transform: rotate(360deg);
   }
-}
+}

frontend/src/Components/AiComponents/ChatWindow.js

@@ -168,10 +168,7 @@ function ChatWindow({
               )}
               {renderSourcesRead() !== null && (
                 <div className="sources-read-container">
-                  <p
-                    className="sources-read"
-                    onClick={() => openRightSidebar("sources", blockId)}
-                  >
+                  <p className="sources-read">
                     Sources Read: {renderSourcesRead()}
                   </p>
                 </div>

frontend/src/Components/AiComponents/Graph.css

@@ -21,7 +21,7 @@
   width: 45% !important;
   max-width: 100% !important;
   background-color: #1e1e1e !important;
-  max-height: 80vh !important;
+  min-height: 80vh !important;
   overflow: hidden !important; /* Prevent scrolling */
 }
 
@@ -62,7 +62,7 @@
 .graph-dialog-content {
   padding: 0 !important;
   background-color: #1e1e1e !important;
-  height: 750px !important;
+  height: 550px !important;
   overflow: hidden !important;
 }
 

frontend/src/Components/AiComponents/Sources.css

@@ -6,6 +6,10 @@
   padding: 0 !important;
 }
 
+.loading-sources {
+  padding: 1rem !important;
+}
+
 /* Styling for each Card component */
 .source-card {
   background-color: #3e3e3eec !important;

frontend/src/Components/AiComponents/Sources.js

@@ -65,7 +65,7 @@ export default function Sources({ sources, handleSourceClick }) {
   if (loading) {
     return (
       <Box className="sources-container">
-        <Typography variant="body2">Loading Sources...</Typography>
+        <Typography className="loading-sources" variant="body2">Loading Sources...</Typography>
       </Box>
     );
   }

main.py

@@ -445,6 +445,10 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
 
                     contents = data["contents"]
                     current_search_contents.extend(contents)
+                
+                elif event_type == "search_process_completed":
+                    await sse_queue.put(("step", "Processing final graph tasks..."))
+                    await asyncio.sleep(0.01)   # Sleep for a short time to allow the message to be sent
 
             state['graph_rag'].set_on_event_callback(on_event_callback)
 

requirements.txt

@@ -18,7 +18,8 @@ langchain_xai==0.1.1
 langgraph==0.2.62
 model2vec==0.3.3
 neo4j==5.26.0
-rustworkx==0.16.0
+rustworkx===0.16.0
+rank_bm25==0.2.2
 openai==1.59.3
 protobuf==4.23.4
 PyPDF2==3.0.1
@@ -32,4 +33,4 @@ transformers==4.46.2
 xformers==0.0.29.post1
 # Intall the following seperately:-
 # latest torch cuda version compatible with xformers' version
-# python -m spacy download en_core_web_sm
+# python -m spacy download en_core_web_sm

src/rag/graph_rag.py

@@ -17,6 +17,9 @@ from src.search.search_engine import SearchEngine
 from src.crawl.crawler import CustomCrawler #, Crawler
 from sentence_transformers import SentenceTransformer
 from bert_score.scorer import BERTScorer
+from openai import RateLimitError
+from anthropic import RateLimitError as AnthropicRateLimitError
+from google.api_core.exceptions import ResourceExhausted
 
 class GraphRAG:
     def __init__(self, num_workers: int = 1):
@@ -54,6 +57,9 @@ class GraphRAG:
         self.sub_node_counter = 0
         self.cross_connections = set()
 
+        # Semaphore protection
+        self.semaphore = asyncio.Semaphore(min(num_workers * 2, 12))
+
         # Thread pool
         self.executor = ThreadPoolExecutor(max_workers=self.num_workers)
 
@@ -379,10 +385,13 @@ class GraphRAG:
         )
 
     async def build_graph(self, query: str, data: str = None, parent_node_id: str = None, 
-                          depth: int = 0, threshold: float = 0.8, recurse: bool = True, 
-                          context: list = None, session_id: str = None, max_tokens_allowed: int = 128000):
+                    depth: int = 0, threshold: float = 0.8, recurse: bool = True, 
+                    context: list = None, session_id: str = None, max_tokens_allowed: int = 128000,
+                    node_data_futures: dict = None, sub_nodes_info: list = None, 
+                    sub_query_ids: list = None, pre_req_nodes: list = None):
         """Build a new graph structure in memory."""
-        async def process_node(node_id: str, sub_query: str, session_id: str, future: asyncio.Future, depth=depth, max_tokens_allowed=max_tokens_allowed):
+        async def process_node(node_id: str, sub_query: str, session_id: str, 
+                               future: asyncio.Future, max_tokens_allowed: int = max_tokens_allowed):
             try:
                 optimized_query = await self.search_engine.generate_optimized_query(sub_query)
                 results = await self.search_engine.search(
@@ -427,14 +436,7 @@ class GraphRAG:
                     elif content:
                         contents += f"Document {k}:\n{content}\n\n"
 
-                if contents.strip():
-                    if depth == 0:
-                        await self.emit_event("sub_query_processed", {
-                            "node_id": node_id,
-                            "sub_query": sub_query,
-                            "contents": contents
-                        })
-
+                if contents.strip():                    
                     token_count = self.llm.get_num_tokens(contents)
                     if token_count > max_tokens_allowed:
                         contents = await self.chunking.chunker(
@@ -444,13 +446,6 @@ class GraphRAG:
                         )
                         print(f"Number of tokens in the answer: {token_count}")
                         print(f"Number of tokens in the content: {self.llm.get_num_tokens(contents)}")
-                else:
-                    if depth == 0:
-                        await self.emit_event("sub_query_failed", {
-                            "node_id": node_id,
-                            "sub_query": sub_query,
-                            "contents": contents
-                        })
 
                 graph_data = self._get_current_graph_data()
                 graph = graph_data["graph"]
@@ -460,13 +455,15 @@ class GraphRAG:
                     idx = node_map[node_id]
                     node_data = graph.get_node_data(idx)
                     node_data["data"] = contents
-                future.set_result(contents)
+                if not future.done():
+                    future.set_result(contents)
             except Exception as e:
                 print(f"Error processing node {node_id}: {str(e)}")
-                future.set_exception(e)
-                raise
+                if not future.done():
+                    future.set_exception(e)
+                raise e
 
-        async def process_dependent_node(node_id: str, sub_query: str, depth, dep_futures: list, future):
+        async def process_dependent_node(node_id: str, sub_query: str, dep_futures: list, future):
             try:
                 dep_data = [await f for f in dep_futures]
                 modified_query = await self.query_processor.modify_query(
@@ -490,16 +487,16 @@ class GraphRAG:
                     node_data["embedding"] = embedding.tolist() if hasattr(embedding, "tolist") else embedding
                 try:
                     if not future.done():
-                        await process_node(node_id, modified_query, session_id, future, depth, max_tokens_allowed)
+                        await process_node(node_id, modified_query, session_id, future, max_tokens_allowed)
                 except Exception as e:
                     if not future.done():
                         future.set_exception(e)
-                    raise
+                    raise e
             except Exception as e:
                 print(f"Error processing dependent node {node_id}: {str(e)}")
                 if not future.done():
                     future.set_exception(e)
-                raise
+                raise e
 
         def create_cross_connections():
             try:
@@ -548,7 +545,15 @@ class GraphRAG:
         
         if context is None:
             context = []
-        node_data_futures = {}
+
+        if node_data_futures is None:
+            node_data_futures = {}
+        if sub_nodes_info is None:
+            sub_nodes_info = []
+        if sub_query_ids is None:
+            sub_query_ids = []
+        if pre_req_nodes is None:
+            pre_req_nodes = {}
 
         if parent_node_id is None:
             self.add_node(self.root_node_id, query, data)
@@ -565,9 +570,6 @@ class GraphRAG:
             context.append(response_data)
 
         if len(sub_queries) > 1 and sub_queries[0] != query:
-            sub_query_ids = []
-            pre_req_nodes = {}
-
             for idx, (sub_query, role, dependency) in enumerate(zip(sub_queries, roles, dependencies)):
                 if depth == 0:
                     await self.emit_event("sub_query_created", {
@@ -589,13 +591,13 @@ class GraphRAG:
                 self.add_node(sub_node_id, sub_query, role=role)
                 future = asyncio.Future()
                 node_data_futures[sub_node_id] = future
+                sub_nodes_info.append((sub_node_id, sub_query, role, dependency, future, depth))
 
                 if role.lower() in ['pre-requisite', 'prerequisite']:
                     pre_req_nodes[idx] = sub_node_id
 
                 if role.lower() in ('pre-requisite', 'prerequisite', 'independent'):
                     self.add_edge(parent_node_id, sub_node_id, relationship_type='hierarchical')
-
                 elif role.lower() == 'dependent':
                     if isinstance(dependency, list) and (len(dependency) == 2 and all(isinstance(d, list) for d in dependency)):
                         print(f"Dependency: {dependency}")
@@ -618,12 +620,11 @@ class GraphRAG:
 
                         if current_deps not in [None, []]:
                             for dep_idx in current_deps:
-                                if dep_idx < len(sub_queries):
+                                if dep_idx < len(sub_query_ids):
                                     dep_node_id = sub_query_ids[dep_idx]
                                     self.add_edge(dep_node_id, sub_node_id, relationship_type='logical')
                                 else:
                                     raise ValueError(f"Invalid dependency index: {dep_idx}")
-                                
                     elif len(dependency) > 0:
                         for dep_idx in dependency:
                             if dep_idx < len(sub_queries):
@@ -635,66 +636,6 @@ class GraphRAG:
                         raise ValueError(f"Invalid dependency: {dependency}")
                 else:
                     raise ValueError(f"Unexpected role: {role}")
-                
-            tasks = []
-            for idx in range(len(sub_queries)):
-                node_id = sub_query_ids[idx]
-                future = node_data_futures[node_id]
-
-                if roles[idx].lower() in ('pre-requisite', 'prerequisite', 'independent'):
-                    tasks.append(process_node(node_id, sub_queries[idx], session_id, future, depth, max_tokens_allowed))
-
-            for idx in range(len(sub_queries)):
-                node_id = sub_query_ids[idx]
-                future = node_data_futures[node_id]
-
-                if roles[idx].lower() == 'dependent':
-                    dep_futures = []
-
-                    if isinstance(dependencies[idx], list) and len(dependencies[idx]) == 2:
-                        prev_deps, current_deps = dependencies[idx]
-                        if context and prev_deps not in [None, []]:
-                            for context_idx, context_data in enumerate(context):
-                                if isinstance(prev_deps, list) and context_idx < len(prev_deps):
-                                    context_dep = prev_deps[context_idx]
-                                    if context_dep is not None and isinstance(context_data, dict) and 'subqueries' in context_data:
-                                        if context_dep < len(context_data['subqueries']):
-                                            dep_query = context_data['subqueries'][context_dep]['subquery']
-                                            matching_nodes = self.find_nodes_by_properties(query=dep_query)
-                                            if matching_nodes not in [None, []]:
-                                                dep_node_id = matching_nodes[0].get('node_id', None)
-                                                score = float(matching_nodes[0].get('score', 0))
-                                                if score == 1.0 and dep_node_id in node_data_futures:
-                                                    dep_futures.append(node_data_futures[dep_node_id])
-
-                                elif isinstance(prev_deps, int):
-                                    if prev_deps < len(context_data['subqueries']):
-                                        dep_query = context_data['subqueries'][prev_deps]['subquery']
-                                        matching_nodes = self.find_nodes_by_properties(query=dep_query)
-                                        if matching_nodes not in [None, []]:
-                                            dep_node_id = matching_nodes[0].get('node_id', None)
-                                            score = matching_nodes[0].get('score', 0)
-                                            if score == 1.0 and dep_node_id in node_data_futures:
-                                                dep_futures.append(node_data_futures[dep_node_id])
-
-                        if current_deps not in [None, []]:
-                            current_deps_list = [current_deps] if isinstance(current_deps, int) else current_deps
-                            for dep_idx in current_deps_list:
-                                if dep_idx < len(sub_queries):
-                                    dep_node_id = sub_query_ids[dep_idx]
-                                    if dep_node_id in node_data_futures:
-                                        dep_futures.append(node_data_futures[dep_node_id])
-
-                    tasks.append(process_dependent_node(node_id, sub_queries[idx], depth, dep_futures, future))
-
-            if depth == 0:
-                await self.emit_event("search_process_started", {
-                    "depth": depth,
-                    "sub_queries": sub_queries,
-                    "roles": roles
-                })
-
-            await asyncio.gather(*tasks)
 
             if recurse:
                 recursion_tasks = []
@@ -710,7 +651,11 @@ class GraphRAG:
                                 threshold=threshold,
                                 recurse=recurse,
                                 context=context,
-                                session_id=session_id
+                                session_id=session_id,
+                                node_data_futures=node_data_futures,
+                                sub_nodes_info=sub_nodes_info,
+                                sub_query_ids=sub_query_ids,
+                                pre_req_nodes=pre_req_nodes
                             )
                         )
                     except Exception as e:
@@ -721,9 +666,154 @@ class GraphRAG:
                     try:
                         await asyncio.gather(*recursion_tasks)
                     except Exception as e:
-                        raise Exception(f"Error during recursive processing: {e}")
-                    
+                        print(f"Error during recursive processing: {e}")
+                        raise e
+
+            futures = {}
+            all_child_futures = {}
+            process_tasks = []
+            graph_data = self._get_current_graph_data()
+            graph = graph_data["graph"]
+            node_map = graph_data["node_map"]
+
+            for (sub_node_id, sub_query, role, dependency, future, local_depth) in sub_nodes_info:
+                idx = node_map.get(sub_node_id)
+                has_children = False
+                child_futures = []
+                if idx is not None:
+                    for (_, child_idx, edge_data) in graph.out_edges(idx):
+                        if edge_data.get("type") == "hierarchical":
+                            has_children = True
+                            child_future = node_data_futures.get(graph.get_node_data(child_idx).get("id"))
+                            if child_future:
+                                child_futures.append(child_future)
+                if local_depth == 0:
+                    futures[sub_query] = future
+                    all_child_futures[sub_query] = child_futures
+                    if has_children:
+                        if not future.done():
+                            future.set_result("")
+                    else:
+                        if role.lower() in ('pre-requisite', 'prerequisite', 'independent'):
+                            process_tasks.append(process_node(sub_node_id, sub_query, session_id, future, max_tokens_allowed))
+                        elif role.lower() == 'dependent':
+                            dep_futures = []
+                            if isinstance(dependency, list) and len(dependency) == 2:
+                                prev_deps, current_deps = dependency
+                                if context and prev_deps not in [None, []]:
+                                    for context_idx, context_data in enumerate(context):
+                                        if isinstance(prev_deps, list) and context_idx < len(prev_deps):
+                                            context_dep = prev_deps[context_idx]
+                                            if (context_dep is not None and isinstance(context_data, dict)
+                                                    and 'subqueries' in context_data):
+                                                if context_dep < len(context_data['subqueries']):
+                                                    dep_query = context_data['subqueries'][context_dep]['subquery']
+                                                    matching_nodes = self.find_nodes_by_properties(query=dep_query)
+                                                    if matching_nodes not in [None, []]:
+                                                        dep_node_id = matching_nodes[0].get('node_id', None)
+                                                        score = float(matching_nodes[0].get('score', 0))
+                                                        if score == 1.0 and dep_node_id in node_data_futures:
+                                                            dep_futures.append(node_data_futures[dep_node_id])
+                                        elif isinstance(prev_deps, int):
+                                            if context_idx < len(context_data['subqueries']):
+                                                dep_query = context_data['subqueries'][prev_deps]['subquery']
+                                                matching_nodes = self.find_nodes_by_properties(query=dep_query)
+                                                if matching_nodes not in [None, []]:
+                                                    dep_node_id = matching_nodes[0].get('node_id', None)
+                                                    score = matching_nodes[0].get('score', 0)
+                                                    if score == 1.0 and dep_node_id in node_data_futures:
+                                                        dep_futures.append(node_data_futures[dep_node_id])
+                                if current_deps not in [None, []]:
+                                    current_deps_list = [current_deps] if isinstance(current_deps, int) else current_deps
+                                    for dep_idx in current_deps_list:
+                                        if dep_idx < len(sub_query_ids):
+                                            dep_node_id = sub_query_ids[dep_idx]
+                                            if dep_node_id in node_data_futures:
+                                                dep_futures.append(node_data_futures[dep_node_id])
+                            process_tasks.append(process_dependent_node(sub_node_id, sub_query, dep_futures, future))
+                else:
+                    if role.lower() in ('pre-requisite', 'prerequisite', 'independent'):
+                        process_tasks.append(process_node(sub_node_id, sub_query, session_id, future, max_tokens_allowed))
+                    elif role.lower() == 'dependent':
+                        dep_futures = []
+                        if isinstance(dependency, list) and len(dependency) == 2:
+                            prev_deps, current_deps = dependency
+                            if context and prev_deps not in [None, []]:
+                                for context_idx, context_data in enumerate(context):
+                                    if isinstance(prev_deps, list) and context_idx < len(prev_deps):
+                                        context_dep = prev_deps[context_idx]
+                                        if (context_dep is not None and isinstance(context_data, dict)
+                                                and 'subqueries' in context_data):
+                                            if context_dep < len(context_data['subqueries']):
+                                                dep_query = context_data['subqueries'][context_dep]['subquery']
+                                                matching_nodes = self.find_nodes_by_properties(query=dep_query)
+                                                if matching_nodes not in [None, []]:
+                                                    dep_node_id = matching_nodes[0].get('node_id', None)
+                                                    score = float(matching_nodes[0].get('score', 0))
+                                                    if score == 1.0 and dep_node_id in node_data_futures:
+                                                        dep_futures.append(node_data_futures[dep_node_id])
+                                    elif isinstance(prev_deps, int):
+                                        if context_idx < len(context_data['subqueries']):
+                                            dep_query = context_data['subqueries'][prev_deps]['subquery']
+                                            matching_nodes = self.find_nodes_by_properties(query=dep_query)
+                                            if matching_nodes not in [None, []]:
+                                                dep_node_id = matching_nodes[0].get('node_id', None)
+                                                score = matching_nodes[0].get('score', 0)
+                                                if score == 1.0 and dep_node_id in node_data_futures:
+                                                    dep_futures.append(node_data_futures[dep_node_id])
+                            if current_deps not in [None, []]:
+                                current_deps_list = [current_deps] if isinstance(current_deps, int) else current_deps
+                                for dep_idx in current_deps_list:
+                                    if dep_idx < len(sub_query_ids):
+                                        dep_node_id = sub_query_ids[dep_idx]
+                                        if dep_node_id in node_data_futures:
+                                            dep_futures.append(node_data_futures[dep_node_id])
+                        process_tasks.append(process_dependent_node(sub_node_id, sub_query, dep_futures, future))
+
+            if process_tasks:
+                await self.emit_event("search_process_started", {
+                    "depth": depth,
+                    "sub_queries": sub_queries,
+                    "roles": roles
+                })
+
+                for sub_query, future in futures.items():
+                    try:
+                        parent_content = future.result().strip()
+                    except:
+                        parent_content = ""
+
+                    child_futures = all_child_futures.get(sub_query)
+                    any_child_done = any(cf.done() and cf.result().strip() for cf in child_futures)
+
+                    if parent_content or any_child_done:
+                        await self.emit_event("sub_query_processed", {"sub_query": sub_query})
+
+                await asyncio.gather(*process_tasks)
+
             if depth == 0:
+                for sub_query, future in futures.items():
+                    try:
+                        parent_content = future.result().strip()
+                    except:
+                        parent_content = ""
+
+                    child_futures = all_child_futures.get(sub_query)
+                    no_child_done = not any(cf.done() and cf.result().strip() for cf in child_futures)
+
+                    if no_child_done:
+                        await self.emit_event("sub_query_failed", {"sub_query": sub_query})
+
+                for idx, (sub_query, future) in enumerate(futures.items(), 1):
+                    if future.done() and future.result().strip():
+                        print(f"Sub-query {idx} processed successfully")
+                    else:
+                        child_futures = all_child_futures.get(sub_query)
+                        if any(cf.done() and cf.result().strip() for cf in child_futures):
+                            print(f"Sub-query {idx} processed successfully because of child nodes")
+                        else:
+                            print(f"Sub-query {idx} failed to process because of child nodes")
+
                 print("Graph building complete, processing final tasks...")
                 await self.emit_event("search_process_completed", {
                     "depth": depth,
@@ -916,6 +1006,8 @@ class GraphRAG:
                 self.update_pagerank()
                 self.verify_graph_integrity()
                 self.verify_graph_consistency()
+        except (RateLimitError, ResourceExhausted, AnthropicRateLimitError):
+            pass
         except Exception as e:
             print(f"Error in process_graph: {str(e)}")
             raise
@@ -1183,7 +1275,7 @@ class GraphRAG:
         graph_data = self._get_current_graph_data()
         graph = graph_data["graph"]
         node_map = graph_data["node_map"]
-        net = Network(height="465px", width="100%", directed=True, bgcolor="#222222", font_color="white")
+        net = Network(height="530px", width="100%", directed=True, bgcolor="#222222", font_color="white")
         net.options = {"physics": {"enabled": False}}
         all_nodes = set()
         all_edges = []
@@ -1220,17 +1312,11 @@ class GraphRAG:
         net.options["layout"] = {"improvedLayout": True}
         net.options["interaction"] = {"dragNodes": True}
 
-        original_dir = os.getcwd()
-        os.chdir(os.getenv("WRITABLE_DIR", "/tmp"))
-
         net.save_graph("temp_graph.html")
 
         with open("temp_graph.html", "r", encoding="utf-8") as f:
             html_str = f.read()
-            
         os.remove("temp_graph.html")
-        os.chdir(original_dir)
-
         return html_str
 
     def verify_graph_integrity(self):