From 7c191bd38c2c0f028bc1e90bc61ac4a4be3eb993 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Tue, 16 Jul 2024 19:51:08 +0100
Subject: [PATCH] Configurable RAG query

---
 trustgraph/graph_rag.py     | 11 +++++++----
 trustgraph/rag/graph/rag.py | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/trustgraph/graph_rag.py b/trustgraph/graph_rag.py
index d43ece48..b106e99e 100644
--- a/trustgraph/graph_rag.py
+++ b/trustgraph/graph_rag.py
@@ -15,7 +15,10 @@ class GraphRag:
             graph_hosts=None,
             pulsar_host="pulsar://pulsar:6650",
             vector_store="http://milvus:19530",
-            verbose=False
+            verbose=False,
+            entity_limit=50,
+            triple_limit=30,
+            max_sg_size=3000,
     ):
 
         self.verbose=verbose
@@ -32,9 +35,9 @@ class GraphRag:
 
         self.vecstore = TripleVectors(vector_store)
 
-        self.entity_limit=50
-        self.query_limit=30
-        self.max_sg_size=3000
+        self.entity_limit=entity_limit
+        self.query_limit=triple_limit
+        self.max_sg_size=max_sg_size
 
         self.label_cache = {}
 
diff --git a/trustgraph/rag/graph/rag.py b/trustgraph/rag/graph/rag.py
index ba4e1335..0ac7b12e 100755
--- a/trustgraph/rag/graph/rag.py
+++ b/trustgraph/rag/graph/rag.py
@@ -34,6 +34,9 @@ class Processor:
             log_level=LogLevel.INFO,
             graph_hosts=default_graph_hosts,
             vector_store=default_vector_store,
+            entity_limit=50,
+            triple_limit=30,
+            max_sg_size=3000,
     ):
 
         self.client = None
@@ -58,6 +61,9 @@ class Processor:
             graph_hosts=graph_hosts,
             vector_store=vector_store,
             verbose=True,
+            entity_limit=entity_limit,
+            triple_limit=triple_limit,
+            max_sg_size=max_sg_size,
         )
 
     def run(self):
@@ -102,7 +108,7 @@ class Processor:
 def run():
 
     parser = argparse.ArgumentParser(
-        prog='llm-ollama-text',
+        prog='graph-rag',
         description=__doc__,
     )
 
@@ -150,6 +156,27 @@ def run():
         help=f'Vector host (default: http://milvus:19530)'
     )
 
+    parser.add_argument(
+        '-e', '--entity-limit',
+        type=int,
+        default=50,
+        help=f'Entity vector fetch limit (default: 50)'
+    )
+
+    parser.add_argument(
+        '-t', '--triple-limit',
+        type=int,
+        default=30,
+        help=f'Triple query limit, per query (default: 30)'
+    )
+
+    parser.add_argument(
+        '-u', '--max-subgraph-size',
+        type=int,
+        default=3000,
+        help=f'Max subgraph size (default: 3000)'
+    )
+
     args = parser.parse_args()
     
     while True:
@@ -164,6 +191,9 @@ def run():
                 log_level=args.log_level,
                 graph_hosts=args.graph_hosts.split(","),
                 vector_store=args.vector_store,
+                entity_limit=args.entity_limit,
+                triple_limit=args.triple_limit,
+                max_sg_size=args.max_subgraph_size,
             )
 
             p.run()