added api respose for context extraction in interactive mode

BaranziniLab · Dec 4, 2023 · b4a31e6 · b4a31e6
1 parent 7012cb4
commit b4a31e6
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/kg_rag/utility.py b/kg_rag/utility.py
@@ -315,7 +315,10 @@ def interactive(question, vectorstore, node_context_df, embedding_function_for_c
     question_embedding = embedding_function_for_context_retrieval.embed_query(question)
     node_context_extracted = ""
     for node_name in node_hits:
-        node_context = node_context_df[node_context_df.node_name == node_name].node_context.values[0]
+        if not api:
+            node_context = node_context_df[node_context_df.node_name == node_name].node_context.values[0]
+        else:
+            node_context = get_context_using_spoke_api(node_name)
         node_context_list = node_context.split(". ")        
         node_context_embeddings = embedding_function_for_context_retrieval.embed_documents(node_context_list)
         similarities = [cosine_similarity(np.array(question_embedding).reshape(1, -1), np.array(node_context_embedding).reshape(1, -1)) for node_context_embedding in node_context_embeddings]