trustgraph-flow 2.6.3__tar.gz → 2.6.5__tar.gz

{trustgraph_flow-2.6.3 → trustgraph_flow-2.6.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 2.6.3
+Version: 2.6.5
 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
@@ -17,6 +17,7 @@ Requires-Dist: cryptography
 Requires-Dist: faiss-cpu
 Requires-Dist: falkordb
 Requires-Dist: fastembed
+Requires-Dist: flashrank
 Requires-Dist: ibis
 Requires-Dist: jsonschema
 Requires-Dist: langchain

{trustgraph_flow-2.6.3 → trustgraph_flow-2.6.5}/pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "faiss-cpu",
     "falkordb",
     "fastembed",
+    "flashrank",
     "ibis",
     "jsonschema",
     "langchain",
@@ -83,6 +84,7 @@ graph-embeddings-write-pinecone = "trustgraph.storage.graph_embeddings.pinecone:
 graph-embeddings-write-qdrant = "trustgraph.storage.graph_embeddings.qdrant:run"
 graph-embeddings = "trustgraph.embeddings.graph_embeddings:run"
 graph-rag = "trustgraph.retrieval.graph_rag:run"
+reranker-flashrank = "trustgraph.reranker.flashrank:run"
 kg-extract-agent = "trustgraph.extract.kg.agent:run"
 kg-extract-definitions = "trustgraph.extract.kg.definitions:run"
 kg-extract-rows = "trustgraph.extract.kg.rows:run"

{trustgraph_flow-2.6.3 → trustgraph_flow-2.6.5}/trustgraph/bootstrap/initialisers/default_flow_start.py

@@ -18,6 +18,10 @@ description : str (default "Default")
     Human-readable description passed to flow-svc.
 parameters : dict (optional)
     Optional parameter overrides passed to start-flow.
+list_timeout : int (default 10)
+    Timeout in seconds for the list-flows request.
+start_timeout : int (default 30)
+    Timeout in seconds for the start-flow request.
 """
 
 from trustgraph.schema import FlowRequest
@@ -34,6 +38,8 @@ class DefaultFlowStart(Initialiser):
             blueprint=None,
             description="Default",
             parameters=None,
+            list_timeout=10,
+            start_timeout=30,
             **kwargs,
     ):
         super().__init__(**kwargs)
@@ -46,6 +52,8 @@ class DefaultFlowStart(Initialiser):
         self.blueprint = blueprint
         self.description = description
         self.parameters = dict(parameters) if parameters else {}
+        self.list_timeout = list_timeout
+        self.start_timeout = start_timeout
 
     async def run(self, ctx, old_flag, new_flag):
 
@@ -70,7 +78,7 @@ class DefaultFlowStart(Initialiser):
                 FlowRequest(
                     operation="list-flows",
                 ),
-                timeout=10,
+                timeout=self.list_timeout,
             )
             if list_resp.error:
                 raise RuntimeError(
@@ -99,7 +107,7 @@ class DefaultFlowStart(Initialiser):
                     description=self.description,
                     parameters=self.parameters,
                 ),
-                timeout=30,
+                timeout=self.start_timeout,
             )
             if resp.error:
                 raise RuntimeError(

{trustgraph_flow-2.6.3 → trustgraph_flow-2.6.5}/trustgraph/bootstrap/initialisers/workspace_init.py

@@ -14,7 +14,9 @@ seed_file : str (required when source=="seed-file")
     Path to a JSON seed file with the same shape TemplateSeed consumes.
 overwrite : bool (default False)
     On re-run (flag change), if True overwrite all keys; if False,
-    upsert-missing-only (preserves in-workspace customisations).
+    upsert-missing-only (preserves in-workspace customisations)
+iam_timeout : int (default 10)
+    Timeout in seconds for the IAM create-workspace request.
 
 Raises (in ``run``)
 -------------------
@@ -41,7 +43,9 @@ class WorkspaceInit(Initialiser):
             source="template",
             seed_file=None,
             overwrite=False,
+            iam_timeout=10,
             **kwargs,
+            
     ):
         super().__init__(**kwargs)
 
@@ -59,6 +63,7 @@ class WorkspaceInit(Initialiser):
         self.source = source
         self.seed_file = seed_file
         self.overwrite = overwrite
+        self.iam_timeout = iam_timeout
 
     async def run(self, ctx, old_flag, new_flag):
         await self._create_workspace(ctx)
@@ -120,10 +125,10 @@ class WorkspaceInit(Initialiser):
                     workspace_record=WorkspaceInput(
                         id=self.workspace,
                         name=self.workspace.title(),
-                        enabled=True,
+                        enabled=True,  
                     ),
                 ),
-                timeout=10,
+                timeout=self.iam_timeout,
             )
             if resp.error:
                 if resp.error.type == "duplicate":

trustgraph_flow-2.6.5/trustgraph/flow_version.py

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

{trustgraph_flow-2.6.3 → trustgraph_flow-2.6.5}/trustgraph/gateway/dispatch/manager.py

@@ -37,6 +37,7 @@ from . graph_embeddings_query import GraphEmbeddingsQueryRequestor
 from . document_embeddings_query import DocumentEmbeddingsQueryRequestor
 from . row_embeddings_query import RowEmbeddingsQueryRequestor
 from . mcp_tool import McpToolRequestor
+from . reranker import RerankerRequestor
 from . text_load import TextLoad
 from . document_load import DocumentLoad
 
@@ -74,6 +75,7 @@ request_response_dispatchers = {
     "structured-diag": StructuredDiagRequestor,
     "row-embeddings": RowEmbeddingsQueryRequestor,
     "sparql": SparqlQueryRequestor,
+    "reranker": RerankerRequestor,
 }
 
 system_dispatchers = {

trustgraph_flow-2.6.5/trustgraph/gateway/dispatch/reranker.py

@@ -0,0 +1,31 @@
+
+from ... schema import RerankerRequest, RerankerResponse
+from ... messaging import TranslatorRegistry
+
+from . requestor import ServiceRequestor
+
+class RerankerRequestor(ServiceRequestor):
+    def __init__(
+            self, backend, request_queue, response_queue, timeout,
+            consumer, subscriber,
+    ):
+
+        super(RerankerRequestor, self).__init__(
+            backend=backend,
+            request_queue=request_queue,
+            response_queue=response_queue,
+            request_schema=RerankerRequest,
+            response_schema=RerankerResponse,
+            subscription = subscriber,
+            consumer_name = consumer,
+            timeout=timeout,
+        )
+
+        self.request_translator = TranslatorRegistry.get_request_translator("reranker")
+        self.response_translator = TranslatorRegistry.get_response_translator("reranker")
+
+    def to_request(self, body):
+        return self.request_translator.decode(body)
+
+    def from_response(self, message):
+        return self.response_translator.encode_with_completion(message)

{trustgraph_flow-2.6.3 → trustgraph_flow-2.6.5}/trustgraph/gateway/registry.py

@@ -518,6 +518,7 @@ _FLOW_SERVICES = {
     "structured-diag": "structured-query:read",
     "row-embeddings": "row-embeddings:read",
     "sparql": "sparql:read",
+    "reranker": "reranker",
 }
 for _kind, _cap in _FLOW_SERVICES.items():
     _register_flow_kind("flow-service", _kind, _cap)

{trustgraph_flow-2.6.3 → trustgraph_flow-2.6.5}/trustgraph/iam/service/iam.py

@@ -72,6 +72,7 @@ _READER_CAPS = {
     "row-embeddings:read",
     "llm",
     "embeddings",
+    "reranker",
     "mcp",
     "config:read",
     "flows:read",

{trustgraph_flow-2.6.3 → trustgraph_flow-2.6.5}/trustgraph/model/text_completion/openai/llm.py

@@ -10,6 +10,7 @@ import logging
 
 from .... exceptions import TooManyRequests, LlmError
 from .... base import LlmService, LlmResult, LlmChunk
+from . variants import get_variant, DEFAULT_VARIANT, VARIANTS
 
 # Module logger
 logger = logging.getLogger(__name__)
@@ -21,6 +22,7 @@ default_temperature = 0.0
 default_max_output = 4096
 default_api_key = os.getenv("OPENAI_TOKEN")
 default_base_url = os.getenv("OPENAI_BASE_URL")
+default_thinking = "off"
 
 if default_base_url is None or default_base_url == "":
     default_base_url = "https://api.openai.com/v1"
@@ -28,16 +30,21 @@ if default_base_url is None or default_base_url == "":
 class Processor(LlmService):
 
     def __init__(self, **params):
-    
+
         model = params.get("model", default_model)
         api_key = params.get("api_key", default_api_key)
         base_url = params.get("url", default_base_url)
         temperature = params.get("temperature", default_temperature)
         max_output = params.get("max_output", default_max_output)
+        thinking = params.get("thinking", default_thinking)
+        variant_name = params.get("variant", DEFAULT_VARIANT)
 
         if not api_key:
             api_key = "not-set"
 
+        self.variant = get_variant(variant_name)
+        self.thinking = thinking
+
         super(Processor, self).__init__(
             **params | {
                 "model": model,
@@ -56,13 +63,28 @@ class Processor(LlmService):
         else:
             self.openai = OpenAI(api_key=api_key)
 
-        logger.info("OpenAI LLM service initialized")
+        logger.info(
+            f"OpenAI LLM service initialized "
+            f"(variant={self.variant.name}, thinking={self.thinking})"
+        )
+
+    def _build_kwargs(self, model_name, temperature):
+        """Build API call kwargs using the active variant."""
+        return self.variant.completion_kwargs(
+            max_output=self.max_output,
+            temperature=temperature,
+            thinking=self.thinking,
+        )
+
+    def _extract_content(self, message):
+        """Extract visible content from a response message."""
+        if hasattr(self.variant, "extract_content"):
+            return self.variant.extract_content(message)
+        return message.content
 
     async def generate_content(self, system, prompt, model=None, temperature=None):
 
-        # Use provided model or fall back to default
         model_name = model or self.default_model
-        # Use provided temperature or fall back to default
         effective_temperature = temperature if temperature is not None else self.temperature
 
         logger.debug(f"Using model: {model_name}")
@@ -72,6 +94,8 @@ class Processor(LlmService):
 
         try:
 
+            api_kwargs = self._build_kwargs(model_name, effective_temperature)
+
             resp = self.openai.chat.completions.create(
                 model=model_name,
                 messages=[
@@ -85,18 +109,23 @@ class Processor(LlmService):
                         ]
                     }
                 ],
-                temperature=effective_temperature,
-                max_completion_tokens=self.max_output,
+                **api_kwargs,
             )
-            
+
             inputtokens = resp.usage.prompt_tokens
             outputtokens = resp.usage.completion_tokens
-            logger.debug(f"LLM response: {resp.choices[0].message.content}")
+
+            content = self._extract_content(resp.choices[0].message)
+            thinking = self.variant.extract_thinking(resp.choices[0].message)
+
+            logger.debug(f"LLM response: {content}")
+            if thinking:
+                logger.debug(f"LLM thinking: {thinking[:200]}...")
             logger.info(f"Input Tokens: {inputtokens}")
             logger.info(f"Output Tokens: {outputtokens}")
 
             resp = LlmResult(
-                text = resp.choices[0].message.content,
+                text = content,
                 in_token = inputtokens,
                 out_token = outputtokens,
                 model = model_name
@@ -136,9 +165,7 @@ class Processor(LlmService):
         Stream content generation from OpenAI.
         Yields LlmChunk objects with is_final=True on the last chunk.
         """
-        # Use provided model or fall back to default
         model_name = model or self.default_model
-        # Use provided temperature or fall back to default
         effective_temperature = temperature if temperature is not None else self.temperature
 
         logger.debug(f"Using model (streaming): {model_name}")
@@ -147,6 +174,8 @@ class Processor(LlmService):
         prompt = system + "\n\n" + prompt
 
         try:
+            api_kwargs = self._build_kwargs(model_name, effective_temperature)
+
             response = self.openai.chat.completions.create(
                 model=model_name,
                 messages=[
@@ -160,16 +189,14 @@ class Processor(LlmService):
                         ]
                     }
                 ],
-                temperature=effective_temperature,
-                max_completion_tokens=self.max_output,
                 stream=True,
-                stream_options={"include_usage": True}
+                stream_options={"include_usage": True},
+                **api_kwargs,
             )
 
             total_input_tokens = 0
             total_output_tokens = 0
 
-            # Stream chunks
             for chunk in response:
                 if chunk.choices and chunk.choices[0].delta.content:
                     yield LlmChunk(
@@ -254,6 +281,20 @@ class Processor(LlmService):
             help=f'LLM max output tokens (default: {default_max_output})'
         )
 
+        parser.add_argument(
+            '--thinking',
+            choices=["off", "low", "medium", "high"],
+            default=default_thinking,
+            help=f'Thinking/reasoning effort level (default: {default_thinking})'
+        )
+
+        parser.add_argument(
+            '--variant',
+            choices=sorted(VARIANTS.keys()),
+            default=DEFAULT_VARIANT,
+            help=f'API variant (default: {DEFAULT_VARIANT})'
+        )
+
 def run():
-    
+
     Processor.launch(default_ident, __doc__)

trustgraph_flow-2.6.5/trustgraph/model/text_completion/openai/variants.py

@@ -0,0 +1,176 @@
+"""
+OpenAI API variant profiles.
+
+Different providers expose OpenAI-compatible APIs with subtle differences
+in parameter names, thinking/reasoning support, and temperature handling.
+Each variant encapsulates those quirks so the processor doesn't need
+provider-specific conditionals.
+"""
+
+import re
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class Variant:
+    """Base variant — defines the interface all variants implement."""
+
+    name = None
+    token_param = "max_completion_tokens"
+    temperature_with_thinking = False
+
+    def completion_kwargs(self, max_output, temperature, thinking):
+        """Build provider-specific kwargs for chat.completions.create().
+
+        Parameters
+        ----------
+        max_output : int
+            Configured max output tokens.
+        temperature : float
+            Configured temperature.
+        thinking : str
+            Thinking effort level: "off", "low", "medium", "high".
+
+        Returns
+        -------
+        dict
+            Extra kwargs to spread into the API call.
+        """
+        kwargs = {self.token_param: max_output}
+
+        if thinking != "off":
+            kwargs.update(self.thinking_kwargs(thinking))
+            if not self.temperature_with_thinking:
+                kwargs["temperature"] = 1.0
+            else:
+                kwargs["temperature"] = temperature
+        else:
+            kwargs["temperature"] = temperature
+
+        return kwargs
+
+    def thinking_kwargs(self, effort):
+        """Return kwargs to enable thinking at the given effort level."""
+        return {}
+
+    def extract_thinking(self, message):
+        """Extract thinking/reasoning content from a response message."""
+        return getattr(message, "reasoning_content", None)
+
+    def extract_thinking_stream(self, delta):
+        """Extract thinking content from a streaming delta."""
+        return getattr(delta, "reasoning_content", None)
+
+
+class OpenAIVariant(Variant):
+    """Standard OpenAI API (GPT-4o, o1, o3, etc.)."""
+
+    name = "openai"
+    token_param = "max_completion_tokens"
+    temperature_with_thinking = False
+
+    def thinking_kwargs(self, effort):
+        return {"reasoning_effort": effort}
+
+
+class DeepSeekVariant(Variant):
+    """DeepSeek API (R1, V3, etc.)."""
+
+    name = "deepseek"
+    token_param = "max_completion_tokens"
+    temperature_with_thinking = True
+
+    def completion_kwargs(self, max_output, temperature, thinking):
+        enabled = "enabled" if thinking != "off" else "disabled"
+        kwargs = {
+            self.token_param: max_output,
+            "temperature": temperature,
+            "extra_body": {
+                "thinking": {"type": enabled},
+            },
+        }
+        return kwargs
+
+    def thinking_kwargs(self, effort):
+        return {}
+
+
+class QwenVariant(Variant):
+    """Qwen / Alibaba Cloud API."""
+
+    name = "qwen"
+    token_param = "max_completion_tokens"
+    temperature_with_thinking = True
+
+    def completion_kwargs(self, max_output, temperature, thinking):
+        enabled = thinking != "off"
+        kwargs = {
+            self.token_param: max_output,
+            "temperature": temperature,
+            "extra_body": {
+                "enable_thinking": enabled,
+            },
+        }
+        return kwargs
+
+    def thinking_kwargs(self, effort):
+        return {}
+
+
+class MistralVariant(Variant):
+    """Mistral API (Mistral Large, etc.)."""
+
+    name = "mistral"
+    token_param = "max_tokens"
+    temperature_with_thinking = False
+
+    def thinking_kwargs(self, effort):
+        return {"reasoning_effort": effort}
+
+
+class LlamaVariant(Variant):
+    """Llama models via OpenAI-compatible servers (vLLM, Ollama, etc.).
+
+    Thinking is typically always-on or always-off depending on the model.
+    When present, thinking appears inline as <think>...</think> tags.
+    """
+
+    name = "llama"
+    token_param = "max_tokens"
+    temperature_with_thinking = True
+
+    def thinking_kwargs(self, effort):
+        return {}
+
+    def extract_thinking(self, message):
+        content = message.content or ""
+        match = re.search(r"<think>(.*?)</think>", content, re.DOTALL)
+        return match.group(1).strip() if match else None
+
+    def extract_content(self, message):
+        """Strip think tags from visible content."""
+        content = message.content or ""
+        return re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL).strip()
+
+
+VARIANTS = {
+    "openai": OpenAIVariant,
+    "deepseek": DeepSeekVariant,
+    "qwen": QwenVariant,
+    "mistral": MistralVariant,
+    "llama": LlamaVariant,
+}
+
+DEFAULT_VARIANT = "openai"
+
+
+def get_variant(name):
+    """Look up a variant by name, raising ValueError if unknown."""
+    cls = VARIANTS.get(name)
+    if cls is None:
+        raise ValueError(
+            f"Unknown variant {name!r}. "
+            f"Available: {', '.join(sorted(VARIANTS))}"
+        )
+    return cls()

trustgraph_flow-2.6.5/trustgraph/reranker/flashrank/__init__.py

@@ -0,0 +1,2 @@
+
+from . processor import *

trustgraph_flow-2.6.5/trustgraph/reranker/flashrank/__main__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+
+from . processor import run
+
+if __name__ == '__main__':
+    run()

trustgraph_flow-2.6.5/trustgraph/reranker/flashrank/processor.py

@@ -0,0 +1,109 @@
+
+"""
+Reranker service using flashrank.
+Scores query-document pairs and returns the top results ranked by
+relevance.
+"""
+
+import asyncio
+import logging
+
+from ... base import RerankerService
+from ... schema import RerankerResult
+
+from flashrank import Ranker, RerankRequest
+
+logger = logging.getLogger(__name__)
+
+default_ident = "reranker"
+
+default_model = "ms-marco-MiniLM-L-12-v2"
+
+class Processor(RerankerService):
+
+    def __init__(self, **params):
+
+        model = params.get("model", default_model)
+
+        super(Processor, self).__init__(
+            **params | { "model": model }
+        )
+
+        self.default_model = model
+
+        self.cached_model_name = None
+        self.ranker = None
+
+        self._load_model(model)
+
+    def _load_model(self, model_name):
+        if self.cached_model_name != model_name:
+            logger.info(f"Loading flashrank model: {model_name}")
+            self.ranker = Ranker(model_name=model_name)
+            self.cached_model_name = model_name
+            logger.info(f"flashrank model {model_name} loaded successfully")
+        else:
+            logger.debug(f"Using cached model: {model_name}")
+
+    def _run_rerank(self, query, passages):
+        request = RerankRequest(query=query, passages=passages)
+        return self.ranker.rerank(request)
+
+    async def on_rerank(self, queries, documents, limit, model=None):
+
+        if not queries or not documents:
+            return []
+
+        use_model = model or self.default_model
+
+        if self.cached_model_name != use_model:
+            await asyncio.to_thread(self._load_model, use_model)
+
+        passages = [
+            {"id": d.document_id, "text": d.document_text}
+            for d in documents
+        ]
+
+        best_scores = {}
+
+        for q in queries:
+            ranked = await asyncio.to_thread(
+                self._run_rerank, q.query_text, passages,
+            )
+
+            for r in ranked:
+                doc_id = r["id"]
+                score = r["score"]
+                score = float(score)
+                if doc_id not in best_scores or score > best_scores[doc_id][1]:
+                    best_scores[doc_id] = (q.query_id, score)
+
+        results = sorted(
+            best_scores.items(),
+            key=lambda x: x[1][1],
+            reverse=True,
+        )[:limit]
+
+        return [
+            RerankerResult(
+                document_id=doc_id,
+                query_id=query_id,
+                score=score,
+            )
+            for doc_id, (query_id, score) in results
+        ]
+
+    @staticmethod
+    def add_args(parser):
+
+        RerankerService.add_args(parser)
+
+        parser.add_argument(
+            '-m', '--model',
+            default=default_model,
+            help=f'Reranker model (default: {default_model})'
+        )
+
+def run():
+
+    Processor.launch(default_ident, __doc__)