Phase 2

trustgraph-cli/trustgraph/cli/invoke_agent.py

@@ -5,12 +5,10 @@ Uses the agent service to answer a question
 import argparse
 import os
 import textwrap
-import uuid
+from trustgraph.api import Api
-import asyncio
-import json
-from websockets.asyncio.client import connect
 
-default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
+default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 default_user = 'trustgraph'
 default_collection = 'default'
 
@@ -99,79 +97,47 @@ def output(text, prefix="> ", width=78):
     )
     print(out)
 
-async def question(
+def question(
         url, question, flow_id, user, collection,
-        plan=None, state=None, group=None, verbose=False, streaming=True
+        plan=None, state=None, group=None, verbose=False, streaming=True,
+        token=None
 ):
 
-    if not url.endswith("/"):
-        url += "/"
-
-    url = url + "api/v1/socket"
-
     if verbose:
         output(wrap(question), "\U00002753 ")
         print()
 
-    # Track last chunk type and current outputter for streaming
+    # Create API client
-    last_chunk_type = None
+    api = Api(url=url, token=token)
-    current_outputter = None
+    socket = api.socket()
+    flow = socket.flow(flow_id)
 
-    def think(x):
+    # Prepare request parameters
-        if verbose:
+    request_params = {
-            output(wrap(x), "\U0001f914 ")
+        "question": question,
-            print()
+        "user": user,
+        "streaming": streaming,
+    }
 
-    def observe(x):
+    # Only add optional fields if they have values
-        if verbose:
+    if state is not None:
-            output(wrap(x), "\U0001f4a1 ")
+        request_params["state"] = state
-            print()
+    if group is not None:
+        request_params["group"] = group
 
-    mid = str(uuid.uuid4())
+    try:
+        # Call agent
+        response = flow.agent(**request_params)
 
-    async with connect(url) as ws:
+        # Handle streaming response
+        if streaming:
+            # Track last chunk type and current outputter for streaming
+            last_chunk_type = None
+            current_outputter = None
 
-        req = {
+            for chunk in response:
-            "id": mid,
+                chunk_type = chunk.chunk_type
-            "service": "agent",
+                content = chunk.content
-            "flow": flow_id,
-            "request": {
-                "question": question,
-                "user": user,
-                "history": [],
-                "streaming": streaming
-            }
-        }
-
-        # Only add optional fields if they have values
-        if state is not None:
-            req["request"]["state"] = state
-        if group is not None:
-            req["request"]["group"] = group
-
-        req = json.dumps(req)
-
-        await ws.send(req)
-
-        while True:
-
-            msg = await ws.recv()
-
-            obj = json.loads(msg)
-
-            if "error" in obj:
-                raise RuntimeError(obj["error"])
-
-            if obj["id"] != mid:
-                print("Ignore message")
-                continue
-
-            response = obj["response"]
-
-            # Handle streaming format (new format with chunk_type)
-            if "chunk_type" in response:
-                chunk_type = response["chunk_type"]
-                content = response.get("content", "")
 
                 # Check if we're switching to a new message type
                 if last_chunk_type != chunk_type:
@@ -195,33 +161,27 @@ async def question(
                 # Output the chunk
                 if current_outputter:
                     current_outputter.output(content)
-                elif chunk_type == "answer":
+                elif chunk_type == "final-answer":
                     print(content, end="", flush=True)
-            else:
-                # Handle legacy format (backward compatibility)
-                if "thought" in response:
-                    think(response["thought"])
 
-                if "observation" in response:
+            # Close any remaining outputter
-                    observe(response["observation"])
+            if current_outputter:
+                current_outputter.__exit__(None, None, None)
+                current_outputter = None
+            # Add final newline if we were outputting answer
+            elif last_chunk_type == "final-answer":
+                print()
 
-                if "answer" in response:
+        else:
-                    print(response["answer"])
+            # Non-streaming response
+            if "answer" in response:
+                print(response["answer"])
+            if "error" in response:
+                raise RuntimeError(response["error"])
 
-                if "error" in response:
+    finally:
-                    raise RuntimeError(response["error"])
+        # Clean up socket connection
-
+        socket.close()
-            if obj["complete"]:
-                # Close any remaining outputter
-                if current_outputter:
-                    current_outputter.__exit__(None, None, None)
-                    current_outputter = None
-                # Add final newline if we were outputting answer
-                elif last_chunk_type == "answer":
-                    print()
-                break
-
-        await ws.close()
 
 def main():
 
@@ -236,6 +196,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-f', '--flow-id',
         default="default",
@@ -292,19 +258,18 @@ def main():
 
     try:
 
-        asyncio.run(
+        question(
-            question(
+            url = args.url,
-                url = args.url,
+            flow_id = args.flow_id,
-                flow_id = args.flow_id,
+            question = args.question,
-                question = args.question,
+            user = args.user,
-                user = args.user,
+            collection = args.collection,
-                collection = args.collection,
+            plan = args.plan,
-                plan = args.plan,
+            state = args.state,
-                state = args.state,
+            group = args.group,
-                group = args.group,
+            verbose = args.verbose,
-                verbose = args.verbose,
+            streaming = not args.no_streaming,
-                streaming = not args.no_streaming,
+            token = args.token,
-            )
         )
 
     except Exception as e:

trustgraph-cli/trustgraph/cli/invoke_document_rag.py

@@ -4,89 +4,50 @@ Uses the DocumentRAG service to answer a question
 
 import argparse
 import os
-import asyncio
-import json
-import uuid
-from websockets.asyncio.client import connect
 from trustgraph.api import Api
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 default_user = 'trustgraph'
 default_collection = 'default'
 default_doc_limit = 10
 
-async def question_streaming(url, flow_id, question, user, collection, doc_limit):
+def question(url, flow_id, question, user, collection, doc_limit, streaming=True, token=None):
-    """Streaming version using websockets"""
 
-    # Convert http:// to ws://
+    # Create API client
-    if url.startswith('http://'):
+    api = Api(url=url, token=token)
-        url = 'ws://' + url[7:]
-    elif url.startswith('https://'):
-        url = 'wss://' + url[8:]
 
-    if not url.endswith("/"):
+    if streaming:
-        url += "/"
+        # Use socket client for streaming
+        socket = api.socket()
+        flow = socket.flow(flow_id)
 
-    url = url + "api/v1/socket"
+        try:
+            response = flow.document_rag(
+                question=question,
+                user=user,
+                collection=collection,
+                doc_limit=doc_limit,
+                streaming=True
+            )
 
-    mid = str(uuid.uuid4())
+            # Stream output
+            for chunk in response:
+                print(chunk.content, end="", flush=True)
+            print()  # Final newline
 
-    async with connect(url) as ws:
+        finally:
-        req = {
+            socket.close()
-            "id": mid,
+    else:
-            "service": "document-rag",
+        # Use REST API for non-streaming
-            "flow": flow_id,
+        flow = api.flow().id(flow_id)
-            "request": {
+        resp = flow.document_rag(
-                "query": question,
+            question=question,
-                "user": user,
+            user=user,
-                "collection": collection,
+            collection=collection,
-                "doc-limit": doc_limit,
+            doc_limit=doc_limit,
-                "streaming": True
+        )
-            }
+        print(resp)
-        }
-
-        req = json.dumps(req)
-        await ws.send(req)
-
-        while True:
-            msg = await ws.recv()
-            obj = json.loads(msg)
-
-            if "error" in obj:
-                raise RuntimeError(obj["error"])
-
-            if obj["id"] != mid:
-                print("Ignore message")
-                continue
-
-            response = obj["response"]
-
-            # Handle streaming format (chunk)
-            if "chunk" in response:
-                chunk = response["chunk"]
-                print(chunk, end="", flush=True)
-            elif "response" in response:
-                # Final response with complete text
-                # Already printed via chunks, just add newline
-                pass
-
-            if obj["complete"]:
-                print()  # Final newline
-                break
-
-        await ws.close()
-
-def question_non_streaming(url, flow_id, question, user, collection, doc_limit):
-    """Non-streaming version using HTTP API"""
-
-    api = Api(url).flow().id(flow_id)
-
-    resp = api.document_rag(
-        question=question, user=user, collection=collection,
-        doc_limit=doc_limit,
-    )
-
-    print(resp)
 
 def main():
 
@@ -101,6 +62,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-f', '--flow-id',
         default="default",
@@ -127,6 +94,7 @@ def main():
 
     parser.add_argument(
         '-d', '--doc-limit',
+        type=int,
         default=default_doc_limit,
         help=f'Document limit (default: {default_doc_limit})'
     )
@@ -141,30 +109,20 @@ def main():
 
     try:
 
-        if not args.no_streaming:
+        question(
-            asyncio.run(
+            url=args.url,
-                question_streaming(
+            flow_id=args.flow_id,
-                    url=args.url,
+            question=args.question,
-                    flow_id=args.flow_id,
+            user=args.user,
-                    question=args.question,
+            collection=args.collection,
-                    user=args.user,
+            doc_limit=args.doc_limit,
-                    collection=args.collection,
+            streaming=not args.no_streaming,
-                    doc_limit=args.doc_limit,
+            token=args.token,
-                )
+        )
-            )
-        else:
-            question_non_streaming(
-                url=args.url,
-                flow_id=args.flow_id,
-                question=args.question,
-                user=args.user,
-                collection=args.collection,
-                doc_limit=args.doc_limit,
-            )
 
     except Exception as e:
 
         print("Exception:", e, flush=True)
 
 if __name__ == "__main__":
-    main()
+    main()

trustgraph-cli/trustgraph/cli/invoke_graph_rag.py

@@ -4,13 +4,10 @@ Uses the GraphRAG service to answer a question
 
 import argparse
 import os
-import asyncio
-import json
-import uuid
-from websockets.asyncio.client import connect
 from trustgraph.api import Api
 
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 default_user = 'trustgraph'
 default_collection = 'default'
 default_entity_limit = 50
@@ -18,89 +15,51 @@ default_triple_limit = 30
 default_max_subgraph_size = 150
 default_max_path_length = 2
 
-async def question_streaming(
+def question(
         url, flow_id, question, user, collection, entity_limit, triple_limit,
-        max_subgraph_size, max_path_length
+        max_subgraph_size, max_path_length, streaming=True, token=None
 ):
-    """Streaming version using websockets"""
 
-    # Convert http:// to ws://
+    # Create API client
-    if url.startswith('http://'):
+    api = Api(url=url, token=token)
-        url = 'ws://' + url[7:]
-    elif url.startswith('https://'):
-        url = 'wss://' + url[8:]
 
-    if not url.endswith("/"):
+    if streaming:
-        url += "/"
+        # Use socket client for streaming
+        socket = api.socket()
+        flow = socket.flow(flow_id)
 
-    url = url + "api/v1/socket"
+        try:
+            response = flow.graph_rag(
+                question=question,
+                user=user,
+                collection=collection,
+                entity_limit=entity_limit,
+                triple_limit=triple_limit,
+                max_subgraph_size=max_subgraph_size,
+                max_path_length=max_path_length,
+                streaming=True
+            )
 
-    mid = str(uuid.uuid4())
+            # Stream output
+            for chunk in response:
+                print(chunk.content, end="", flush=True)
+            print()  # Final newline
 
-    async with connect(url) as ws:
+        finally:
-        req = {
+            socket.close()
-            "id": mid,
+    else:
-            "service": "graph-rag",
+        # Use REST API for non-streaming
-            "flow": flow_id,
+        flow = api.flow().id(flow_id)
-            "request": {
+        resp = flow.graph_rag(
-                "query": question,
+            question=question,
-                "user": user,
+            user=user,
-                "collection": collection,
+            collection=collection,
-                "entity-limit": entity_limit,
+            entity_limit=entity_limit,
-                "triple-limit": triple_limit,
+            triple_limit=triple_limit,
-                "max-subgraph-size": max_subgraph_size,
+            max_subgraph_size=max_subgraph_size,
-                "max-path-length": max_path_length,
+            max_path_length=max_path_length
-                "streaming": True
+        )
-            }
+        print(resp)
-        }
-
-        req = json.dumps(req)
-        await ws.send(req)
-
-        while True:
-            msg = await ws.recv()
-            obj = json.loads(msg)
-
-            if "error" in obj:
-                raise RuntimeError(obj["error"])
-
-            if obj["id"] != mid:
-                print("Ignore message")
-                continue
-
-            response = obj["response"]
-
-            # Handle streaming format (chunk)
-            if "chunk" in response:
-                chunk = response["chunk"]
-                print(chunk, end="", flush=True)
-            elif "response" in response:
-                # Final response with complete text
-                # Already printed via chunks, just add newline
-                pass
-
-            if obj["complete"]:
-                print()  # Final newline
-                break
-
-        await ws.close()
-
-def question_non_streaming(
-        url, flow_id, question, user, collection, entity_limit, triple_limit,
-        max_subgraph_size, max_path_length
-):
-    """Non-streaming version using HTTP API"""
-
-    api = Api(url).flow().id(flow_id)
-
-    resp = api.graph_rag(
-        question=question, user=user, collection=collection,
-        entity_limit=entity_limit, triple_limit=triple_limit,
-        max_subgraph_size=max_subgraph_size,
-        max_path_length=max_path_length
-    )
-
-    print(resp)
 
 def main():
 
@@ -115,6 +74,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-f', '--flow-id',
         default="default",
@@ -141,24 +106,28 @@ def main():
 
     parser.add_argument(
         '-e', '--entity-limit',
+        type=int,
         default=default_entity_limit,
         help=f'Entity limit (default: {default_entity_limit})'
     )
 
     parser.add_argument(
-        '-t', '--triple-limit',
+        '--triple-limit',
+        type=int,
         default=default_triple_limit,
         help=f'Triple limit (default: {default_triple_limit})'
     )
 
     parser.add_argument(
         '-s', '--max-subgraph-size',
+        type=int,
         default=default_max_subgraph_size,
         help=f'Max subgraph size (default: {default_max_subgraph_size})'
     )
 
     parser.add_argument(
         '-p', '--max-path-length',
+        type=int,
         default=default_max_path_length,
         help=f'Max path length (default: {default_max_path_length})'
     )
@@ -173,36 +142,23 @@ def main():
 
     try:
 
-        if not args.no_streaming:
+        question(
-            asyncio.run(
+            url=args.url,
-                question_streaming(
+            flow_id=args.flow_id,
-                    url=args.url,
+            question=args.question,
-                    flow_id=args.flow_id,
+            user=args.user,
-                    question=args.question,
+            collection=args.collection,
-                    user=args.user,
+            entity_limit=args.entity_limit,
-                    collection=args.collection,
+            triple_limit=args.triple_limit,
-                    entity_limit=args.entity_limit,
+            max_subgraph_size=args.max_subgraph_size,
-                    triple_limit=args.triple_limit,
+            max_path_length=args.max_path_length,
-                    max_subgraph_size=args.max_subgraph_size,
+            streaming=not args.no_streaming,
-                    max_path_length=args.max_path_length,
+            token=args.token,
-                )
+        )
-            )
-        else:
-            question_non_streaming(
-                url=args.url,
-                flow_id=args.flow_id,
-                question=args.question,
-                user=args.user,
-                collection=args.collection,
-                entity_limit=args.entity_limit,
-                triple_limit=args.triple_limit,
-                max_subgraph_size=args.max_subgraph_size,
-                max_path_length=args.max_path_length,
-            )
 
     except Exception as e:
 
         print("Exception:", e, flush=True)
 
 if __name__ == "__main__":
-    main()
+    main()

trustgraph-cli/trustgraph/cli/invoke_llm.py

@@ -5,64 +5,39 @@ and user prompt.  Both arguments are required.
 
 import argparse
 import os
-import json
+from trustgraph.api import Api
-import uuid
-import asyncio
-from websockets.asyncio.client import connect
 
-default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
+default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-async def query(url, flow_id, system, prompt, streaming=True):
+def query(url, flow_id, system, prompt, streaming=True, token=None):
 
-    if not url.endswith("/"):
+    # Create API client
-        url += "/"
+    api = Api(url=url, token=token)
+    socket = api.socket()
+    flow = socket.flow(flow_id)
 
-    url = url + "api/v1/socket"
+    try:
+        # Call text completion
+        response = flow.text_completion(
+            system=system,
+            prompt=prompt,
+            streaming=streaming
+        )
 
-    mid = str(uuid.uuid4())
+        if streaming:
+            # Stream output to stdout without newline
+            for chunk in response:
+                print(chunk.content, end="", flush=True)
+            # Add final newline after streaming
+            print()
+        else:
+            # Non-streaming: print complete response
+            print(response)
 
-    async with connect(url) as ws:
+    finally:
-
+        # Clean up socket connection
-        req = {
+        socket.close()
-            "id": mid,
-            "service": "text-completion",
-            "flow": flow_id,
-            "request": {
-                "system": system,
-                "prompt": prompt,
-                "streaming": streaming
-            }
-        }
-
-        await ws.send(json.dumps(req))
-
-        while True:
-
-            msg = await ws.recv()
-
-            obj = json.loads(msg)
-
-            if "error" in obj:
-                raise RuntimeError(obj["error"])
-
-            if obj["id"] != mid:
-                continue
-
-            if "response" in obj["response"]:
-                if streaming:
-                    # Stream output to stdout without newline
-                    print(obj["response"]["response"], end="", flush=True)
-                else:
-                    # Non-streaming: print complete response
-                    print(obj["response"]["response"])
-
-            if obj["complete"]:
-                if streaming:
-                    # Add final newline after streaming
-                    print()
-                break
-
-        await ws.close()
 
 def main():
 
@@ -77,6 +52,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         'system',
         nargs=1,
@@ -105,17 +86,18 @@ def main():
 
     try:
 
-        asyncio.run(query(
+        query(
             url=args.url,
             flow_id=args.flow_id,
             system=args.system[0],
             prompt=args.prompt[0],
-            streaming=not args.no_streaming
+            streaming=not args.no_streaming,
-        ))
+            token=args.token,
+        )
 
     except Exception as e:
 
         print("Exception:", e, flush=True)
 
 if __name__ == "__main__":
-    main()
+    main()

trustgraph-cli/trustgraph/cli/invoke_prompt.py

@@ -10,76 +10,61 @@ using key=value arguments on the command line, and these replace
 import argparse
 import os
 import json
-import uuid
+from trustgraph.api import Api
-import asyncio
-from websockets.asyncio.client import connect
 
-default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
+default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
+default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-async def query(url, flow_id, template_id, variables, streaming=True):
+def query(url, flow_id, template_id, variables, streaming=True, token=None):
 
-    if not url.endswith("/"):
+    # Create API client
-        url += "/"
+    api = Api(url=url, token=token)
+    socket = api.socket()
+    flow = socket.flow(flow_id)
 
-    url = url + "api/v1/socket"
+    try:
+        # Call prompt
+        response = flow.prompt(
+            id=template_id,
+            variables=variables,
+            streaming=streaming
+        )
 
-    mid = str(uuid.uuid4())
+        if streaming:
+            full_response = {"text": "", "object": ""}
 
-    async with connect(url) as ws:
+            # Stream output
+            for chunk in response:
+                content = chunk.content
+                if content:
+                    print(content, end="", flush=True)
+                    full_response["text"] += content
 
-        req = {
+                # Check if this is an object response (JSON)
-            "id": mid,
+                if hasattr(chunk, 'object') and chunk.object:
-            "service": "prompt",
+                    full_response["object"] = chunk.object
-            "flow": flow_id,
-            "request": {
-                "id": template_id,
-                "variables": variables,
-                "streaming": streaming
-            }
-        }
 
-        await ws.send(json.dumps(req))
+            # Handle final output
+            if full_response["text"]:
+                # Add final newline after streaming text
+                print()
+            elif full_response["object"]:
+                # Print JSON object (pretty-printed)
+                print(json.dumps(json.loads(full_response["object"]), indent=4))
 
-        full_response = {"text": "", "object": ""}
+        else:
-
+            # Non-streaming: handle response
-        while True:
+            if isinstance(response, str):
-
+                print(response)
-            msg = await ws.recv()
+            elif isinstance(response, dict):
-
+                if "text" in response:
-            obj = json.loads(msg)
-
-            if "error" in obj:
-                raise RuntimeError(obj["error"])
-
-            if obj["id"] != mid:
-                continue
-
-            response = obj["response"]
-
-            # Handle text responses (streaming)
-            if "text" in response and response["text"]:
-                if streaming:
-                    # Stream output to stdout without newline
-                    print(response["text"], end="", flush=True)
-                    full_response["text"] += response["text"]
-                else:
-                    # Non-streaming: print complete response
                     print(response["text"])
+                elif "object" in response:
+                    print(json.dumps(json.loads(response["object"]), indent=4))
 
-            # Handle object responses (JSON, never streamed)
+    finally:
-            if "object" in response and response["object"]:
+        # Clean up socket connection
-                full_response["object"] = response["object"]
+        socket.close()
-
-            if obj["complete"]:
-                if streaming and full_response["text"]:
-                    # Add final newline after streaming text
-                    print()
-                elif full_response["object"]:
-                    # Print JSON object (pretty-printed)
-                    print(json.dumps(json.loads(full_response["object"]), indent=4))
-                break
-
-        await ws.close()
 
 def main():
 
@@ -94,6 +79,12 @@ def main():
         help=f'API URL (default: {default_url})',
     )
 
+    parser.add_argument(
+        '-t', '--token',
+        default=default_token,
+        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
+    )
+
     parser.add_argument(
         '-f', '--flow-id',
         default="default",
@@ -135,17 +126,18 @@ specified multiple times''',
 
     try:
 
-        asyncio.run(query(
+        query(
             url=args.url,
             flow_id=args.flow_id,
             template_id=args.id[0],
             variables=variables,
-            streaming=not args.no_streaming
+            streaming=not args.no_streaming,
-        ))
+            token=args.token,
+        )
 
     except Exception as e:
 
         print("Exception:", e, flush=True)
 
 if __name__ == "__main__":
-    main()
+    main()