diff --git a/tests/test-agent b/tests/test-agent index 4782bbae..b1420098 100755 --- a/tests/test-agent +++ b/tests/test-agent @@ -20,7 +20,11 @@ def output(text, prefix="> ", width=78): ) print(out) -p = AgentClient(pulsar_host="pulsar://localhost:6650") +p = AgentClient( + pulsar_host="pulsar://pulsar:6650", + input_queue = "non-persistent://tg/request/agent:0000", + output_queue = "non-persistent://tg/response/agent:0000", +) 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." diff --git a/tests/test-graph-rag b/tests/test-graph-rag index ac407f42..b62f890c 100755 --- a/tests/test-graph-rag +++ b/tests/test-graph-rag @@ -10,9 +10,11 @@ rag = GraphRagClient( output_queue = "non-persistent://tg/response/graph-rag:default", ) -query=""" -This knowledge graph describes the Space Shuttle disaster. -Present 20 facts which are present in the knowledge graph.""" +#query=""" +#This knowledge graph describes the Space Shuttle disaster. +#Present 20 facts which are present in the knowledge graph.""" + +query = "How many cats does Mark have?" resp = rag.request(query) diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index ac3bb766..2accbb21 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/trustgraph-base/trustgraph/base/__init__.py @@ -26,4 +26,6 @@ from . document_embeddings_query_service import DocumentEmbeddingsQueryService from . graph_embeddings_client import GraphEmbeddingsClientSpec from . triples_client import TriplesClientSpec from . document_embeddings_client import DocumentEmbeddingsClientSpec +from . agent_service import AgentService +from . graph_rag_client import GraphRagClientSpec diff --git a/trustgraph-base/trustgraph/base/agent_client.py b/trustgraph-base/trustgraph/base/agent_client.py new file mode 100644 index 00000000..76e1adff --- /dev/null +++ b/trustgraph-base/trustgraph/base/agent_client.py @@ -0,0 +1,39 @@ + +from . request_response_spec import RequestResponse, RequestResponseSpec +from .. schema import AgentRequest, AgentResponse +from .. knowledge import Uri, Literal + +class AgentClient(RequestResponse): + async def request(self, recipient, question, plan=None, state=None, + history=[], timeout=300): + + resp = await self.request( + AgentRequest( + question = question, + plan = plan, + state = state, + history = history, + ), + recipient=recipient, + timeout=timeout, + ) + + print(resp, flush=True) + + if resp.error: + raise RuntimeError(resp.error.message) + + return resp + +class GraphEmbeddingsClientSpec(RequestResponseSpec): + def __init__( + self, request_name, response_name, + ): + super(GraphEmbeddingsClientSpec, self).__init__( + request_name = request_name, + request_schema = GraphEmbeddingsRequest, + response_name = response_name, + response_schema = GraphEmbeddingsResponse, + impl = GraphEmbeddingsClient, + ) + diff --git a/trustgraph-base/trustgraph/base/agent_service.py b/trustgraph-base/trustgraph/base/agent_service.py new file mode 100644 index 00000000..0dbe728e --- /dev/null +++ b/trustgraph-base/trustgraph/base/agent_service.py @@ -0,0 +1,100 @@ + +""" +Agent manager service completion base class +""" + +import time +from prometheus_client import Histogram + +from .. schema import AgentRequest, AgentResponse, Error +from .. exceptions import TooManyRequests +from .. base import FlowProcessor, ConsumerSpec, ProducerSpec + +default_ident = "agent-manager" + +class AgentService(FlowProcessor): + + def __init__(self, **params): + + id = params.get("id") + + super(AgentService, self).__init__(**params | { "id": id }) + + self.register_specification( + ConsumerSpec( + name = "request", + schema = AgentRequest, + handler = self.on_request + ) + ) + + self.register_specification( + ProducerSpec( + name = "next", + schema = AgentRequest + ) + ) + + self.register_specification( + ProducerSpec( + name = "response", + schema = AgentResponse + ) + ) + + async def on_request(self, msg, consumer, flow): + + try: + + request = msg.value() + + # Sender-produced ID + id = msg.properties()["id"] + + async def respond(resp): + + await flow("response").send( + resp, + properties={"id": id} + ) + + async def next(resp): + + await flow("next").send( + resp, + properties={"id": id} + ) + + await self.agent_request( + request = request, respond = respond, next = next, + flow = flow + ) + + except TooManyRequests as e: + raise e + + except Exception as e: + + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"on_request Exception: {e}") + + print("Send error response...", flush=True) + + await flow.producer["response"].send( + AgentResponse( + error=Error( + type = "agent-error", + message = str(e), + ), + thought = None, + observation = None, + answer = None, + ), + properties={"id": id} + ) + + @staticmethod + def add_args(parser): + + FlowProcessor.add_args(parser) + diff --git a/trustgraph-base/trustgraph/base/consumer.py b/trustgraph-base/trustgraph/base/consumer.py index ddedd93a..57b940ac 100644 --- a/trustgraph-base/trustgraph/base/consumer.py +++ b/trustgraph-base/trustgraph/base/consumer.py @@ -89,7 +89,7 @@ class Consumer: except Exception as e: - print("Exception:", e, flush=True) + print("consumer subs Exception:", e, flush=True) await asyncio.sleep(self.reconnect_time) continue @@ -107,7 +107,7 @@ class Consumer: except Exception as e: - print("Exception:", e, flush=True) + print("consumer loop exception:", e, flush=True) self.consumer.close() self.consumer = None await asyncio.sleep(self.reconnect_time) @@ -184,7 +184,7 @@ class Consumer: except Exception as e: - print("Exception:", e, flush=True) + print("consume exception:", e, flush=True) # Message failed to be processed, this causes it to # be retried diff --git a/trustgraph-base/trustgraph/base/document_embeddings_query_service.py b/trustgraph-base/trustgraph/base/document_embeddings_query_service.py old mode 100755 new mode 100644 diff --git a/trustgraph-base/trustgraph/base/document_embeddings_store_service.py b/trustgraph-base/trustgraph/base/document_embeddings_store_service.py old mode 100755 new mode 100644 diff --git a/trustgraph-base/trustgraph/base/embeddings_service.py b/trustgraph-base/trustgraph/base/embeddings_service.py old mode 100755 new mode 100644 diff --git a/trustgraph-base/trustgraph/base/flow_processor.py b/trustgraph-base/trustgraph/base/flow_processor.py index d96c9345..e6460fe3 100644 --- a/trustgraph-base/trustgraph/base/flow_processor.py +++ b/trustgraph-base/trustgraph/base/flow_processor.py @@ -67,26 +67,27 @@ class FlowProcessor(AsyncProcessor): # Get my flow config flow_config = json.loads(config["flows-active"][self.id]) - # Get list of flows which should be running and are currently - # running - wanted_flows = flow_config.keys() - current_flows = self.flows.keys() - - # Start all the flows which arent currently running - for flow in wanted_flows: - if flow not in current_flows: - await self.start_flow(flow, flow_config[flow]) - - # Stop all the unwanted flows which are due to be stopped - for flow in current_flows: - if flow not in wanted_flows: - await self.stop_flow(flow) - - print("Handled config update") - else: - print("No configuration settings for me!", flush=True) + print("No configuration settings for me.", flush=True) + flow_config = {} + + # Get list of flows which should be running and are currently + # running + wanted_flows = flow_config.keys() + current_flows = self.flows.keys() + + # Start all the flows which arent currently running + for flow in wanted_flows: + if flow not in current_flows: + await self.start_flow(flow, flow_config[flow]) + + # Stop all the unwanted flows which are due to be stopped + for flow in current_flows: + if flow not in wanted_flows: + await self.stop_flow(flow) + + print("Handled config update") # Start threads, just call parent async def start(self): diff --git a/trustgraph-base/trustgraph/base/graph_embeddings_query_service.py b/trustgraph-base/trustgraph/base/graph_embeddings_query_service.py old mode 100755 new mode 100644 diff --git a/trustgraph-base/trustgraph/base/graph_embeddings_store_service.py b/trustgraph-base/trustgraph/base/graph_embeddings_store_service.py old mode 100755 new mode 100644 diff --git a/trustgraph-base/trustgraph/base/graph_rag_client.py b/trustgraph-base/trustgraph/base/graph_rag_client.py new file mode 100644 index 00000000..c4f3f7ab --- /dev/null +++ b/trustgraph-base/trustgraph/base/graph_rag_client.py @@ -0,0 +1,33 @@ + +from . request_response_spec import RequestResponse, RequestResponseSpec +from .. schema import GraphRagQuery, GraphRagResponse + +class GraphRagClient(RequestResponse): + async def rag(self, query, user="trustgraph", collection="default", + timeout=600): + resp = await self.request( + GraphRagQuery( + query = query, + user = user, + collection = collection, + ), + timeout=timeout + ) + + if resp.error: + raise RuntimeError(resp.error.message) + + return resp.response + +class GraphRagClientSpec(RequestResponseSpec): + def __init__( + self, request_name, response_name, + ): + super(GraphRagClientSpec, self).__init__( + request_name = request_name, + request_schema = GraphRagQuery, + response_name = response_name, + response_schema = GraphRagResponse, + impl = GraphRagClient, + ) + diff --git a/trustgraph-base/trustgraph/base/llm_service.py b/trustgraph-base/trustgraph/base/llm_service.py old mode 100755 new mode 100644 index d22fe564..39323db7 --- a/trustgraph-base/trustgraph/base/llm_service.py +++ b/trustgraph-base/trustgraph/base/llm_service.py @@ -62,8 +62,6 @@ class LlmService(FlowProcessor): id = msg.properties()["id"] - prompt = request.system + "\n\n" + request.prompt - with __class__.text_completion_metric.labels( id=self.id, flow=f"{flow.name}-{consumer.name}", diff --git a/trustgraph-base/trustgraph/base/prompt_client.py b/trustgraph-base/trustgraph/base/prompt_client.py index 88e1a15f..9e8ab033 100644 --- a/trustgraph-base/trustgraph/base/prompt_client.py +++ b/trustgraph-base/trustgraph/base/prompt_client.py @@ -63,6 +63,22 @@ class PromptClient(RequestResponse): timeout = timeout, ) + async def agent_react(self, variables, timeout=600): + return await self.prompt( + id = "agent-react", + variables = variables, + timeout = timeout, + ) + + async def question(self, question, timeout=600): + return await self.prompt( + id = "question", + variables = { + "question": question, + }, + timeout = timeout, + ) + class PromptClientSpec(RequestResponseSpec): def __init__( self, request_name, response_name, diff --git a/trustgraph-base/trustgraph/base/request_response_spec.py b/trustgraph-base/trustgraph/base/request_response_spec.py index 8c881550..dcfcbf9b 100644 --- a/trustgraph-base/trustgraph/base/request_response_spec.py +++ b/trustgraph-base/trustgraph/base/request_response_spec.py @@ -40,7 +40,7 @@ class RequestResponse(Subscriber): await self.producer.stop() await super(RequestResponse, self).stop() - async def request(self, req, timeout=300): + async def request(self, req, timeout=300, recipient=None): id = str(uuid.uuid4()) @@ -55,14 +55,40 @@ class RequestResponse(Subscriber): properties={"id": id} ) - resp = await asyncio.wait_for( - q.get(), - timeout=timeout - ) + except Exception as e: - print("Got response.", flush=True) + print("Exception:", e) + raise e - return resp + + try: + + while True: + + resp = await asyncio.wait_for( + q.get(), + timeout=timeout + ) + + print("Got response.", flush=True) + + if recipient is None: + + # If no recipient handler, just return the first + # response we get + return resp + else: + + # Recipient handler gets to decide when we're done b + # returning a boolean + fin = await recipient(resp) + + # If done, return the last result otherwise loop round for + # next response + if fin: + return resp + else: + continue except Exception as e: @@ -108,4 +134,3 @@ class RequestResponseSpec(Spec): flow.consumer[self.request_name] = rr - diff --git a/trustgraph-base/trustgraph/base/triples_query_service.py b/trustgraph-base/trustgraph/base/triples_query_service.py old mode 100755 new mode 100644 diff --git a/trustgraph-base/trustgraph/base/triples_store_service.py b/trustgraph-base/trustgraph/base/triples_store_service.py old mode 100755 new mode 100644 diff --git a/trustgraph-base/trustgraph/schema/agent.py b/trustgraph-base/trustgraph/schema/agent.py index 9bcdde51..ee20a9aa 100644 --- a/trustgraph-base/trustgraph/schema/agent.py +++ b/trustgraph-base/trustgraph/schema/agent.py @@ -26,12 +26,5 @@ class AgentResponse(Record): thought = String() observation = String() -agent_request_queue = topic( - 'agent', kind='non-persistent', namespace='request' -) -agent_response_queue = topic( - 'agent', kind='non-persistent', namespace='response' -) - ############################################################################ diff --git a/trustgraph-flow/trustgraph/agent/react/agent_manager.py b/trustgraph-flow/trustgraph/agent/react/agent_manager.py index a195bd80..d20b86f7 100644 --- a/trustgraph-flow/trustgraph/agent/react/agent_manager.py +++ b/trustgraph-flow/trustgraph/agent/react/agent_manager.py @@ -8,12 +8,11 @@ logger = logging.getLogger(__name__) class AgentManager: - def __init__(self, context, tools, additional_context=None): - self.context = context + def __init__(self, tools, additional_context=None): self.tools = tools self.additional_context = additional_context - def reason(self, question, history): + async def reason(self, question, history, context): tools = self.tools @@ -56,10 +55,7 @@ class AgentManager: logger.info(f"prompt: {variables}") - obj = self.context.prompt.request( - "agent-react", - variables - ) + obj = await context("prompt-request").agent_react(variables) print(json.dumps(obj, indent=4), flush=True) @@ -85,9 +81,13 @@ class AgentManager: return a - async def react(self, question, history, think, observe): + async def react(self, question, history, think, observe, context): - act = self.reason(question, history) + act = await self.reason( + question = question, + history = history, + context = context, + ) logger.info(f"act: {act}") if isinstance(act, Final): @@ -104,7 +104,12 @@ class AgentManager: else: raise RuntimeError(f"No action for {act.name}!") - resp = action.implementation.invoke(**act.arguments) + print("TOOL>>>", act) + resp = await action.implementation(context).invoke( + **act.arguments + ) + + print("RSETUL", resp) resp = resp.strip() diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py index 5b5e97b5..beb17fd4 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/trustgraph-flow/trustgraph/agent/react/service.py @@ -6,105 +6,68 @@ import json import re import sys -from pulsar.schema import JsonSchema +from ... base import AgentService, TextCompletionClientSpec, PromptClientSpec +from ... base import GraphRagClientSpec -from ... base import ConsumerProducer -from ... schema import Error -from ... schema import AgentRequest, AgentResponse, AgentStep -from ... schema import agent_request_queue, agent_response_queue -from ... schema import prompt_request_queue as pr_request_queue -from ... schema import prompt_response_queue as pr_response_queue -from ... schema import graph_rag_request_queue as gr_request_queue -from ... schema import graph_rag_response_queue as gr_response_queue -from ... clients.prompt_client import PromptClient -from ... clients.llm_client import LlmClient -from ... clients.graph_rag_client import GraphRagClient +from ... schema import AgentRequest, AgentResponse, AgentStep, Error from . tools import KnowledgeQueryImpl, TextCompletionImpl from . agent_manager import AgentManager from . types import Final, Action, Tool, Argument -module = "agent" +default_ident = "agent-manager" +default_max_iterations = 10 -default_input_queue = agent_request_queue -default_output_queue = agent_response_queue -default_subscriber = module -default_max_iterations = 15 - -class Processor(ConsumerProducer): +class Processor(AgentService): def __init__(self, **params): + id = params.get("id") + self.max_iterations = int( params.get("max_iterations", default_max_iterations) ) - tools = {} - - input_queue = params.get("input_queue", default_input_queue) - output_queue = params.get("output_queue", default_output_queue) - subscriber = params.get("subscriber", default_subscriber) - prompt_request_queue = params.get( - "prompt_request_queue", pr_request_queue - ) - prompt_response_queue = params.get( - "prompt_response_queue", pr_response_queue - ) - graph_rag_request_queue = params.get( - "graph_rag_request_queue", gr_request_queue - ) - graph_rag_response_queue = params.get( - "graph_rag_response_queue", gr_response_queue - ) - self.config_key = params.get("config_type", "agent") super(Processor, self).__init__( **params | { - "input_queue": input_queue, - "output_queue": output_queue, - "subscriber": subscriber, - "input_schema": AgentRequest, - "output_schema": AgentResponse, - "prompt_request_queue": prompt_request_queue, - "prompt_response_queue": prompt_response_queue, - "graph_rag_request_queue": gr_request_queue, - "graph_rag_response_queue": gr_response_queue, + "id": id, + "max_iterations": self.max_iterations, + "config_type": self.config_key, } ) - self.prompt = PromptClient( - subscriber=subscriber, - input_queue=prompt_request_queue, - output_queue=prompt_response_queue, - pulsar_host = self.pulsar_host, - pulsar_api_key=self.pulsar_api_key, - ) - - self.graph_rag = GraphRagClient( - subscriber=subscriber, - input_queue=graph_rag_request_queue, - output_queue=graph_rag_response_queue, - pulsar_host = self.pulsar_host, - pulsar_api_key=self.pulsar_api_key, - ) - - # Need to be able to feed requests to myself - self.recursive_input = self.client.create_producer( - topic=input_queue, - schema=JsonSchema(AgentRequest), - ) - self.agent = AgentManager( - context=self, tools=[], additional_context="", ) - self.config_handlers.append(self.on_config) + self.config_handlers.append(self.on_tools_config) - async def on_config(self, version, config): + self.register_specification( + TextCompletionClientSpec( + request_name = "text-completion-request", + response_name = "text-completion-response", + ) + ) + + self.register_specification( + GraphRagClientSpec( + request_name = "graph-rag-request", + response_name = "graph-rag-response", + ) + ) + + self.register_specification( + PromptClientSpec( + request_name = "prompt-request", + response_name = "prompt-response", + ) + ) + + async def on_tools_config(self, config, version): print("Loading configuration version", version) @@ -140,9 +103,9 @@ class Processor(ConsumerProducer): impl_id = data.get("type") if impl_id == "knowledge-query": - impl = KnowledgeQueryImpl(self) + impl = KnowledgeQueryImpl elif impl_id == "text-completion": - impl = TextCompletionImpl(self) + impl = TextCompletionImpl else: raise RuntimeError( f"Tool-kind {impl_id} not known" @@ -157,7 +120,6 @@ class Processor(ConsumerProducer): ) self.agent = AgentManager( - context=self, tools=tools, additional_context=additional ) @@ -166,19 +128,14 @@ class Processor(ConsumerProducer): except Exception as e: - print("Exception:", e, flush=True) + print("on_tools_config Exception:", e, flush=True) print("Configuration reload failed", flush=True) - async def handle(self, msg): + async def agent_request(self, request, respond, next, flow): try: - v = msg.value() - - # Sender-produced ID - id = msg.properties()["id"] - - if v.history: + if request.history: history = [ Action( thought=h.thought, @@ -186,12 +143,12 @@ class Processor(ConsumerProducer): arguments=h.arguments, observation=h.observation ) - for h in v.history + for h in request.history ] else: history = [] - print(f"Question: {v.question}", flush=True) + print(f"Question: {request.question}", flush=True) if len(history) >= self.max_iterations: raise RuntimeError("Too many agent iterations") @@ -209,7 +166,7 @@ class Processor(ConsumerProducer): observation=None, ) - await self.send(r, properties={"id": id}) + await respond(r) async def observe(x): @@ -222,15 +179,21 @@ class Processor(ConsumerProducer): observation=x, ) - await self.send(r, properties={"id": id}) + await respond(r) - act = await self.agent.react(v.question, history, think, observe) + act = await self.agent.react( + question = request.question, + history = history, + think = think, + observe = observe, + context = flow, + ) print(f"Action: {act}", flush=True) - print("Send response...", flush=True) + if isinstance(act, Final): - if type(act) == Final: + print("Send final response...", flush=True) r = AgentResponse( answer=act.final, @@ -238,18 +201,20 @@ class Processor(ConsumerProducer): thought=None, ) - await self.send(r, properties={"id": id}) + await respond(r) print("Done.", flush=True) return + print("Send next...", flush=True) + history.append(act) r = AgentRequest( - question=v.question, - plan=v.plan, - state=v.state, + question=request.question, + plan=request.plan, + state=request.state, history=[ AgentStep( thought=h.thought, @@ -261,7 +226,7 @@ class Processor(ConsumerProducer): ] ) - self.recursive_input.send(r, properties={"id": id}) + await next(r) print("Done.", flush=True) @@ -269,7 +234,7 @@ class Processor(ConsumerProducer): except Exception as e: - print(f"Exception: {e}") + print(f"agent_request Exception: {e}") print("Send error response...", flush=True) @@ -281,39 +246,12 @@ class Processor(ConsumerProducer): response=None, ) - await self.send(r, properties={"id": id}) + await respond(r) @staticmethod def add_args(parser): - ConsumerProducer.add_args( - parser, default_input_queue, default_subscriber, - default_output_queue, - ) - - parser.add_argument( - '--prompt-request-queue', - default=pr_request_queue, - help=f'Prompt request queue (default: {pr_request_queue})', - ) - - parser.add_argument( - '--prompt-response-queue', - default=pr_response_queue, - help=f'Prompt response queue (default: {pr_response_queue})', - ) - - parser.add_argument( - '--graph-rag-request-queue', - default=gr_request_queue, - help=f'Graph RAG request queue (default: {gr_request_queue})', - ) - - parser.add_argument( - '--graph-rag-response-queue', - default=gr_response_queue, - help=f'Graph RAG response queue (default: {gr_response_queue})', - ) + AgentService.add_args(parser) parser.add_argument( '--max-iterations', @@ -329,5 +267,5 @@ class Processor(ConsumerProducer): def run(): - Processor.launch(module, __doc__) + Processor.launch(default_ident, __doc__) diff --git a/trustgraph-flow/trustgraph/agent/react/tools.py b/trustgraph-flow/trustgraph/agent/react/tools.py index 023abc02..31568b25 100644 --- a/trustgraph-flow/trustgraph/agent/react/tools.py +++ b/trustgraph-flow/trustgraph/agent/react/tools.py @@ -4,16 +4,22 @@ class KnowledgeQueryImpl: def __init__(self, context): self.context = context - def invoke(self, **arguments): - return self.context.graph_rag.request(arguments.get("question")) + async def invoke(self, **arguments): + client = self.context("graph-rag-request") + print("Graph RAG question...", flush=True) + return await client.rag( + arguments.get("question") + ) # This tool implementation knows how to do text completion. This uses # the prompt service, rather than talking to TextCompletion directly. class TextCompletionImpl: def __init__(self, context): self.context = context - def invoke(self, **arguments): - return self.context.prompt.request( - "question", { "question": arguments.get("question") } + async def invoke(self, **arguments): + client = self.context("prompt-request") + print("Prompt question...", flush=True) + return await client.question( + arguments.get("question") ) diff --git a/trustgraph-flow/trustgraph/gateway/agent.py b/trustgraph-flow/trustgraph/gateway/agent.py index 150b970e..5a54931b 100644 --- a/trustgraph-flow/trustgraph/gateway/agent.py +++ b/trustgraph-flow/trustgraph/gateway/agent.py @@ -39,4 +39,3 @@ class AgentRequestor(ServiceRequestor): # The 2nd boolean expression indicates whether we're done responding return resp, (message.answer is not None) -