From c4479afbad1dfeb643dae419901c9f0f4e6d786a Mon Sep 17 00:00:00 2001
From: Parker Erickson <parker.erickson@tigergraph.com>
Date: Fri, 22 Nov 2024 08:27:38 -0600
Subject: [PATCH] add graphrag to ui

---
 copilot-ui/src/components/Bot.tsx       |  2 +-
 copilot/app/agent/agent_graph.py        | 30 ++++++++++++++++++++++++-
 eventual-consistency-service/Dockerfile |  2 +-
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/copilot-ui/src/components/Bot.tsx b/copilot-ui/src/components/Bot.tsx
index 274aeeae..19ac9dad 100644
--- a/copilot-ui/src/components/Bot.tsx
+++ b/copilot-ui/src/components/Bot.tsx
@@ -67,7 +67,7 @@ const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getCo
               <DropdownMenuLabel>Select a RAG Pattern</DropdownMenuLabel>
               <DropdownMenuSeparator />
               <DropdownMenuGroup>
-                {["HNSW", "HNSW_Overlap", "Sibling"].map((f, i) => (
+                {["HNSW", "HNSW_Overlap", "Sibling", "GraphRAG"].map((f, i) => (
                   <DropdownMenuItem key={i} onSelect={() => handleSelectRag(f)}>
                     {/* <User className="mr-2 h-4 w-4" /> */}
                     <span>{f}</span>
diff --git a/copilot/app/agent/agent_graph.py b/copilot/app/agent/agent_graph.py
index c9781733..93afb79c 100644
--- a/copilot/app/agent/agent_graph.py
+++ b/copilot/app/agent/agent_graph.py
@@ -11,7 +11,7 @@
 from langgraph.graph import END, StateGraph
 from pyTigerGraph.pyTigerGraphException import TigerGraphException
 from supportai.retrievers import (HNSWOverlapRetriever, HNSWRetriever,
-                                  HNSWSiblingRetriever)
+                                  HNSWSiblingRetriever, GraphRAG)
 from tools import MapQuestionToSchemaException
 from typing_extensions import TypedDict
 
@@ -272,6 +272,32 @@ def sibling_search(self, state):
         state["lookup_source"] = "supportai"
         return state
     
+    def graphrag_search(self, state):
+        """
+        Run the agent graphrag search.
+        """
+        self.emit_progress("Searching the knowledge graph")
+        retriever = GraphRAG(
+            self.embedding_model,
+            self.embedding_store,
+            self.llm_provider.model,
+            self.db_connection,
+        )
+        step = retriever.search(
+            state["question"],
+            community_level=3
+        )
+
+        state["context"] = {
+            "function_call": "GraphRAG",
+            "result": {"@@final_retrieval": step[0]},
+            "query_output_format": self.db_connection.getQueryMetadata(
+                "GraphRAG_Community_Retriever"
+            )["output"],
+        }
+        state["lookup_source"] = "supportai"
+        return state
+    
     def supportai_search(self, state):
         """
         Run the agent supportai search.
@@ -283,6 +309,8 @@ def supportai_search(self, state):
             return self.hnsw_search(state)
         elif self.supportai_retriever == "sibling":
             return self.sibling_search(state)
+        elif self.supportai_retriever == "graphrag":
+            return self.graphrag_search(state)
         else:
             raise ValueError(f"Invalid supportai retriever: {self.supportai_retriever}")
     
diff --git a/eventual-consistency-service/Dockerfile b/eventual-consistency-service/Dockerfile
index bd31c0fe..8694f2c0 100644
--- a/eventual-consistency-service/Dockerfile
+++ b/eventual-consistency-service/Dockerfile
@@ -1,4 +1,4 @@
-FROM python:3.11.8
+FROM python:3.11.8-bullseye
 WORKDIR /code
 
 COPY common/requirements.txt requirements.txt