From 906cab2cff96d514699b05cb6e3093d553239e91 Mon Sep 17 00:00:00 2001
From: emrgnt-cmplxty <owen@algofi.org>
Date: Wed, 4 Dec 2024 16:47:12 -0800
Subject: [PATCH] fix llm formatting for rag reply

---
 py/core/pipes/retrieval/search_rag_pipe.py | 26 +++++++++++++---------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/py/core/pipes/retrieval/search_rag_pipe.py b/py/core/pipes/retrieval/search_rag_pipe.py
index f9d922755..dd9416cbf 100644
--- a/py/core/pipes/retrieval/search_rag_pipe.py
+++ b/py/core/pipes/retrieval/search_rag_pipe.py
@@ -7,6 +7,7 @@
     AsyncState,
     CompletionProvider,
     DatabaseProvider,
+    KGSearchResultType
 )
 from core.base.abstractions import GenerationConfig, RAGCompletion
 
@@ -59,7 +60,6 @@ async def _run_logic(  # type: ignore
             )
             context += context_piece
             search_iteration += 1
-
         messages = (
             await self.database_provider.prompt_handler.get_message_payload(
                 system_prompt_name=self.config.system_prompt,
@@ -103,15 +103,21 @@ async def _collect_context(
         if results.graph_search_results:
             context += f"Knowledge Graph ({iteration}):\n"
             it = total_results + 1
-            for search_results in results.graph_search_results:  # [1]:
-                if associated_query := search_results.metadata.get(
-                    "associated_query"
-                ):
-                    context += f"Query: {associated_query}\n\n"
-                context += f"Results:\n"
-                for search_result in search_results:
-                    context += f"[{it}]: {search_result}\n\n"
-                    it += 1
+            for search_result in results.graph_search_results:  # [1]:
+                # if associated_query := search_results.metadata.get(
+                #     "associated_query"
+                # ):
+                #     context += f"Query: {associated_query}\n\n"
+                # context += f"Results:\n"
+                if search_result.result_type == KGSearchResultType.ENTITY:
+                    context += f"[{it}]: Entity Name - {search_result.content.name}\n\nDescription - {search_result.content.description}\n\n"
+                elif search_result.result_type == KGSearchResultType.RELATIONSHIP:
+                    context += f"[{it}]: Relationship - {search_result.content.subject} - {search_result.content.predicate} - {search_result.content.object}\n\n"
+                else:
+                        context += f"[{it}]: Community Name - {search_result.content.name}\n\nDescription - {search_result.content.summary}\n\n"
+
+
+                it += 1
             total_results = (
                 it - 1
             )  # Update total_results based on the last index used