trustgraph-flow 2.7.3__tar.gz → 2.7.4__tar.gz

{trustgraph_flow-2.7.3 → trustgraph_flow-2.7.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 2.7.3
+Version: 2.7.4
 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

trustgraph_flow-2.7.4/trustgraph/flow_version.py

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

{trustgraph_flow-2.7.3 → trustgraph_flow-2.7.4}/trustgraph/model/text_completion/azure/llm.py

@@ -55,7 +55,10 @@ class Processor(LlmService):
         self.max_output = max_output
         self.default_model = model
 
-    def build_prompt(self, system, content, temperature=None, stream=False, model=None):
+    def build_prompt(
+        self, system, content, temperature=None, stream=False,
+        model=None, response_format=None, schema=None,
+    ):
         # Use provided temperature or fall back to default
         effective_temperature = temperature if temperature is not None else self.temperature
         model_name = model or self.default_model
@@ -79,6 +82,17 @@ class Processor(LlmService):
             data["stream"] = True
             data["stream_options"] = {"include_usage": True}
 
+        if response_format == "json" and schema is not None:
+            data["response_format"] = {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "response",
+                    "schema": schema,
+                    "strict": True,
+                },
+            }
+            logger.debug("Structured output enabled")
+
         body = json.dumps(data)
 
         return body
@@ -109,7 +123,10 @@ class Processor(LlmService):
 
         return result
 
-    async def generate_content(self, system, prompt, model=None, temperature=None):
+    async def generate_content(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
 
         # Use provided model or fall back to default
         model_name = model or self.default_model
@@ -125,7 +142,9 @@ class Processor(LlmService):
                 system,
                 prompt,
                 effective_temperature,
-                model=model_name
+                model=model_name,
+                response_format=response_format,
+                schema=schema,
             )
 
             response = self.call_llm(prompt)
@@ -169,7 +188,10 @@ class Processor(LlmService):
         """Azure serverless endpoints support streaming"""
         return True
 
-    async def generate_content_stream(self, system, prompt, model=None, temperature=None):
+    async def generate_content_stream(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
         """Stream content generation from Azure serverless endpoint"""
         model_name = model or self.default_model
         effective_temperature = temperature if temperature is not None else self.temperature
@@ -178,7 +200,11 @@ class Processor(LlmService):
         logger.debug(f"Using temperature: {effective_temperature}")
 
         try:
-            body = self.build_prompt(system, prompt, effective_temperature, stream=True, model=model_name)
+            body = self.build_prompt(
+                system, prompt, effective_temperature, stream=True,
+                model=model_name, response_format=response_format,
+                schema=schema,
+            )
 
             url = self.endpoint
             api_key = self.token

{trustgraph_flow-2.7.3 → trustgraph_flow-2.7.4}/trustgraph/model/text_completion/azure_openai/llm.py

@@ -62,7 +62,10 @@ class Processor(LlmService):
             azure_endpoint = endpoint,
         )
 
-    async def generate_content(self, system, prompt, model=None, temperature=None):
+    async def generate_content(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
 
         # Use provided model or fall back to default
         model_name = model or self.default_model
@@ -76,6 +79,18 @@ class Processor(LlmService):
 
         try:
 
+            kwargs = {}
+            if response_format == "json" and schema is not None:
+                kwargs["response_format"] = {
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "response",
+                        "schema": schema,
+                        "strict": True,
+                    },
+                }
+                logger.debug("Structured output enabled")
+
             resp = self.openai.chat.completions.create(
                 model=model_name,
                 messages=[
@@ -92,6 +107,7 @@ class Processor(LlmService):
                 temperature=effective_temperature,
                 max_completion_tokens=self.max_output,
                 top_p=1,
+                **kwargs,
             )
 
             inputtokens = resp.usage.prompt_tokens
@@ -129,7 +145,10 @@ class Processor(LlmService):
         """Azure OpenAI supports streaming"""
         return True
 
-    async def generate_content_stream(self, system, prompt, model=None, temperature=None):
+    async def generate_content_stream(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
         """
         Stream content generation from Azure OpenAI.
         Yields LlmChunk objects with is_final=True on the last chunk.

{trustgraph_flow-2.7.3 → trustgraph_flow-2.7.4}/trustgraph/model/text_completion/claude/llm.py

@@ -48,7 +48,10 @@ class Processor(LlmService):
 
         logger.info("Claude LLM service initialized")
 
-    async def generate_content(self, system, prompt, model=None, temperature=None):
+    async def generate_content(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
 
         # Use provided model or fall back to default
         model_name = model or self.default_model
@@ -60,11 +63,27 @@ class Processor(LlmService):
 
         try:
 
+            kwargs = {}
+            use_tool_extract = False
+
+            if response_format == "json" and schema is not None:
+                kwargs["tools"] = [{
+                    "name": "structured_response",
+                    "description": "Return the structured response",
+                    "input_schema": schema,
+                }]
+                kwargs["tool_choice"] = {
+                    "type": "tool",
+                    "name": "structured_response",
+                }
+                use_tool_extract = True
+                logger.debug("Structured output enabled (tool-use)")
+
             response = message = self.claude.messages.create(
                 model=model_name,
                 max_tokens=self.max_output,
                 temperature=effective_temperature,
-                system = system,
+                system=system,
                 messages=[
                     {
                         "role": "user",
@@ -75,10 +94,22 @@ class Processor(LlmService):
                             }
                         ]
                     }
-                ]
+                ],
+                **kwargs,
             )
 
-            resp = response.content[0].text
+            if use_tool_extract:
+                import json
+                tool_block = next(
+                    (b for b in response.content if b.type == "tool_use"),
+                    None,
+                )
+                if tool_block:
+                    resp = json.dumps(tool_block.input)
+                else:
+                    resp = response.content[0].text
+            else:
+                resp = response.content[0].text
             inputtokens = response.usage.input_tokens
             outputtokens = response.usage.output_tokens
             logger.debug(f"LLM response: {resp}")
@@ -110,7 +141,10 @@ class Processor(LlmService):
         """Claude/Anthropic supports streaming"""
         return True
 
-    async def generate_content_stream(self, system, prompt, model=None, temperature=None):
+    async def generate_content_stream(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
         """Stream content generation from Claude"""
         model_name = model or self.default_model
         effective_temperature = temperature if temperature is not None else self.temperature

{trustgraph_flow-2.7.3 → trustgraph_flow-2.7.4}/trustgraph/model/text_completion/cohere/llm.py

@@ -46,7 +46,10 @@ class Processor(LlmService):
 
         logger.info("Cohere LLM service initialized")
 
-    async def generate_content(self, system, prompt, model=None, temperature=None):
+    async def generate_content(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
 
         # Use provided model or fall back to default
         model_name = model or self.default_model
@@ -104,7 +107,10 @@ class Processor(LlmService):
         """Cohere supports streaming"""
         return True
 
-    async def generate_content_stream(self, system, prompt, model=None, temperature=None):
+    async def generate_content_stream(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
         """Stream content generation from Cohere"""
         model_name = model or self.default_model
         effective_temperature = temperature if temperature is not None else self.temperature

{trustgraph_flow-2.7.3 → trustgraph_flow-2.7.4}/trustgraph/model/text_completion/llamafile/llm.py

@@ -50,7 +50,10 @@ class Processor(LlmService):
 
         logger.info("Llamafile LLM service initialized")
 
-    async def generate_content(self, system, prompt, model=None, temperature=None):
+    async def generate_content(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
 
         # Use provided model or fall back to default
         model_name = model or self.default_model
@@ -64,6 +67,18 @@ class Processor(LlmService):
 
         try:
 
+            if response_format == "json" and schema is not None:
+                fmt = {
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "response",
+                        "schema": schema,
+                    },
+                }
+                logger.debug("Structured output enabled")
+            else:
+                fmt = {"type": "text"}
+
             resp = self.openai.chat.completions.create(
                 model=model_name,
                 messages=[
@@ -74,9 +89,7 @@ class Processor(LlmService):
                 top_p=1,
                 frequency_penalty=0,
                 presence_penalty=0,
-                response_format={
-                    "type": "text"
-                }
+                response_format=fmt,
             )
 
             inputtokens = resp.usage.prompt_tokens
@@ -106,7 +119,10 @@ class Processor(LlmService):
         """LlamaFile supports streaming"""
         return True
 
-    async def generate_content_stream(self, system, prompt, model=None, temperature=None):
+    async def generate_content_stream(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
         """Stream content generation from LlamaFile"""
         model_name = model or self.default_model
         effective_temperature = temperature if temperature is not None else self.temperature
@@ -117,6 +133,18 @@ class Processor(LlmService):
         prompt = system + "\n\n" + prompt
 
         try:
+            if response_format == "json" and schema is not None:
+                fmt = {
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "response",
+                        "schema": schema,
+                    },
+                }
+                logger.debug("Structured output enabled (streaming)")
+            else:
+                fmt = {"type": "text"}
+
             response = self.openai.chat.completions.create(
                 model=model_name,
                 messages=[{"role": "user", "content": prompt}],
@@ -125,7 +153,7 @@ class Processor(LlmService):
                 top_p=1,
                 frequency_penalty=0,
                 presence_penalty=0,
-                response_format={"type": "text"},
+                response_format=fmt,
                 stream=True,
                 stream_options={"include_usage": True}
             )

{trustgraph_flow-2.7.3 → trustgraph_flow-2.7.4}/trustgraph/model/text_completion/lmstudio/llm.py

@@ -50,7 +50,10 @@ class Processor(LlmService):
 
         logger.info("LMStudio LLM service initialized")
 
-    async def generate_content(self, system, prompt, model=None, temperature=None):
+    async def generate_content(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
 
         # Use provided model or fall back to default
         model_name = model or self.default_model
@@ -66,6 +69,18 @@ class Processor(LlmService):
 
             logger.debug(f"Prompt: {prompt}")
 
+            if response_format == "json" and schema is not None:
+                fmt = {
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "response",
+                        "schema": schema,
+                    },
+                }
+                logger.debug("Structured output enabled")
+            else:
+                fmt = {"type": "text"}
+
             resp = self.openai.chat.completions.create(
                 model=model_name,
                 messages=[
@@ -76,9 +91,7 @@ class Processor(LlmService):
                 top_p=1,
                 frequency_penalty=0,
                 presence_penalty=0,
-                response_format={
-                    "type": "text"
-                }
+                response_format=fmt,
             )
 
             logger.debug(f"Full response: {resp}")
@@ -110,7 +123,10 @@ class Processor(LlmService):
         """LM Studio supports streaming"""
         return True
 
-    async def generate_content_stream(self, system, prompt, model=None, temperature=None):
+    async def generate_content_stream(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
         """Stream content generation from LM Studio"""
         model_name = model or self.default_model
         effective_temperature = temperature if temperature is not None else self.temperature
@@ -121,6 +137,18 @@ class Processor(LlmService):
         prompt = system + "\n\n" + prompt
 
         try:
+            if response_format == "json" and schema is not None:
+                fmt = {
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "response",
+                        "schema": schema,
+                    },
+                }
+                logger.debug("Structured output enabled (streaming)")
+            else:
+                fmt = {"type": "text"}
+
             response = self.openai.chat.completions.create(
                 model=model_name,
                 messages=[{"role": "user", "content": prompt}],
@@ -129,7 +157,7 @@ class Processor(LlmService):
                 top_p=1,
                 frequency_penalty=0,
                 presence_penalty=0,
-                response_format={"type": "text"},
+                response_format=fmt,
                 stream=True,
                 stream_options={"include_usage": True}
             )

{trustgraph_flow-2.7.3 → trustgraph_flow-2.7.4}/trustgraph/model/text_completion/mistral/llm.py

@@ -48,7 +48,10 @@ class Processor(LlmService):
 
         logger.info("Mistral LLM service initialized")
 
-    async def generate_content(self, system, prompt, model=None, temperature=None):
+    async def generate_content(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
 
         # Use provided model or fall back to default
         model_name = model or self.default_model
@@ -62,6 +65,19 @@ class Processor(LlmService):
 
         try:
 
+            if response_format == "json" and schema is not None:
+                fmt = {
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "response",
+                        "schema": schema,
+                        "strict": True,
+                    },
+                }
+                logger.debug("Structured output enabled")
+            else:
+                fmt = {"type": "text"}
+
             resp = self.mistral.chat.complete(
                 model=model_name,
                 messages=[
@@ -80,9 +96,7 @@ class Processor(LlmService):
                 top_p=1,
                 frequency_penalty=0,
                 presence_penalty=0,
-                response_format={
-                    "type": "text"
-                }
+                response_format=fmt,
             )
             
             inputtokens = resp.usage.prompt_tokens
@@ -120,7 +134,10 @@ class Processor(LlmService):
         """Mistral supports streaming"""
         return True
 
-    async def generate_content_stream(self, system, prompt, model=None, temperature=None):
+    async def generate_content_stream(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
         """Stream content generation from Mistral"""
         model_name = model or self.default_model
         effective_temperature = temperature if temperature is not None else self.temperature
@@ -131,6 +148,19 @@ class Processor(LlmService):
         prompt = system + "\n\n" + prompt
 
         try:
+            if response_format == "json" and schema is not None:
+                fmt = {
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "response",
+                        "schema": schema,
+                        "strict": True,
+                    },
+                }
+                logger.debug("Structured output enabled (streaming)")
+            else:
+                fmt = {"type": "text"}
+
             stream = self.mistral.chat.stream(
                 model=model_name,
                 messages=[
@@ -149,7 +179,7 @@ class Processor(LlmService):
                 top_p=1,
                 frequency_penalty=0,
                 presence_penalty=0,
-                response_format={"type": "text"}
+                response_format=fmt,
             )
 
             total_input_tokens = 0

{trustgraph_flow-2.7.3 → trustgraph_flow-2.7.4}/trustgraph/model/text_completion/ollama/llm.py

@@ -62,7 +62,10 @@ class Processor(LlmService):
             else:
                 logger.warning(f"Failed to check Ollama model '{model_name}': {e}")
 
-    async def generate_content(self, system, prompt, model=None, temperature=None):
+    async def generate_content(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
 
         # Use provided model or fall back to default
         model_name = model or self.default_model
@@ -79,7 +82,12 @@ class Processor(LlmService):
 
         try:
 
-            response = await self.llm.generate(model_name, prompt, options={'temperature': effective_temperature})
+            kwargs = {}
+            if response_format == "json" and schema is not None:
+                kwargs["format"] = schema
+                logger.debug("Structured output enabled")
+
+            response = await self.llm.generate(model_name, prompt, options={'temperature': effective_temperature}, **kwargs)
 
             response_text = response['response']
             logger.debug("Sending response...")
@@ -108,7 +116,10 @@ class Processor(LlmService):
         """Ollama supports streaming"""
         return True
 
-    async def generate_content_stream(self, system, prompt, model=None, temperature=None):
+    async def generate_content_stream(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
         """Stream content generation from Ollama"""
         model_name = model or self.default_model
 
@@ -123,11 +134,17 @@ class Processor(LlmService):
         prompt = system + "\n\n" + prompt
 
         try:
+            kwargs = {}
+            if response_format == "json" and schema is not None:
+                kwargs["format"] = schema
+                logger.debug("Structured output enabled (streaming)")
+
             stream = await self.llm.generate(
                 model_name,
                 prompt,
                 options={'temperature': effective_temperature},
-                stream=True
+                stream=True,
+                **kwargs,
             )
 
             total_input_tokens = 0

{trustgraph_flow-2.7.3 → trustgraph_flow-2.7.4}/trustgraph/model/text_completion/openai/llm.py

@@ -82,7 +82,10 @@ class Processor(LlmService):
             return self.variant.extract_content(message)
         return message.content
 
-    async def generate_content(self, system, prompt, model=None, temperature=None):
+    async def generate_content(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
 
         model_name = model or self.default_model
         effective_temperature = temperature if temperature is not None else self.temperature
@@ -96,6 +99,25 @@ class Processor(LlmService):
 
             api_kwargs = self._build_kwargs(model_name, effective_temperature)
 
+            if response_format == "json" and schema is not None:
+                is_top_level_array = schema.get("type") == "array"
+                if is_top_level_array and not self.variant.supports_top_level_array():
+                    logger.debug(
+                        "Variant %s does not support top-level array "
+                        "schemas, falling back to free-text",
+                        self.variant.name,
+                    )
+                else:
+                    api_kwargs["response_format"] = {
+                        "type": "json_schema",
+                        "json_schema": {
+                            "name": "response",
+                            "schema": schema,
+                            "strict": True,
+                        },
+                    }
+                    logger.debug("Structured output enabled")
+
             messages = [
                 {
                     "role": "user",
@@ -160,7 +182,10 @@ class Processor(LlmService):
         """OpenAI supports streaming"""
         return True
 
-    async def generate_content_stream(self, system, prompt, model=None, temperature=None):
+    async def generate_content_stream(
+        self, system, prompt, model=None, temperature=None,
+        response_format=None, schema=None,
+    ):
         """
         Stream content generation from OpenAI.
         Yields LlmChunk objects with is_final=True on the last chunk.
@@ -176,6 +201,25 @@ class Processor(LlmService):
         try:
             api_kwargs = self._build_kwargs(model_name, effective_temperature)
 
+            if response_format == "json" and schema is not None:
+                is_top_level_array = schema.get("type") == "array"
+                if is_top_level_array and not self.variant.supports_top_level_array():
+                    logger.debug(
+                        "Variant %s does not support top-level array "
+                        "schemas, falling back to free-text (streaming)",
+                        self.variant.name,
+                    )
+                else:
+                    api_kwargs["response_format"] = {
+                        "type": "json_schema",
+                        "json_schema": {
+                            "name": "response",
+                            "schema": schema,
+                            "strict": True,
+                        },
+                    }
+                    logger.debug("Structured output enabled (streaming)")
+
             messages = [
                 {
                     "role": "user",

{trustgraph_flow-2.7.3 → trustgraph_flow-2.7.4}/trustgraph/model/text_completion/openai/variants.py

@@ -62,6 +62,10 @@ class Variant:
         """Extract thinking content from a streaming delta."""
         return getattr(delta, "reasoning_content", None)
 
+    def supports_top_level_array(self):
+        """Whether this provider accepts a top-level array JSON schema."""
+        return True
+
     def create_completion(self, client, model, messages, **kwargs):
         """Call the completions API. Override for non-standard SDKs."""
         return client.chat.completions.create(
@@ -84,6 +88,9 @@ class OpenAIVariant(Variant):
     token_param = "max_completion_tokens"
     temperature_with_thinking = False
 
+    def supports_top_level_array(self):
+        return False
+
     def thinking_kwargs(self, effort):
         return {"reasoning_effort": effort}
 
@@ -195,6 +202,12 @@ class LlamaVariant(Variant):
         return re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL).strip()
 
 
+class VllmVariant(LlamaVariant):
+    """vLLM via OpenAI-compatible API. Supports full structured output."""
+
+    name = "vllm"
+
+
 VARIANTS = {
     "openai": OpenAIVariant,
     "deepseek": DeepSeekVariant,
@@ -203,6 +216,7 @@ VARIANTS = {
     "dashscope": DashScopeVariant,
     "glm": GlmVariant,
     "llama": LlamaVariant,
+    "vllm": VllmVariant,
 }
 
 DEFAULT_VARIANT = "openai"