From d75342bae2c112d14d327694ea985d71a88e1d78 Mon Sep 17 00:00:00 2001
From: Mohammad Amin <dadgaramin96@gmail.com>
Date: Tue, 2 Jan 2024 15:34:27 +0330
Subject: [PATCH] feat: Added source node returning!

---
 discord_query.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/discord_query.py b/discord_query.py
index a630f7e..7a61c32 100644
--- a/discord_query.py
+++ b/discord_query.py
@@ -1,4 +1,5 @@
 from llama_index import QueryBundle
+from llama_index.schema import NodeWithScore
 from tc_hivemind_backend.embeddings.cohere import CohereEmbedding
 from utils.query_engine.discord_query_engine import prepare_discord_engine_auto_filter
 
@@ -6,7 +7,7 @@
 def query_discord(
     community_id: str,
     query: str,
-) -> str:
+) -> tuple[str, list[NodeWithScore]]:
     """
     query the llm using the query engine
 
@@ -25,4 +26,4 @@ def query_discord(
         query_str=query, embedding=CohereEmbedding().get_text_embedding(text=query)
     )
     response = query_engine.query(query_bundle)
-    return response
+    return response.response, response.source_nodes