From e96760d1c75b0b38c4b64ae1eb2231eb8e3beb01 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Mon, 23 Feb 2026 20:04:34 +0000
Subject: [PATCH] Add new agent tool limits to CLI

---
 trustgraph-cli/trustgraph/cli/set_tool.py         | 12 +++++++++++-
 trustgraph-cli/trustgraph/cli/show_tools.py       |  2 ++
 trustgraph-flow/trustgraph/agent/react/service.py |  3 ++-
 trustgraph-flow/trustgraph/agent/react/tools.py   |  5 +++--
 4 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/trustgraph-cli/trustgraph/cli/set_tool.py b/trustgraph-cli/trustgraph/cli/set_tool.py
index a647d49d..c6412e48 100644
--- a/trustgraph-cli/trustgraph/cli/set_tool.py
+++ b/trustgraph-cli/trustgraph/cli/set_tool.py
@@ -67,6 +67,7 @@ def set_tool(
         template : str,
         schema_name : str,
         index_name : str,
+        limit : int,
         arguments : List[Argument],
         group : List[str],
         state : str,
@@ -96,6 +97,8 @@ def set_tool(
 
     if index_name: object["index-name"] = index_name
 
+    if limit: object["limit"] = limit
+
     if arguments:
         object["arguments"] = [
             {
@@ -156,7 +159,7 @@ def main():
                        --type row-embeddings-query \\
                        --description "Find customers by name using semantic search" \\
                        --schema-name customers --collection sales \\
-                       --index-name full_name
+                       --index-name full_name --limit 20
 
               %(prog)s --id calc_tool --name calculate --type mcp-tool \\
                        --description "Perform mathematical calculations" \\
@@ -218,6 +221,12 @@ def main():
         help=f'For row-embeddings-query type: specific index to filter search (optional)',
     )
 
+    parser.add_argument(
+        '--limit',
+        type=int,
+        help=f'For row-embeddings-query type: maximum results to return (default: 10)',
+    )
+
     parser.add_argument(
         '--template',
         help=f'For prompt type: template ID to use',
@@ -288,6 +297,7 @@ def main():
             template=args.template,
             schema_name=args.schema_name,
             index_name=args.index_name,
+            limit=args.limit,
             arguments=arguments,
             group=args.group,
             state=args.state,
diff --git a/trustgraph-cli/trustgraph/cli/show_tools.py b/trustgraph-cli/trustgraph/cli/show_tools.py
index 42c79959..d77f1fae 100644
--- a/trustgraph-cli/trustgraph/cli/show_tools.py
+++ b/trustgraph-cli/trustgraph/cli/show_tools.py
@@ -52,6 +52,8 @@ def show_config(url, token=None):
                 table.append(("schema-name", data["schema-name"]))
             if "index-name" in data:
                 table.append(("index-name", data["index-name"]))
+            if "limit" in data:
+                table.append(("limit", data["limit"]))
 
         if tp == "prompt":
             table.append(("template", data["template"]))
diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py
index 68bb5c14..1a44ef9e 100755
--- a/trustgraph-flow/trustgraph/agent/react/service.py
+++ b/trustgraph-flow/trustgraph/agent/react/service.py
@@ -173,7 +173,8 @@ class Processor(AgentService):
                             schema_name=data.get("schema-name"),
                             collection=data.get("collection"),
                             user=None,  # User will be provided dynamically via context
-                            index_name=data.get("index-name")  # Optional filter
+                            index_name=data.get("index-name"),  # Optional filter
+                            limit=int(data.get("limit", 10))  # Max results
                         )
                         arguments = RowEmbeddingsQueryImpl.get_arguments()
                     else:
diff --git a/trustgraph-flow/trustgraph/agent/react/tools.py b/trustgraph-flow/trustgraph/agent/react/tools.py
index c77929d9..2b442a0d 100644
--- a/trustgraph-flow/trustgraph/agent/react/tools.py
+++ b/trustgraph-flow/trustgraph/agent/react/tools.py
@@ -130,12 +130,13 @@ class StructuredQueryImpl:
 
 # This tool implementation knows how to query row embeddings for semantic search
 class RowEmbeddingsQueryImpl:
-    def __init__(self, context, schema_name, collection=None, user=None, index_name=None):
+    def __init__(self, context, schema_name, collection=None, user=None, index_name=None, limit=10):
         self.context = context
         self.schema_name = schema_name
         self.collection = collection
         self.user = user
         self.index_name = index_name  # Optional: filter to specific index
+        self.limit = limit  # Max results to return
 
     @staticmethod
     def get_arguments():
@@ -168,7 +169,7 @@ class RowEmbeddingsQueryImpl:
             user=user,
             collection=self.collection or "default",
             index_name=self.index_name,
-            limit=10
+            limit=self.limit
         )
 
         # Format results for agent consumption