trustgraph-flow 1.1.9__tar.gz → 1.2.0__tar.gz

{trustgraph-flow-1.1.9 → trustgraph-flow-1.2.0}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: trustgraph-flow
-Version: 1.1.9
+Version: 1.2.0
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Home-page: https://github.com/trustgraph-ai/trustgraph
-Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v1.1.9.tar.gz
+Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v1.2.0.tar.gz
 Author: trustgraph.ai
 Author-email: security@trustgraph.ai
 Classifier: Programming Language :: Python :: 3

trustgraph-flow-1.2.0/scripts/kg-extract-agent

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

trustgraph-flow-1.2.0/scripts/prompt-template

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

{trustgraph-flow-1.1.9 → trustgraph-flow-1.2.0}/setup.py

@@ -34,7 +34,7 @@ setuptools.setup(
     python_requires='>=3.8',
     download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz",
     install_requires=[
-        "trustgraph-base>=1.1,<1.2",
+        "trustgraph-base>=1.2,<1.3",
         "aiohttp",
         "anthropic",
         "cassandra-driver",
@@ -96,7 +96,7 @@ setuptools.setup(
         "scripts/graph-rag",
         "scripts/kg-extract-definitions",
         "scripts/kg-extract-relationships",
-        "scripts/kg-extract-topics",
+        "scripts/kg-extract-agent",
         "scripts/kg-store",
         "scripts/kg-manager",
         "scripts/librarian",
@@ -106,7 +106,6 @@ setuptools.setup(
         "scripts/oe-write-milvus",
         "scripts/pdf-decoder",
         "scripts/pdf-ocr-mistral",
-        "scripts/prompt-generic",
         "scripts/prompt-template",
         "scripts/rows-write-cassandra",
         "scripts/run-processing",

{trustgraph-flow-1.1.9 → trustgraph-flow-1.2.0}/trustgraph/agent/mcp_tool/service.py

@@ -47,8 +47,8 @@ class Service(ToolService):
 
             url = self.mcp_services[name]["url"]
 
-            if "name" in self.mcp_services[name]:
-                remote_name = self.mcp_services[name]["name"]
+            if "remote-name" in self.mcp_services[name]:
+                remote_name = self.mcp_services[name]["remote-name"]
             else:
                 remote_name = name
 

trustgraph-flow-1.2.0/trustgraph/agent/react/agent_manager.py

@@ -0,0 +1,280 @@
+
+import logging
+import json
+import re
+
+from . types import Action, Final
+
+logger = logging.getLogger(__name__)
+
+class AgentManager:
+
+    def __init__(self, tools, additional_context=None):
+        self.tools = tools
+        self.additional_context = additional_context
+
+    def parse_react_response(self, text):
+        """Parse text-based ReAct response format.
+        
+        Expected format:
+        Thought: [reasoning about what to do next]
+        Action: [tool_name]
+        Args: {
+            "param": "value"
+        }
+        
+        OR
+        
+        Thought: [reasoning about the final answer]
+        Final Answer: [the answer]
+        """
+        if not isinstance(text, str):
+            raise ValueError(f"Expected string response, got {type(text)}")
+        
+        # Remove any markdown code blocks that might wrap the response
+        text = re.sub(r'^```[^\n]*\n', '', text.strip())
+        text = re.sub(r'\n```$', '', text.strip())
+            
+        lines = text.strip().split('\n')
+        
+        thought = None
+        action = None
+        args = None
+        final_answer = None
+        
+        i = 0
+        while i < len(lines):
+            line = lines[i].strip()
+            
+            # Parse Thought
+            if line.startswith("Thought:"):
+                thought = line[8:].strip()
+                # Handle multi-line thoughts
+                i += 1
+                while i < len(lines):
+                    next_line = lines[i].strip()
+                    if next_line.startswith(("Action:", "Final Answer:", "Args:")):
+                        break
+                    thought += " " + next_line
+                    i += 1
+                continue
+            
+            # Parse Final Answer
+            if line.startswith("Final Answer:"):
+                final_answer = line[13:].strip()
+                # Handle multi-line final answers (including JSON)
+                i += 1
+                
+                # Check if the answer might be JSON
+                if final_answer.startswith('{') or (i < len(lines) and lines[i].strip().startswith('{')):
+                    # Collect potential JSON answer
+                    json_text = final_answer if final_answer.startswith('{') else ""
+                    brace_count = json_text.count('{') - json_text.count('}')
+                    
+                    while i < len(lines) and (brace_count > 0 or not json_text):
+                        current_line = lines[i].strip()
+                        if current_line.startswith(("Thought:", "Action:")) and brace_count == 0:
+                            break
+                        json_text += ("\n" if json_text else "") + current_line
+                        brace_count += current_line.count('{') - current_line.count('}')
+                        i += 1
+                    
+                    # Try to parse as JSON
+                    # try:
+                    #     final_answer = json.loads(json_text)
+                    # except json.JSONDecodeError:
+                    #     # Not valid JSON, treat as regular text
+                    #     final_answer = json_text
+                    final_answer = json_text
+                else:
+                    # Regular text answer
+                    while i < len(lines):
+                        next_line = lines[i].strip()
+                        if next_line.startswith(("Thought:", "Action:")):
+                            break
+                        final_answer += " " + next_line
+                        i += 1
+                    
+                # If we have a final answer, return Final object
+                return Final(
+                    thought=thought or "",
+                    final=final_answer
+                )
+            
+            # Parse Action
+            if line.startswith("Action:"):
+                action = line[7:].strip()
+            
+            # Parse Args
+            if line.startswith("Args:"):
+                # Check if JSON starts on the same line
+                args_on_same_line = line[5:].strip()
+                if args_on_same_line:
+                    args_text = args_on_same_line
+                    brace_count = args_on_same_line.count('{') - args_on_same_line.count('}')
+                else:
+                    args_text = ""
+                    brace_count = 0
+                
+                # Collect all lines that form the JSON arguments
+                i += 1
+                started = bool(args_on_same_line and '{' in args_on_same_line)
+                
+                while i < len(lines) and (not started or brace_count > 0):
+                    current_line = lines[i]
+                    args_text += ("\n" if args_text else "") + current_line
+                    
+                    # Count braces to determine when JSON is complete
+                    for char in current_line:
+                        if char == '{':
+                            brace_count += 1
+                            started = True
+                        elif char == '}':
+                            brace_count -= 1
+                    
+                    # If we've started and braces are balanced, we're done
+                    if started and brace_count == 0:
+                        break
+                    
+                    i += 1
+                
+                # Parse the JSON arguments
+                try:
+                    args = json.loads(args_text.strip())
+                except json.JSONDecodeError as e:
+                    logger.error(f"Failed to parse JSON arguments: {args_text}")
+                    raise ValueError(f"Invalid JSON in Args: {e}")
+            
+            i += 1
+        
+        # If we have an action, return Action object
+        if action:
+            return Action(
+                thought=thought or "",
+                name=action,
+                arguments=args or {},
+                observation=""
+            )
+        
+        # If we only have a thought but no action or final answer
+        if thought and not action and not final_answer:
+            raise ValueError(f"Response has thought but no action or final answer: {text}")
+        
+        raise ValueError(f"Could not parse response: {text}")
+
+    async def reason(self, question, history, context):
+
+        print(f"calling reason: {question}", flush=True)
+
+        tools = self.tools
+
+        print(f"in reason", flush=True)
+        print(tools, flush=True)
+
+        tool_names = ",".join([
+            t for t in self.tools.keys()
+        ])
+
+        print("Tool names:", tool_names, flush=True)
+
+        variables = {
+            "question": question,
+            "tools": [
+                {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "arguments": [
+                        {
+                            "name": arg.name,
+                            "type": arg.type,
+                            "description": arg.description
+                        }
+                        for arg in tool.arguments
+                    ]
+                }
+                for tool in self.tools.values()
+            ],
+            "context": self.additional_context,
+            "question": question,
+            "tool_names": tool_names,
+            "history": [
+                {
+                    "thought": h.thought,
+                    "action": h.name,
+                    "arguments": h.arguments,
+                    "observation": h.observation,
+                }
+                for h in history
+            ]
+        }
+
+        print(json.dumps(variables, indent=4), flush=True)
+
+        logger.info(f"prompt: {variables}")
+
+        # Get text response from prompt service
+        response_text = await context("prompt-request").agent_react(variables)
+
+        print(f"Response text:\n{response_text}", flush=True)
+
+        logger.info(f"response: {response_text}")
+
+        # Parse the text response
+        try:
+            result = self.parse_react_response(response_text)
+            logger.info(f"Parsed result: {result}")
+            return result
+        except ValueError as e:
+            logger.error(f"Failed to parse response: {e}")
+            # Try to provide a helpful error message
+            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):
+
+        logger.info(f"question: {question}")
+        print(f"question: {question}", flush=True)
+
+        act = await self.reason(
+            question = question,
+            history = history,
+            context = context,
+        )
+        logger.info(f"act: {act}")
+
+        if isinstance(act, Final):
+
+            await think(act.thought)
+            return act
+
+        else:
+
+            await think(act.thought)
+
+            if act.name in self.tools:
+                action = self.tools[act.name]
+            else:
+                raise RuntimeError(f"No action for {act.name}!")
+
+            print("TOOL>>>", act, flush=True)
+
+            resp = await action.implementation(context).invoke(
+                **act.arguments
+            )
+
+            if isinstance(resp, str):
+                resp = resp.strip()
+            else:
+                resp = str(resp)
+                resp = resp.strip()
+
+            logger.info(f"resp: {resp}")
+
+            await observe(resp)
+
+            act.observation = resp
+
+            logger.info(f"iter: {act}")
+
+            return act
+

{trustgraph-flow-1.1.9 → trustgraph-flow-1.2.0}/trustgraph/agent/react/service.py

@@ -6,13 +6,17 @@ import json
 import re
 import sys
 import functools
+import logging
+
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
 
 from ... base import AgentService, TextCompletionClientSpec, PromptClientSpec
 from ... base import GraphRagClientSpec, ToolClientSpec
 
 from ... schema import AgentRequest, AgentResponse, AgentStep, Error
 
-from . tools import KnowledgeQueryImpl, TextCompletionImpl, McpToolImpl
+from . tools import KnowledgeQueryImpl, TextCompletionImpl, McpToolImpl, PromptImpl
 from . agent_manager import AgentManager
 
 from . types import Final, Action, Tool, Argument
@@ -79,64 +83,76 @@ class Processor(AgentService):
 
         print("Loading configuration version", version)
 
-        if self.config_key not in config:
-            print(f"No key {self.config_key} in config", flush=True)
-            return
-
-        config = config[self.config_key]
-
         try:
 
-            # This is some extra stuff to put in the prompt
-            additional = config.get("additional-context", None)
-
-            ix = json.loads(config["tool-index"])
-
             tools = {}
 
-            for k in ix:
-
-                pc = config[f"tool.{k}"]
-                data = json.loads(pc)
-
-                arguments = {
-                    v.get("name"): Argument(
-                        name = v.get("name"),
-                        type = v.get("type"),
-                        description = v.get("description")
+            # Load tool configurations from the new location
+            if "tool" in config:
+                for tool_id, tool_value in config["tool"].items():
+                    data = json.loads(tool_value)
+                    
+                    impl_id = data.get("type")
+                    name = data.get("name")
+                    
+                    # Create the appropriate implementation
+                    if impl_id == "knowledge-query":
+                        impl = functools.partial(
+                            KnowledgeQueryImpl, 
+                            collection=data.get("collection")
+                        )
+                        arguments = KnowledgeQueryImpl.get_arguments()
+                    elif impl_id == "text-completion":
+                        impl = TextCompletionImpl
+                        arguments = TextCompletionImpl.get_arguments()
+                    elif impl_id == "mcp-tool":
+                        impl = functools.partial(
+                            McpToolImpl, 
+                            mcp_tool_id=data.get("mcp-tool")
+                        )
+                        arguments = McpToolImpl.get_arguments()
+                    elif impl_id == "prompt":
+                        # For prompt tools, arguments come from config
+                        config_args = data.get("arguments", [])
+                        arguments = [
+                            Argument(
+                                name=arg.get("name"),
+                                type=arg.get("type"),
+                                description=arg.get("description")
+                            )
+                            for arg in config_args
+                        ]
+                        impl = functools.partial(
+                            PromptImpl,
+                            template_id=data.get("template"),
+                            arguments=arguments
+                        )
+                    else:
+                        raise RuntimeError(
+                            f"Tool type {impl_id} not known"
+                        )
+                    
+                    tools[name] = Tool(
+                        name=name,
+                        description=data.get("description"),
+                        implementation=impl,
+                        config=data,  # Store full config for reference
+                        arguments=arguments,
                     )
-                    for v in data["arguments"]
-                }
-
-                impl_id = data.get("type")
-
-                name = data.get("name")
-
-                if impl_id == "knowledge-query":
-                    impl = KnowledgeQueryImpl
-                elif impl_id == "text-completion":
-                    impl = TextCompletionImpl
-                elif impl_id == "mcp-tool":
-                    impl = functools.partial(McpToolImpl, name=k)
-                else:
-                    raise RuntimeError(
-                        f"Tool-kind {impl_id} not known"
-                    )
-
-                tools[data.get("name")] = Tool(
-                    name = name,
-                    description = data.get("description"),
-                    implementation = impl,
-                    config=data.get("config", {}),
-                    arguments = arguments,
-                )
-
+            
+            # Load additional context from agent config if it exists
+            additional = None
+            if self.config_key in config:
+                agent_config = config[self.config_key]
+                additional = agent_config.get("additional-context", None)
+            
             self.agent = AgentManager(
                 tools=tools,
                 additional_context=additional
             )
 
-            print("Prompt configuration reloaded.", flush=True)
+            print(f"Loaded {len(tools)} tools", flush=True)
+            print("Tool configuration reloaded.", flush=True)
 
         except Exception as e:
 
@@ -209,6 +225,11 @@ class Processor(AgentService):
 
                 print("Send final response...", flush=True)
 
+                if isinstance(act.final, str):
+                    f = act.final
+                else:
+                    f = json.dumps(act.final)
+
                 r = AgentResponse(
                     answer=act.final,
                     error=None,
@@ -280,6 +301,5 @@ class Processor(AgentService):
         )
 
 def run():
-
     Processor.launch(default_ident, __doc__)
 

trustgraph-flow-1.2.0/trustgraph/agent/react/tools.py

@@ -0,0 +1,101 @@
+
+import json
+from .types import Argument
+
+# This tool implementation knows how to put a question to the graph RAG
+# service
+class KnowledgeQueryImpl:
+    def __init__(self, context, collection=None):
+        self.context = context
+        self.collection = collection
+    
+    @staticmethod
+    def get_arguments():
+        return [
+            Argument(
+                name="question",
+                type="string",
+                description="The question to ask the knowledge base"
+            )
+        ]
+    
+    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
+    
+    @staticmethod
+    def get_arguments():
+        return [
+            Argument(
+                name="question",
+                type="string",
+                description="The text prompt or question for completion"
+            )
+        ]
+    
+    async def invoke(self, **arguments):
+        client = self.context("prompt-request")
+        print("Prompt question...", flush=True)
+        return await client.question(
+            arguments.get("question")
+        )
+
+# This tool implementation knows how to do MCP tool invocation.  This uses
+# the mcp-tool service.
+class McpToolImpl:
+
+    def __init__(self, context, mcp_tool_id):
+        self.context = context
+        self.mcp_tool_id = mcp_tool_id
+    
+    @staticmethod
+    def get_arguments():
+        # MCP tools define their own arguments dynamically
+        # For now, we return empty list and let the MCP service handle validation
+        return []
+
+    async def invoke(self, **arguments):
+
+        client = self.context("mcp-tool-request")
+
+        print(f"MCP tool invocation: {self.mcp_tool_id}...", flush=True)
+        output = await client.invoke(
+            name = self.mcp_tool_id,
+            parameters = arguments,  # Pass the actual arguments
+        )
+
+        print(output)
+
+        if isinstance(output, str):
+            return output
+        else:
+            return json.dumps(output)
+
+
+# This tool implementation knows how to execute prompt templates
+class PromptImpl:
+    def __init__(self, context, template_id, arguments=None):
+        self.context = context
+        self.template_id = template_id
+        self.arguments = arguments or []  # These come from config
+    
+    def get_arguments(self):
+        # For prompt tools, arguments are defined in configuration
+        return self.arguments
+    
+    async def invoke(self, **arguments):
+        client = self.context("prompt-request")
+        print(f"Prompt template invocation: {self.template_id}...", flush=True)
+        return await client.prompt(
+            id=self.template_id,
+            variables=arguments
+        )

trustgraph-flow-1.2.0/trustgraph/extract/kg/agent/__init__.py

@@ -0,0 +1 @@
+from .extract import *

trustgraph-flow-1.2.0/trustgraph/extract/kg/agent/__main__.py

@@ -0,0 +1,4 @@
+from .extract import Processor
+
+if __name__ == "__main__":
+    Processor.run()