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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 2.2.26
+Version: 2.3.1
 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.1}/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.1}/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.1}/trustgraph/agent/orchestrator/pattern_base.py

@@ -25,6 +25,7 @@ from trustgraph.provenance import (
     agent_plan_uri,
     agent_step_result_uri,
     agent_synthesis_uri,
+    agent_pattern_decision_uri,
     agent_session_triples,
     agent_iteration_triples,
     agent_observation_triples,
@@ -34,6 +35,7 @@ from trustgraph.provenance import (
     agent_plan_triples,
     agent_step_result_triples,
     agent_synthesis_triples,
+    agent_pattern_decision_triples,
     set_graph,
     GRAPH_RETRIEVAL,
 )
@@ -65,6 +67,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.
@@ -151,7 +184,7 @@ class PatternBase:
             logger.debug(f"Think: {x} (is_final={is_final})")
             if streaming:
                 r = AgentResponse(
-                    chunk_type="thought",
+                    message_type="thought",
                     content=x,
                     end_of_message=is_final,
                     end_of_dialog=False,
@@ -159,7 +192,7 @@ class PatternBase:
                 )
             else:
                 r = AgentResponse(
-                    chunk_type="thought",
+                    message_type="thought",
                     content=x,
                     end_of_message=True,
                     end_of_dialog=False,
@@ -174,7 +207,7 @@ class PatternBase:
             logger.debug(f"Observe: {x} (is_final={is_final})")
             if streaming:
                 r = AgentResponse(
-                    chunk_type="observation",
+                    message_type="observation",
                     content=x,
                     end_of_message=is_final,
                     end_of_dialog=False,
@@ -182,7 +215,7 @@ class PatternBase:
                 )
             else:
                 r = AgentResponse(
-                    chunk_type="observation",
+                    message_type="observation",
                     content=x,
                     end_of_message=True,
                     end_of_dialog=False,
@@ -197,7 +230,7 @@ class PatternBase:
             logger.debug(f"Answer: {x}")
             if streaming:
                 r = AgentResponse(
-                    chunk_type="answer",
+                    message_type="answer",
                     content=x,
                     end_of_message=False,
                     end_of_dialog=False,
@@ -205,7 +238,7 @@ class PatternBase:
                 )
             else:
                 r = AgentResponse(
-                    chunk_type="answer",
+                    message_type="answer",
                     content=x,
                     end_of_message=True,
                     end_of_dialog=False,
@@ -239,16 +272,43 @@ class PatternBase:
         logger.debug(f"Emitted session triples for {session_uri}")
 
         await respond(AgentResponse(
-            chunk_type="explain",
+            message_type="explain",
             content="",
             explain_id=session_uri,
             explain_graph=GRAPH_RETRIEVAL,
             explain_triples=triples,
         ))
 
+    async def emit_pattern_decision_triples(
+        self, flow, session_id, session_uri, pattern, task_type,
+        user, collection, respond,
+    ):
+        """Emit provenance triples for a meta-router pattern decision."""
+        uri = agent_pattern_decision_uri(session_id)
+        triples = set_graph(
+            agent_pattern_decision_triples(
+                uri, session_uri, pattern, task_type,
+            ),
+            GRAPH_RETRIEVAL,
+        )
+        await flow("explainability").send(Triples(
+            metadata=Metadata(id=uri, user=user, collection=collection),
+            triples=triples,
+        ))
+        await respond(AgentResponse(
+            message_type="explain", content="",
+            explain_id=uri, explain_graph=GRAPH_RETRIEVAL,
+            explain_triples=triples,
+        ))
+        return uri
+
     async def emit_iteration_triples(self, flow, session_id, iteration_num,
                                      session_uri, act, request, respond,
-                                     streaming):
+                                     streaming, tool_candidates=None,
+                                     step_number=None,
+                                     llm_duration_ms=None,
+                                     in_token=None, out_token=None,
+                                     model=None):
         """Emit provenance triples for an iteration (Analysis+ToolUse)."""
         iteration_uri = agent_iteration_uri(session_id, iteration_num)
 
@@ -288,6 +348,12 @@ class PatternBase:
                 arguments=act.arguments,
                 thought_uri=thought_entity_uri if thought_doc_id else None,
                 thought_document_id=thought_doc_id,
+                tool_candidates=tool_candidates,
+                step_number=step_number,
+                llm_duration_ms=llm_duration_ms,
+                in_token=in_token,
+                out_token=out_token,
+                model=model,
             ),
             GRAPH_RETRIEVAL,
         )
@@ -302,7 +368,7 @@ class PatternBase:
         logger.debug(f"Emitted iteration triples for {iteration_uri}")
 
         await respond(AgentResponse(
-            chunk_type="explain",
+            message_type="explain",
             content="",
             explain_id=iteration_uri,
             explain_graph=GRAPH_RETRIEVAL,
@@ -311,7 +377,9 @@ class PatternBase:
 
     async def emit_observation_triples(self, flow, session_id, iteration_num,
                                        observation_text, request, respond,
-                                       context=None):
+                                       context=None,
+                                       tool_duration_ms=None,
+                                       tool_error=None):
         """Emit provenance triples for a standalone Observation entity."""
         iteration_uri = agent_iteration_uri(session_id, iteration_num)
         observation_entity_uri = agent_observation_uri(session_id, iteration_num)
@@ -344,6 +412,8 @@ class PatternBase:
                 observation_entity_uri,
                 parent_uri,
                 document_id=observation_doc_id,
+                tool_duration_ms=tool_duration_ms,
+                tool_error=tool_error,
             ),
             GRAPH_RETRIEVAL,
         )
@@ -358,7 +428,7 @@ class PatternBase:
         logger.debug(f"Emitted observation triples for {observation_entity_uri}")
 
         await respond(AgentResponse(
-            chunk_type="explain",
+            message_type="explain",
             content="",
             explain_id=observation_entity_uri,
             explain_graph=GRAPH_RETRIEVAL,
@@ -367,7 +437,7 @@ class PatternBase:
 
     async def emit_final_triples(self, flow, session_id, iteration_num,
                                  session_uri, answer_text, request, respond,
-                                 streaming):
+                                 streaming, termination_reason=None):
         """Emit provenance triples for the final answer and save to librarian."""
         final_uri = agent_final_uri(session_id)
 
@@ -401,6 +471,7 @@ class PatternBase:
                 question_uri=final_question_uri,
                 previous_uri=final_previous_uri,
                 document_id=answer_doc_id,
+                termination_reason=termination_reason,
             ),
             GRAPH_RETRIEVAL,
         )
@@ -415,7 +486,7 @@ class PatternBase:
         logger.debug(f"Emitted final triples for {final_uri}")
 
         await respond(AgentResponse(
-            chunk_type="explain",
+            message_type="explain",
             content="",
             explain_id=final_uri,
             explain_graph=GRAPH_RETRIEVAL,
@@ -439,7 +510,7 @@ class PatternBase:
             triples=triples,
         ))
         await respond(AgentResponse(
-            chunk_type="explain", content="",
+            message_type="explain", content="",
             explain_id=uri, explain_graph=GRAPH_RETRIEVAL,
             explain_triples=triples,
         ))
@@ -478,7 +549,7 @@ class PatternBase:
             triples=triples,
         ))
         await respond(AgentResponse(
-            chunk_type="explain", content="",
+            message_type="explain", content="",
             explain_id=uri, explain_graph=GRAPH_RETRIEVAL,
             explain_triples=triples,
         ))
@@ -498,7 +569,7 @@ class PatternBase:
             triples=triples,
         ))
         await respond(AgentResponse(
-            chunk_type="explain", content="",
+            message_type="explain", content="",
             explain_id=uri, explain_graph=GRAPH_RETRIEVAL,
             explain_triples=triples,
         ))
@@ -531,14 +602,14 @@ class PatternBase:
             triples=triples,
         ))
         await respond(AgentResponse(
-            chunk_type="explain", content="",
+            message_type="explain", content="",
             explain_id=uri, explain_graph=GRAPH_RETRIEVAL,
             explain_triples=triples,
         ))
 
     async def emit_synthesis_triples(
         self, flow, session_id, previous_uris, answer_text, user, collection,
-        respond, streaming,
+        respond, streaming, termination_reason=None,
     ):
         """Emit provenance for a synthesis answer."""
         uri = agent_synthesis_uri(session_id)
@@ -555,7 +626,10 @@ class PatternBase:
             doc_id = None
 
         triples = set_graph(
-            agent_synthesis_triples(uri, previous_uris, doc_id),
+            agent_synthesis_triples(
+                uri, previous_uris, doc_id,
+                termination_reason=termination_reason,
+            ),
             GRAPH_RETRIEVAL,
         )
         await flow("explainability").send(Triples(
@@ -563,7 +637,7 @@ class PatternBase:
             triples=triples,
         ))
         await respond(AgentResponse(
-            chunk_type="explain", content="",
+            message_type="explain", content="",
             explain_id=uri, explain_graph=GRAPH_RETRIEVAL,
             explain_triples=triples,
         ))
@@ -571,7 +645,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.
 
@@ -584,29 +659,35 @@ class PatternBase:
                 if text:
                     accumulated.append(text)
                     await respond(AgentResponse(
-                        chunk_type="answer",
+                        message_type="answer",
                         content=text,
                         end_of_message=False,
                         end_of_dialog=False,
                         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,33 +695,44 @@ 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:
                 await respond(AgentResponse(
-                    chunk_type="answer",
+                    message_type="answer",
                     content=answer_text,
                     end_of_message=False,
                     end_of_dialog=False,
                     message_id=message_id,
                 ))
         if streaming:
-            # End-of-dialog marker
+            # End-of-dialog marker with usage
             await respond(AgentResponse(
-                chunk_type="answer",
+                message_type="answer",
                 content="",
                 end_of_message=True,
                 end_of_dialog=True,
                 message_id=message_id,
+                **usage_kwargs,
             ))
         else:
             await respond(AgentResponse(
-                chunk_type="answer",
+                message_type="answer",
                 content=answer_text,
                 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.1}/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,11 @@ 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,
+                      pattern_decision_uri=None):
+
+        if usage is None:
+            usage = UsageTracker()
 
         streaming = getattr(request, 'streaming', False)
         session_id = getattr(request, 'session_id', '') or str(uuid.uuid4())
@@ -63,17 +67,19 @@ class PlanThenExecutePattern(PatternBase):
         # Determine current phase by checking history for a plan step
         plan = self._extract_plan(request.history)
 
+        derive_from_uri = pattern_decision_uri or session_uri
+
         if plan is None:
             await self._planning_iteration(
                 request, respond, next, flow,
-                session_id, collection, streaming, session_uri,
-                iteration_num,
+                session_id, collection, streaming, derive_from_uri,
+                iteration_num, usage=usage,
             )
         else:
             await self._execution_iteration(
                 request, respond, next, flow,
-                session_id, collection, streaming, session_uri,
-                iteration_num, plan,
+                session_id, collection, streaming, derive_from_uri,
+                iteration_num, plan, usage=usage,
             )
 
     def _extract_plan(self, history):
@@ -98,7 +104,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 +119,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 +130,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 +196,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 +208,7 @@ class PlanThenExecutePattern(PatternBase):
                 request, respond, next, flow,
                 session_id, collection, streaming,
                 session_uri, iteration_num, plan,
+                usage=usage,
             )
             return
 
@@ -240,7 +251,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 +269,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 +344,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 +380,7 @@ class PlanThenExecutePattern(PatternBase):
             respond=respond,
             streaming=streaming,
             message_id=synthesis_msg_id,
+            usage=usage,
         )
 
         # Emit synthesis provenance (links back to last step result)
@@ -372,6 +388,7 @@ class PlanThenExecutePattern(PatternBase):
         await self.emit_synthesis_triples(
             flow, session_id, last_step_uri,
             response_text, request.user, collection, respond, streaming,
+            termination_reason="plan-complete",
         )
 
         if self.is_subagent(request):
@@ -380,4 +397,5 @@ class PlanThenExecutePattern(PatternBase):
             await self.send_final_response(
                 respond, streaming, response_text, already_streamed=streaming,
                 message_id=synthesis_msg_id,
+                usage=usage,
             )