trustgraph-flow 2.2.26__tar.gz → 2.3.0__tar.gz

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 2.2.26
+Version: 2.3.0
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Author-email: "trustgraph.ai" <security@trustgraph.ai>
 Project-URL: Homepage, https://github.com/trustgraph-ai/trustgraph
@@ -8,7 +8,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
-Requires-Dist: trustgraph-base<2.3,>=2.2
+Requires-Dist: trustgraph-base<2.4,>=2.3
 Requires-Dist: aiohttp
 Requires-Dist: anthropic
 Requires-Dist: scylla-driver

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/pyproject.toml

@@ -10,7 +10,7 @@ description = "TrustGraph provides a means to run a pipeline of flexible AI proc
 readme = "README.md"
 requires-python = ">=3.8"
 dependencies = [
-    "trustgraph-base>=2.2,<2.3",
+    "trustgraph-base>=2.3,<2.4",
     "aiohttp",
     "anthropic",
     "scylla-driver",

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/trustgraph/agent/orchestrator/meta_router.py

@@ -53,7 +53,7 @@ class MetaRouter:
                 "general": {"name": "general", "description": "General queries", "valid_patterns": ["react"], "framing": ""},
             }
 
-    async def identify_task_type(self, question, context):
+    async def identify_task_type(self, question, context, usage=None):
         """
         Use the LLM to classify the question into one of the known task types.
 
@@ -71,7 +71,7 @@ class MetaRouter:
 
         try:
             client = context("prompt-request")
-            response = await client.prompt(
+            result = await client.prompt(
                 id="task-type-classify",
                 variables={
                     "question": question,
@@ -81,7 +81,9 @@ class MetaRouter:
                     ],
                 },
             )
-            selected = response.strip().lower().replace('"', '').replace("'", "")
+            if usage:
+                usage.track(result)
+            selected = result.text.strip().lower().replace('"', '').replace("'", "")
 
             if selected in self.task_types:
                 framing = self.task_types[selected].get("framing", DEFAULT_FRAMING)
@@ -100,7 +102,7 @@ class MetaRouter:
         )
         return DEFAULT_TASK_TYPE, framing
 
-    async def select_pattern(self, question, task_type, context):
+    async def select_pattern(self, question, task_type, context, usage=None):
         """
         Use the LLM to select the best execution pattern for this task type.
 
@@ -120,7 +122,7 @@ class MetaRouter:
 
         try:
             client = context("prompt-request")
-            response = await client.prompt(
+            result = await client.prompt(
                 id="pattern-select",
                 variables={
                     "question": question,
@@ -133,7 +135,9 @@ class MetaRouter:
                     ],
                 },
             )
-            selected = response.strip().lower().replace('"', '').replace("'", "")
+            if usage:
+                usage.track(result)
+            selected = result.text.strip().lower().replace('"', '').replace("'", "")
 
             if selected in valid_patterns:
                 logger.info(f"MetaRouter: selected pattern '{selected}'")
@@ -148,19 +152,20 @@ class MetaRouter:
             logger.warning(f"MetaRouter: pattern selection failed: {e}")
             return valid_patterns[0] if valid_patterns else DEFAULT_PATTERN
 
-    async def route(self, question, context):
+    async def route(self, question, context, usage=None):
         """
         Full routing pipeline: identify task type, then select pattern.
 
         Args:
             question: The user's query.
             context: UserAwareContext (flow wrapper).
+            usage: Optional UsageTracker for token counting.
 
         Returns:
             (pattern, task_type, framing) tuple.
         """
-        task_type, framing = await self.identify_task_type(question, context)
-        pattern = await self.select_pattern(question, task_type, context)
+        task_type, framing = await self.identify_task_type(question, context, usage=usage)
+        pattern = await self.select_pattern(question, task_type, context, usage=usage)
         logger.info(
             f"MetaRouter: route result — "
             f"pattern={pattern}, task_type={task_type}, framing={framing!r}"

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/trustgraph/agent/orchestrator/pattern_base.py

@@ -65,6 +65,37 @@ class UserAwareContext:
         return client
 
 
+class UsageTracker:
+    """Accumulates token usage across multiple prompt calls."""
+
+    def __init__(self):
+        self.total_in = 0
+        self.total_out = 0
+        self.last_model = None
+
+    def track(self, result):
+        """Track usage from a PromptResult."""
+        if result is not None:
+            if getattr(result, "in_token", None) is not None:
+                self.total_in += result.in_token
+            if getattr(result, "out_token", None) is not None:
+                self.total_out += result.out_token
+            if getattr(result, "model", None) is not None:
+                self.last_model = result.model
+
+    @property
+    def in_token(self):
+        return self.total_in if self.total_in > 0 else None
+
+    @property
+    def out_token(self):
+        return self.total_out if self.total_out > 0 else None
+
+    @property
+    def model(self):
+        return self.last_model
+
+
 class PatternBase:
     """
     Shared infrastructure for all agent patterns.
@@ -571,7 +602,8 @@ class PatternBase:
     # ---- Response helpers ---------------------------------------------------
 
     async def prompt_as_answer(self, client, prompt_id, variables,
-                               respond, streaming, message_id=""):
+                               respond, streaming, message_id="",
+                               usage=None):
         """Call a prompt template, forwarding chunks as answer
         AgentResponse messages when streaming is enabled.
 
@@ -591,22 +623,28 @@ class PatternBase:
                         message_id=message_id,
                     ))
 
-            await client.prompt(
+            result = await client.prompt(
                 id=prompt_id,
                 variables=variables,
                 streaming=True,
                 chunk_callback=on_chunk,
             )
+            if usage:
+                usage.track(result)
 
             return "".join(accumulated)
         else:
-            return await client.prompt(
+            result = await client.prompt(
                 id=prompt_id,
                 variables=variables,
             )
+            if usage:
+                usage.track(result)
+            return result.text
 
     async def send_final_response(self, respond, streaming, answer_text,
-                                  already_streamed=False, message_id=""):
+                                  already_streamed=False, message_id="",
+                                  usage=None):
         """Send the answer content and end-of-dialog marker.
 
         Args:
@@ -614,7 +652,16 @@ class PatternBase:
                 via streaming callbacks (e.g. ReactPattern). Only the
                 end-of-dialog marker is emitted.
             message_id: Provenance URI for the answer entity.
+            usage: UsageTracker with accumulated token counts.
         """
+        usage_kwargs = {}
+        if usage:
+            usage_kwargs = {
+                "in_token": usage.in_token,
+                "out_token": usage.out_token,
+                "model": usage.model,
+            }
+
         if streaming and not already_streamed:
             # Answer wasn't streamed yet — send it as a chunk first
             if answer_text:
@@ -626,13 +673,14 @@ class PatternBase:
                     message_id=message_id,
                 ))
         if streaming:
-            # End-of-dialog marker
+            # End-of-dialog marker with usage
             await respond(AgentResponse(
                 chunk_type="answer",
                 content="",
                 end_of_message=True,
                 end_of_dialog=True,
                 message_id=message_id,
+                **usage_kwargs,
             ))
         else:
             await respond(AgentResponse(
@@ -641,6 +689,7 @@ class PatternBase:
                 end_of_message=True,
                 end_of_dialog=True,
                 message_id=message_id,
+                **usage_kwargs,
             ))
 
     def build_next_request(self, request, history, session_id, collection,

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/trustgraph/agent/orchestrator/plan_pattern.py

@@ -18,7 +18,7 @@ from trustgraph.provenance import (
     agent_synthesis_uri,
 )
 
-from . pattern_base import PatternBase
+from . pattern_base import PatternBase, UsageTracker
 
 logger = logging.getLogger(__name__)
 
@@ -35,7 +35,10 @@ class PlanThenExecutePattern(PatternBase):
     Subsequent calls execute the next pending plan step via ReACT.
     """
 
-    async def iterate(self, request, respond, next, flow):
+    async def iterate(self, request, respond, next, flow, usage=None):
+
+        if usage is None:
+            usage = UsageTracker()
 
         streaming = getattr(request, 'streaming', False)
         session_id = getattr(request, 'session_id', '') or str(uuid.uuid4())
@@ -67,13 +70,13 @@ class PlanThenExecutePattern(PatternBase):
             await self._planning_iteration(
                 request, respond, next, flow,
                 session_id, collection, streaming, session_uri,
-                iteration_num,
+                iteration_num, usage=usage,
             )
         else:
             await self._execution_iteration(
                 request, respond, next, flow,
                 session_id, collection, streaming, session_uri,
-                iteration_num, plan,
+                iteration_num, plan, usage=usage,
             )
 
     def _extract_plan(self, history):
@@ -98,7 +101,7 @@ class PlanThenExecutePattern(PatternBase):
 
     async def _planning_iteration(self, request, respond, next, flow,
                                   session_id, collection, streaming,
-                                  session_uri, iteration_num):
+                                  session_uri, iteration_num, usage=None):
         """Ask the LLM to produce a structured plan."""
 
         think = self.make_think_callback(respond, streaming)
@@ -113,7 +116,7 @@ class PlanThenExecutePattern(PatternBase):
         client = context("prompt-request")
 
         # Use the plan-create prompt template
-        plan_steps = await client.prompt(
+        result = await client.prompt(
             id="plan-create",
             variables={
                 "question": request.question,
@@ -124,7 +127,10 @@ class PlanThenExecutePattern(PatternBase):
                 ],
             },
         )
+        if usage:
+            usage.track(result)
 
+        plan_steps = result.objects
         # Validate we got a list
         if not isinstance(plan_steps, list) or not plan_steps:
             logger.warning("plan-create returned invalid result, falling back to single step")
@@ -187,7 +193,8 @@ class PlanThenExecutePattern(PatternBase):
 
     async def _execution_iteration(self, request, respond, next, flow,
                                    session_id, collection, streaming,
-                                   session_uri, iteration_num, plan):
+                                   session_uri, iteration_num, plan,
+                                   usage=None):
         """Execute the next pending plan step via single-shot tool call."""
 
         pending_idx = self._find_next_pending_step(plan)
@@ -198,6 +205,7 @@ class PlanThenExecutePattern(PatternBase):
                 request, respond, next, flow,
                 session_id, collection, streaming,
                 session_uri, iteration_num, plan,
+                usage=usage,
             )
             return
 
@@ -240,7 +248,7 @@ class PlanThenExecutePattern(PatternBase):
         client = context("prompt-request")
 
         # Single-shot: ask LLM which tool + arguments to use for this goal
-        tool_call = await client.prompt(
+        result = await client.prompt(
             id="plan-step-execute",
             variables={
                 "goal": goal,
@@ -258,7 +266,10 @@ class PlanThenExecutePattern(PatternBase):
                 ],
             },
         )
+        if usage:
+            usage.track(result)
 
+        tool_call = result.object
         tool_name = tool_call.get("tool", "")
         tool_arguments = tool_call.get("arguments", {})
 
@@ -330,7 +341,8 @@ class PlanThenExecutePattern(PatternBase):
 
     async def _synthesise(self, request, respond, next, flow,
                           session_id, collection, streaming,
-                          session_uri, iteration_num, plan):
+                          session_uri, iteration_num, plan,
+                          usage=None):
         """Synthesise a final answer from all completed plan step results."""
 
         think = self.make_think_callback(respond, streaming)
@@ -365,6 +377,7 @@ class PlanThenExecutePattern(PatternBase):
             respond=respond,
             streaming=streaming,
             message_id=synthesis_msg_id,
+            usage=usage,
         )
 
         # Emit synthesis provenance (links back to last step result)
@@ -380,4 +393,5 @@ class PlanThenExecutePattern(PatternBase):
             await self.send_final_response(
                 respond, streaming, response_text, already_streamed=streaming,
                 message_id=synthesis_msg_id,
+                usage=usage,
             )

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/trustgraph/agent/orchestrator/react_pattern.py

@@ -23,7 +23,7 @@ from ..react.agent_manager import AgentManager
 from ..react.types import Action, Final
 from ..tool_filter import get_next_state
 
-from . pattern_base import PatternBase
+from . pattern_base import PatternBase, UsageTracker
 
 logger = logging.getLogger(__name__)
 
@@ -37,7 +37,10 @@ class ReactPattern(PatternBase):
     result is appended to history and a next-request is emitted.
     """
 
-    async def iterate(self, request, respond, next, flow):
+    async def iterate(self, request, respond, next, flow, usage=None):
+
+        if usage is None:
+            usage = UsageTracker()
 
         streaming = getattr(request, 'streaming', False)
         session_id = getattr(request, 'session_id', '') or str(uuid.uuid4())
@@ -121,6 +124,7 @@ class ReactPattern(PatternBase):
             context=context,
             streaming=streaming,
             on_action=on_action,
+            usage=usage,
         )
 
         logger.debug(f"Action: {act}")
@@ -144,6 +148,7 @@ class ReactPattern(PatternBase):
                 await self.send_final_response(
                     respond, streaming, f, already_streamed=streaming,
                     message_id=answer_msg_id,
+                    usage=usage,
                 )
             return
 

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/trustgraph/agent/orchestrator/service.py

@@ -23,6 +23,7 @@ from ... base import Consumer, Producer
 from ... base import ConsumerMetrics, ProducerMetrics
 
 from ... schema import AgentRequest, AgentResponse, AgentStep, Error
+from ..orchestrator.pattern_base import UsageTracker
 from ... schema import Triples, Metadata
 from ... schema import LibrarianRequest, LibrarianResponse, DocumentMetadata
 from ... schema import librarian_request_queue, librarian_response_queue
@@ -493,6 +494,8 @@ class Processor(AgentService):
 
     async def agent_request(self, request, respond, next, flow):
 
+        usage = UsageTracker()
+
         try:
 
             # Intercept subagent completion messages
@@ -516,7 +519,7 @@ class Processor(AgentService):
 
                 if self.meta_router:
                     pattern, task_type, framing = await self.meta_router.route(
-                        request.question, context,
+                        request.question, context, usage=usage,
                     )
                 else:
                     pattern = "react"
@@ -536,16 +539,16 @@ class Processor(AgentService):
             # Dispatch to the selected pattern
             if pattern == "plan-then-execute":
                 await self.plan_pattern.iterate(
-                    request, respond, next, flow,
+                    request, respond, next, flow, usage=usage,
                 )
             elif pattern == "supervisor":
                 await self.supervisor_pattern.iterate(
-                    request, respond, next, flow,
+                    request, respond, next, flow, usage=usage,
                 )
             else:
                 # Default to react
                 await self.react_pattern.iterate(
-                    request, respond, next, flow,
+                    request, respond, next, flow, usage=usage,
                 )
 
         except Exception as e:

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/trustgraph/agent/orchestrator/supervisor_pattern.py

@@ -22,7 +22,7 @@ from trustgraph.provenance import (
     agent_synthesis_uri,
 )
 
-from . pattern_base import PatternBase
+from . pattern_base import PatternBase, UsageTracker
 
 logger = logging.getLogger(__name__)
 
@@ -38,7 +38,10 @@ class SupervisorPattern(PatternBase):
       - "synthesise": triggered by aggregator with results in subagent_results
     """
 
-    async def iterate(self, request, respond, next, flow):
+    async def iterate(self, request, respond, next, flow, usage=None):
+
+        if usage is None:
+            usage = UsageTracker()
 
         streaming = getattr(request, 'streaming', False)
         session_id = getattr(request, 'session_id', '') or str(uuid.uuid4())
@@ -72,17 +75,19 @@ class SupervisorPattern(PatternBase):
                 request, respond, next, flow,
                 session_id, collection, streaming,
                 session_uri, iteration_num,
+                usage=usage,
             )
         else:
             await self._decompose_and_fanout(
                 request, respond, next, flow,
                 session_id, collection, streaming,
                 session_uri, iteration_num,
+                usage=usage,
             )
 
     async def _decompose_and_fanout(self, request, respond, next, flow,
                                     session_id, collection, streaming,
-                                    session_uri, iteration_num):
+                                    session_uri, iteration_num, usage=None):
         """Decompose the question into sub-goals and fan out subagents."""
 
         decompose_msg_id = agent_decomposition_uri(session_id)
@@ -100,7 +105,7 @@ class SupervisorPattern(PatternBase):
         client = context("prompt-request")
 
         # Use the supervisor-decompose prompt template
-        goals = await client.prompt(
+        result = await client.prompt(
             id="supervisor-decompose",
             variables={
                 "question": request.question,
@@ -112,7 +117,10 @@ class SupervisorPattern(PatternBase):
                 ],
             },
         )
+        if usage:
+            usage.track(result)
 
+        goals = result.objects
         # Validate result
         if not isinstance(goals, list):
             goals = []
@@ -175,7 +183,7 @@ class SupervisorPattern(PatternBase):
 
     async def _synthesise(self, request, respond, next, flow,
                           session_id, collection, streaming,
-                          session_uri, iteration_num):
+                          session_uri, iteration_num, usage=None):
         """Synthesise final answer from subagent results."""
 
         synthesis_msg_id = agent_synthesis_uri(session_id)
@@ -216,6 +224,7 @@ class SupervisorPattern(PatternBase):
             respond=respond,
             streaming=streaming,
             message_id=synthesis_msg_id,
+            usage=usage,
         )
 
         # Emit synthesis provenance (links back to all findings)
@@ -231,4 +240,5 @@ class SupervisorPattern(PatternBase):
         await self.send_final_response(
             respond, streaming, response_text, already_streamed=streaming,
             message_id=synthesis_msg_id,
+            usage=usage,
         )

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/trustgraph/agent/react/agent_manager.py

@@ -170,7 +170,7 @@ class AgentManager:
 
         raise ValueError(f"Could not parse response: {text}")
 
-    async def reason(self, question, history, context, streaming=False, think=None, observe=None, answer=None):
+    async def reason(self, question, history, context, streaming=False, think=None, observe=None, answer=None, usage=None):
 
         logger.debug(f"calling reason: {question}")
 
@@ -255,11 +255,13 @@ class AgentManager:
             client = context("prompt-request")
 
             # Get streaming response
-            response_text = await client.agent_react(
+            prompt_result = await client.agent_react(
                 variables=variables,
                 streaming=True,
                 chunk_callback=on_chunk
             )
+            if usage:
+                usage.track(prompt_result)
 
             # Finalize parser
             parser.finalize()
@@ -275,10 +277,13 @@ class AgentManager:
             # Non-streaming path - get complete text and parse
             client = context("prompt-request")
 
-            response_text = await client.agent_react(
+            prompt_result = await client.agent_react(
                 variables=variables,
                 streaming=False
             )
+            if usage:
+                usage.track(prompt_result)
+            response_text = prompt_result.text
 
             logger.debug(f"Response text:\n{response_text}")
 
@@ -292,7 +297,8 @@ class AgentManager:
                 raise RuntimeError(f"Failed to parse agent response: {e}")
 
     async def react(self, question, history, think, observe, context,
-                    streaming=False, answer=None, on_action=None):
+                    streaming=False, answer=None, on_action=None,
+                    usage=None):
 
         act = await self.reason(
             question = question,
@@ -302,6 +308,7 @@ class AgentManager:
             think = think,
             observe = observe,
             answer = answer,
+            usage = usage,
         )
 
         if isinstance(act, Final):

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/trustgraph/agent/react/tools.py

@@ -78,9 +78,10 @@ class TextCompletionImpl:
     async def invoke(self, **arguments):
         client = self.context("prompt-request")
         logger.debug("Prompt question...")
-        return await client.question(
+        result = await client.question(
             arguments.get("question")
         )
+        return result.text
 
 # This tool implementation knows how to do MCP tool invocation.  This uses
 # the mcp-tool service.
@@ -227,10 +228,11 @@ class PromptImpl:
     async def invoke(self, **arguments):
         client = self.context("prompt-request")
         logger.debug(f"Prompt template invocation: {self.template_id}...")
-        return await client.prompt(
+        result = await client.prompt(
             id=self.template_id,
             variables=arguments
         )
+        return result.text
 
 
 # This tool implementation invokes a dynamically configured tool service

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/trustgraph/extract/kg/definitions/extract.py

@@ -117,10 +117,11 @@ class Processor(FlowProcessor):
 
             try:
 
-                defs = await flow("prompt-request").extract_definitions(
+                result = await flow("prompt-request").extract_definitions(
                     text = chunk
                 )
 
+                defs = result.objects
                 logger.debug(f"Definitions response: {defs}")
 
                 if type(defs) != list:

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/trustgraph/extract/kg/ontology/extract.py

@@ -376,10 +376,11 @@ class Processor(FlowProcessor):
         """
         try:
             # Call prompt service with simplified format prompt
-            extraction_response = await flow("prompt-request").prompt(
+            result = await flow("prompt-request").prompt(
                 id="extract-with-ontologies",
                 variables=prompt_variables
             )
+            extraction_response = result.object
             logger.debug(f"Simplified extraction response: {extraction_response}")
 
             # Parse response into structured format

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/trustgraph/extract/kg/relationships/extract.py

@@ -100,10 +100,11 @@ class Processor(FlowProcessor):
 
             try:
 
-                rels = await flow("prompt-request").extract_relationships(
+                result = await flow("prompt-request").extract_relationships(
                     text = chunk
                 )
 
+                rels = result.objects
                 logger.debug(f"Prompt response: {rels}")
 
                 if type(rels) != list:

{trustgraph_flow-2.2.26 → trustgraph_flow-2.3.0}/trustgraph/extract/kg/rows/processor.py

@@ -148,11 +148,12 @@ class Processor(FlowProcessor):
             schema_dict = row_schema_translator.encode(schema)
             
             # Use prompt client to extract rows based on schema
-            objects = await flow("prompt-request").extract_objects(
+            result = await flow("prompt-request").extract_objects(
                 schema=schema_dict,
                 text=text
             )
-            
+
+            objects = result.objects
             if not isinstance(objects, list):
                 return []
             

trustgraph_flow-2.3.0/trustgraph/flow_version.py

@@ -0,0 +1 @@
+__version__ = "2.3.0"