diff --git a/trustgraph-base/trustgraph/base/prompt_client.py b/trustgraph-base/trustgraph/base/prompt_client.py index 0a98b580..568a25a3 100644 --- a/trustgraph-base/trustgraph/base/prompt_client.py +++ b/trustgraph-base/trustgraph/base/prompt_client.py @@ -6,7 +6,7 @@ from .. schema import PromptRequest, PromptResponse class PromptClient(RequestResponse): - async def prompt(self, id, variables, timeout=600): + async def prompt(self, id, variables, timeout=600, streaming=False): resp = await self.request( PromptRequest( @@ -14,7 +14,8 @@ class PromptClient(RequestResponse): terms = { k: json.dumps(v) for k, v in variables.items() - } + }, + streaming = streaming ), timeout=timeout ) @@ -70,11 +71,12 @@ class PromptClient(RequestResponse): timeout = timeout, ) - async def agent_react(self, variables, timeout=600): + async def agent_react(self, variables, timeout=600, streaming=False): return await self.prompt( id = "agent-react", variables = variables, timeout = timeout, + streaming = streaming, ) async def question(self, question, timeout=600): diff --git a/trustgraph-flow/trustgraph/agent/react/agent_manager.py b/trustgraph-flow/trustgraph/agent/react/agent_manager.py index 9b46bd34..70b7597b 100644 --- a/trustgraph-flow/trustgraph/agent/react/agent_manager.py +++ b/trustgraph-flow/trustgraph/agent/react/agent_manager.py @@ -169,7 +169,7 @@ class AgentManager: raise ValueError(f"Could not parse response: {text}") - async def reason(self, question, history, context): + async def reason(self, question, history, context, streaming=False): logger.debug(f"calling reason: {question}") @@ -220,7 +220,7 @@ class AgentManager: logger.info(f"prompt: {variables}") # Get text response from prompt service - response_text = await context("prompt-request").agent_react(variables) + response_text = await context("prompt-request").agent_react(variables, streaming=streaming) logger.debug(f"Response text:\n{response_text}") @@ -237,7 +237,7 @@ class AgentManager: logger.error(f"Response was: {response_text}") raise RuntimeError(f"Failed to parse agent response: {e}") - async def react(self, question, history, think, observe, context): + async def react(self, question, history, think, observe, context, streaming=False): logger.info(f"question: {question}") @@ -245,6 +245,7 @@ class AgentManager: question = question, history = history, context = context, + streaming = streaming, ) logger.info(f"act: {act}") diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py index f68ed4b2..8edb6a9b 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/trustgraph-flow/trustgraph/agent/react/service.py @@ -307,6 +307,7 @@ class Processor(AgentService): think = think, observe = observe, context = UserAwareContext(flow, request.user), + streaming = streaming, ) logger.debug(f"Action: {act}") diff --git a/trustgraph-flow/trustgraph/librarian/blob_store.py b/trustgraph-flow/trustgraph/librarian/blob_store.py index 78a8fdfd..e4ccfad9 100644 --- a/trustgraph-flow/trustgraph/librarian/blob_store.py +++ b/trustgraph-flow/trustgraph/librarian/blob_store.py @@ -34,9 +34,9 @@ class BlobStore: def ensure_bucket(self): # Make the bucket if it doesn't exist. - found = self.minio.bucket_exists(self.bucket_name) + found = self.minio.bucket_exists(bucket_name=self.bucket_name) if not found: - self.minio.make_bucket(self.bucket_name) + self.minio.make_bucket(bucket_name=self.bucket_name) logger.info(f"Created bucket {self.bucket_name}") else: logger.debug(f"Bucket {self.bucket_name} already exists")