Add new agent tool limits to CLI

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,

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"]))

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:

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