Feature/agent manager (#146)

* Agent schema
* Agent working through client
* Add agent-manager-react command line
* test-agent test script
* Add tg-invoke-agent CLI

tests/test-agent

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+
+import json
+import textwrap
+from trustgraph.clients.agent_client import AgentClient
+
+def wrap(text, width=75):
+
+    if text is None: text = "n/a"
+
+    out = textwrap.wrap(
+        text, width=width
+    )
+    return "\n".join(out)
+
+def output(text, prefix="> ", width=78):
+
+    out = textwrap.indent(
+        text, prefix=prefix
+    )
+    print(out)
+
+p = AgentClient(pulsar_host="pulsar://localhost:6650")
+
+q = "How many cats does Mark have? Calculate that number raised to 0.4 power.  Is that number lower than the numeric part of the mission identifier of the Space Shuttle Challenger on its last mission?  If so, give me an apple pie recipe, otherwise return a poem about cheese."
+
+output(wrap(q), "\U00002753 ")
+print()
+
+def think(x):
+    output(wrap(x), "\U0001f914 ")
+    print()
+
+def observe(x):
+    output(wrap(x), "\U0001f4a1 ")
+    print()
+
+resp = p.request(
+    question=q, think=think, observe=observe,
+)
+
+output(resp, "\U0001f4ac ")
+print()
+

tests/test-llm

@@ -5,9 +5,10 @@ from trustgraph.clients.llm_client import LlmClient
 
 llm = LlmClient(pulsar_host="pulsar://localhost:6650")
 
+system = "You are a lovely assistant."
 prompt="Write a funny limerick about a llama"
 
-resp = llm.request(prompt)
+resp = llm.request(system, prompt)
 
 print(resp)
 

trustgraph-base/trustgraph/base/base_processor.py

@@ -39,8 +39,9 @@ def __init__(self, **params):
 
     def __del__(self):
 
-        if self.client:
-            self.client.close()
+        if hasattr(self, "client"):
+            if self.client:
+                self.client.close()
 
     @staticmethod
     def add_args(parser):

trustgraph-base/trustgraph/clients/agent_client.py

@@ -0,0 +1,64 @@
+
+import _pulsar
+
+from .. schema import AgentRequest, AgentResponse
+from .. schema import agent_request_queue
+from .. schema import agent_response_queue
+from . base import BaseClient
+
+# Ugly
+ERROR=_pulsar.LoggerLevel.Error
+WARN=_pulsar.LoggerLevel.Warn
+INFO=_pulsar.LoggerLevel.Info
+DEBUG=_pulsar.LoggerLevel.Debug
+
+class AgentClient(BaseClient):
+
+    def __init__(
+            self, log_level=ERROR,
+            subscriber=None,
+            input_queue=None,
+            output_queue=None,
+            pulsar_host="pulsar://pulsar:6650",
+    ):
+
+        if input_queue is None: input_queue = agent_request_queue
+        if output_queue is None: output_queue = agent_response_queue
+
+        super(AgentClient, self).__init__(
+            log_level=log_level,
+            subscriber=subscriber,
+            input_queue=input_queue,
+            output_queue=output_queue,
+            pulsar_host=pulsar_host,
+            input_schema=AgentRequest,
+            output_schema=AgentResponse,
+        )
+
+    def request(
+            self,
+            question,
+            think=None,
+            observe=None,
+            timeout=300
+    ):
+
+        def inspect(x):
+
+            if x.thought and think:
+                think(x.thought)
+                return
+
+            if x.observation and observe:
+                observe(x.observation)
+                return
+
+            if x.answer:
+                return True
+
+            return False
+
+        return self.call(
+            question=question, inspect=inspect, timeout=timeout
+        ).answer
+

trustgraph-base/trustgraph/clients/base.py

@@ -59,10 +59,14 @@ def __init__(
     def call(self, **args):
 
         timeout = args.get("timeout", DEFAULT_TIMEOUT)
+        inspect = args.get("inspect", lambda x: True)
 
         if "timeout" in args:
             del args["timeout"]
 
+        if "inspect" in args:
+            del args["inspect"]
+
         id = str(uuid.uuid4())
 
         r = self.input_schema(**args)
@@ -103,6 +107,10 @@ def call(self, **args):
                             f"{value.error.type}: {value.error.message}"
                         )
 
+                complete = inspect(value)
+
+                if not complete: continue
+
                 resp = msg.value()
                 self.consumer.acknowledge(msg)
                 return resp

trustgraph-base/trustgraph/schema/__init__.py

@@ -8,4 +8,5 @@
 from . graph import *
 from . retrieval import *
 from . metadata import *
+from . agent import *
 

trustgraph-base/trustgraph/schema/agent.py

@@ -0,0 +1,37 @@
+
+from pulsar.schema import Record, String, Array, Map
+
+from . topic import topic
+from . types import Error
+
+############################################################################
+
+# Prompt services, abstract the prompt generation
+
+class AgentStep(Record):
+    thought = String()
+    action = String()
+    arguments = Map(String())
+    observation = String()
+
+class AgentRequest(Record):
+    question = String()
+    plan = String()
+    state = String()
+    history = Array(AgentStep())
+
+class AgentResponse(Record):
+    answer = String()
+    error = Error()
+    thought = String()
+    observation = String()
+
+agent_request_queue = topic(
+    'agent', kind='non-persistent', namespace='request'
+)
+agent_response_queue = topic(
+    'agent', kind='non-persistent', namespace='response'
+)
+
+############################################################################
+

trustgraph-cli/scripts/tg-invoke-agent

@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+
+"""
+Uses the GraphRAG service to answer a query
+"""
+
+import argparse
+import os
+import textwrap
+
+from trustgraph.clients.agent_client import AgentClient
+
+default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_user = 'trustgraph'
+default_collection = 'default'
+
+def wrap(text, width=75):
+    if text is None: text = "n/a"
+    out = textwrap.wrap(
+        text, width=width
+    )
+    return "\n".join(out)
+
+def output(text, prefix="> ", width=78):
+    out = textwrap.indent(
+        text, prefix=prefix
+    )
+    print(out)
+
+def query(
+        pulsar_host, query, user, collection,
+        plan=None, state=None, verbose=False
+):
+
+    am = AgentClient(pulsar_host=pulsar_host)
+
+    if verbose:
+        output(wrap(query), "\U00002753 ")
+        print()
+
+    def think(x):
+        if verbose:
+            output(wrap(x), "\U0001f914 ")
+            print()
+
+    def observe(x):
+        if verbose:
+            output(wrap(x), "\U0001f4a1 ")
+            print()
+
+    resp = am.request(
+        question=query, think=think, observe=observe,
+    )
+
+    print(resp)
+
+def main():
+
+    parser = argparse.ArgumentParser(
+        prog='tg-invoke-agent',
+        description=__doc__,
+    )
+
+    parser.add_argument(
+        '-p', '--pulsar-host',
+        default=default_pulsar_host,
+        help=f'Pulsar host (default: {default_pulsar_host})',
+    )
+
+    parser.add_argument(
+        '-q', '--query',
+        required=True,
+        help=f'Query to execute',
+    )
+
+    parser.add_argument(
+        '-u', '--user',
+        default=default_user,
+        help=f'User ID (default: {default_user})'
+    )
+
+    parser.add_argument(
+        '-c', '--collection',
+        default=default_collection,
+        help=f'Collection ID (default: {default_collection})'
+    )
+
+    parser.add_argument(
+        '-l', '--plan',
+        help=f'Agent plan (default: unspecified)'
+    )
+
+    parser.add_argument(
+        '-s', '--state',
+        help=f'Agent initial state (default: unspecified)'
+    )
+
+    parser.add_argument(
+        '-v', '--verbose',
+        action="store_true",
+        help=f'Output thinking/observations'
+    )
+
+    args = parser.parse_args()
+
+    try:
+
+        query(
+            pulsar_host=args.pulsar_host,
+            query=args.query,
+            user=args.user,
+            collection=args.collection,
+            plan=args.plan,
+            state=args.state,
+            verbose=args.verbose,
+        )
+
+    except Exception as e:
+
+        print("Exception:", e, flush=True)
+
+main()
+

trustgraph-cli/setup.py

@@ -51,6 +51,7 @@
         "scripts/tg-query-graph-rag",
         "scripts/tg-init-pulsar",
         "scripts/tg-processor-state",
+        "scripts/tg-invoke-agent",
         "scripts/tg-invoke-prompt",
         "scripts/tg-invoke-llm",
     ]

trustgraph-flow/scripts/agent-manager-react

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+
+from trustgraph.agent.react import run
+
+run()
+

trustgraph-flow/setup.py

@@ -60,6 +60,7 @@
         "jsonschema",
     ],
     scripts=[
+        "scripts/agent-manager-react",
         "scripts/chunker-recursive",
         "scripts/chunker-token",
         "scripts/de-query-milvus",

trustgraph-flow/trustgraph/agent/react/README.md

@@ -0,0 +1,19 @@
+
+agent-manager-react \
+    -p pulsar://localhost:6650 \
+    --tool-type \
+        shuttle=knowledge-query:query \
+        cats=knowledge-query:query \
+        compute=text-completion:computation \
+    --tool-description \
+        shuttle="Query a knowledge base with information about the space shuttle.  The query should be a simple natural language question" \
+        cats="Query a knowledge base with information about Mark's cats.  The query should be a simple natural language question" \
+        compute="A computation engine which can answer questions about maths and computation" \
+    --tool-argument \
+        cats="query:string:The search query string" \
+        shuttle="query:string:The search query string" \
+        compute="computation:string:The computation to solve"
+
+
+    --context 'The space shuttle challenger final mission was 58-L'
+

trustgraph-flow/trustgraph/agent/react/__init__.py

@@ -0,0 +1,3 @@
+
+from . service import *
+

trustgraph-flow/trustgraph/agent/react/__main__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+from . service import run
+
+if __name__ == '__main__':
+    run()
+