trustgraph-flow 1.0.12__tar.gz → 1.0.14__tar.gz

{trustgraph-flow-1.0.12 → trustgraph-flow-1.0.14}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: trustgraph-flow
-Version: 1.0.12
+Version: 1.0.14
 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.0.12.tar.gz
+Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v1.0.14.tar.gz
 Author: trustgraph.ai
 Author-email: security@trustgraph.ai
 Classifier: Programming Language :: Python :: 3

trustgraph-flow-1.0.14/scripts/text-completion-vllm

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

{trustgraph-flow-1.0.12 → trustgraph-flow-1.0.14}/setup.py

@@ -118,6 +118,7 @@ setuptools.setup(
         "scripts/text-completion-ollama",
         "scripts/text-completion-openai",
         "scripts/text-completion-tgi",
+        "scripts/text-completion-vllm",
         "scripts/triples-query-cassandra",
         "scripts/triples-query-falkordb",
         "scripts/triples-query-memgraph",

{trustgraph-flow-1.0.12 → trustgraph-flow-1.0.14}/trustgraph/extract/kg/definitions/extract.py

@@ -21,16 +21,19 @@ RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
 SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True)
 
 default_ident = "kg-extract-definitions"
+default_concurrency = 1
 
 class Processor(FlowProcessor):
 
     def __init__(self, **params):
 
         id = params.get("id")
+        concurrency = params.get("concurrency", 1)
 
         super(Processor, self).__init__(
             **params | {
                 "id": id,
+                "concurrency": concurrency,
             }
         )
 
@@ -38,7 +41,8 @@ class Processor(FlowProcessor):
             ConsumerSpec(
                 name = "input",
                 schema = Chunk,
-                handler = self.on_message
+                handler = self.on_message,
+                concurrency = concurrency,
             )
         )
 
@@ -190,6 +194,13 @@ class Processor(FlowProcessor):
     @staticmethod
     def add_args(parser):
 
+        parser.add_argument(
+            '-c', '--concurrency',
+            type=int,
+            default=default_concurrency,
+            help=f'Concurrent processing threads (default: {default_concurrency})'
+        )
+
         FlowProcessor.add_args(parser)
 
 def run():

{trustgraph-flow-1.0.12 → trustgraph-flow-1.0.14}/trustgraph/extract/kg/relationships/extract.py

@@ -20,16 +20,19 @@ RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
 SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True)
 
 default_ident = "kg-extract-relationships"
+default_concurrency = 1
 
 class Processor(FlowProcessor):
 
     def __init__(self, **params):
 
         id = params.get("id")
+        concurrency = params.get("concurrency", 1)
 
         super(Processor, self).__init__(
             **params | {
                 "id": id,
+                "concurrency": concurrency,
             }
         )
 
@@ -37,7 +40,8 @@ class Processor(FlowProcessor):
             ConsumerSpec(
                 name = "input",
                 schema = Chunk,
-                handler = self.on_message
+                handler = self.on_message,
+                concurrency = concurrency,
             )
         )
 
@@ -192,6 +196,13 @@ class Processor(FlowProcessor):
     @staticmethod
     def add_args(parser):
 
+        parser.add_argument(
+            '-c', '--concurrency',
+            type=int,
+            default=default_concurrency,
+            help=f'Concurrent processing threads (default: {default_concurrency})'
+        )
+
         FlowProcessor.add_args(parser)
 
 def run():

trustgraph-flow-1.0.14/trustgraph/flow_version.py

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

{trustgraph-flow-1.0.12 → trustgraph-flow-1.0.14}/trustgraph/model/prompt/template/service.py

@@ -18,12 +18,14 @@ from .... base import ProducerSpec, ConsumerSpec, TextCompletionClientSpec
 from . prompt_manager import PromptConfiguration, Prompt, PromptManager
 
 default_ident = "prompt"
+default_concurrency = 1
 
 class Processor(FlowProcessor):
 
     def __init__(self, **params):
 
         id = params.get("id")
+        concurrency = params.get("concurrency", 1)
 
         # Config key for prompts
         self.config_key = params.get("config_type", "prompt")
@@ -31,6 +33,7 @@ class Processor(FlowProcessor):
         super(Processor, self).__init__(
             **params | {
                 "id": id,
+                "concurrency": concurrency,
             }
         )
 
@@ -38,7 +41,8 @@ class Processor(FlowProcessor):
             ConsumerSpec(
                 name = "request",
                 schema = PromptRequest,
-                handler = self.on_request
+                handler = self.on_request,
+                concurrency = concurrency,
             )
         )
 
@@ -219,6 +223,13 @@ class Processor(FlowProcessor):
     @staticmethod
     def add_args(parser):
 
+        parser.add_argument(
+            '-c', '--concurrency',
+            type=int,
+            default=default_concurrency,
+            help=f'Concurrent processing threads (default: {default_concurrency})'
+        )
+
         FlowProcessor.add_args(parser)
 
         parser.add_argument(

trustgraph-flow-1.0.14/trustgraph/model/text_completion/vllm/__init__.py

@@ -0,0 +1,3 @@
+
+from . llm import *
+

trustgraph-flow-1.0.14/trustgraph/model/text_completion/vllm/__main__.py

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

trustgraph-flow-1.0.14/trustgraph/model/text_completion/vllm/llm.py

@@ -0,0 +1,138 @@
+
+"""
+Simple LLM service, performs text prompt completion using vLLM
+Input is prompt, output is response.
+"""
+
+import os
+import aiohttp
+
+from .... exceptions import TooManyRequests
+from .... base import LlmService, LlmResult
+
+default_ident = "text-completion"
+
+default_temperature = 0.0
+default_max_output = 2048
+default_base_url = os.getenv("VLLM_BASE_URL")
+default_model = "TheBloke/Mistral-7B-v0.1-AWQ"
+
+if default_base_url == "" or default_base_url is None:
+    default_base_url = "http://vllm-service:8899/v1"
+
+class Processor(LlmService):
+
+    def __init__(self, **params):
+    
+        base_url = params.get("url", default_base_url)
+        temperature = params.get("temperature", default_temperature)
+        max_output = params.get("max_output", default_max_output)
+        model = params.get("model", default_model)
+
+        super(Processor, self).__init__(
+            **params | {
+                "temperature": temperature,
+                "max_output": max_output,
+                "url": base_url,
+                "model": model,
+            }
+        )
+
+        self.base_url = base_url
+        self.temperature = temperature
+        self.max_output = max_output
+        self.model = model
+
+        self.session = aiohttp.ClientSession()
+
+        print("Using vLLM service at", base_url)
+
+        print("Initialised", flush=True)
+
+    async def generate_content(self, system, prompt):
+
+        headers = {
+            "Content-Type": "application/json",
+        }
+
+        request = {
+            "model": self.model,
+            "prompt": system + "\n\n" + prompt,
+            "max_tokens": self.max_output,
+            "temperature": self.temperature,
+        }            
+
+        try:
+
+            url = f"{self.base_url}/completions"
+
+            async with self.session.post(
+                    url,
+                    headers=headers,
+                    json=request,
+            ) as response:
+
+                if response.status != 200:
+                    raise RuntimeError("Bad status: " + str(response.status))
+
+                resp = await response.json()
+            
+            inputtokens = resp["usage"]["prompt_tokens"]
+            outputtokens = resp["usage"]["completion_tokens"]
+            ans = resp["choices"][0]["text"]
+            print(f"Input Tokens: {inputtokens}", flush=True)
+            print(f"Output Tokens: {outputtokens}", flush=True)
+            print(ans, flush=True)
+
+            resp = LlmResult(
+                text = ans,
+                in_token = inputtokens,
+                out_token = outputtokens,
+                model = self.model,
+            )
+
+            return resp
+
+        # FIXME: Assuming vLLM won't produce rate limits?
+
+        except Exception as e:
+
+            # Apart from rate limits, treat all exceptions as unrecoverable
+
+            print(f"Exception: {type(e)} {e}")
+            raise e
+
+    @staticmethod
+    def add_args(parser):
+
+        LlmService.add_args(parser)
+
+        parser.add_argument(
+            '-u', '--url',
+            default=default_base_url,
+            help=f'vLLM service base URL (default: {default_base_url})'
+        )
+
+        parser.add_argument(
+            '-m', '--model',
+            default=default_model,
+            help=f'LLM model (default: {default_model})'
+        )
+
+        parser.add_argument(
+            '-t', '--temperature',
+            type=float,
+            default=default_temperature,
+            help=f'LLM temperature parameter (default: {default_temperature})'
+        )
+
+        parser.add_argument(
+            '-x', '--max-output',
+            type=int,
+            default=default_max_output,
+            help=f'LLM max output tokens (default: {default_max_output})'
+        )
+
+def run():
+    
+    Processor.launch(default_ident, __doc__)

{trustgraph-flow-1.0.12 → trustgraph-flow-1.0.14}/trustgraph/retrieval/graph_rag/rag.py

@@ -11,12 +11,14 @@ from ... base import PromptClientSpec, EmbeddingsClientSpec
 from ... base import GraphEmbeddingsClientSpec, TriplesClientSpec
 
 default_ident = "graph-rag"
+default_concurrency = 1
 
 class Processor(FlowProcessor):
 
     def __init__(self, **params):
 
         id = params.get("id", default_ident)
+        concurrency = params.get("concurrency", 1)
 
         entity_limit = params.get("entity_limit", 50)
         triple_limit = params.get("triple_limit", 30)
@@ -26,6 +28,7 @@ class Processor(FlowProcessor):
         super(Processor, self).__init__(
             **params | {
                 "id": id,
+                "concurrency": concurrency,
                 "entity_limit": entity_limit,
                 "triple_limit": triple_limit,
                 "max_subgraph_size": max_subgraph_size,
@@ -43,6 +46,7 @@ class Processor(FlowProcessor):
                 name = "request",
                 schema = GraphRagQuery,
                 handler = self.on_request,
+                concurrency = concurrency,
             )
         )
 
@@ -157,6 +161,13 @@ class Processor(FlowProcessor):
     @staticmethod
     def add_args(parser):
 
+        parser.add_argument(
+            '-c', '--concurrency',
+            type=int,
+            default=default_concurrency,
+            help=f'Concurrent processing threads (default: {default_concurrency})'
+        )
+
         FlowProcessor.add_args(parser)
 
         parser.add_argument(

{trustgraph-flow-1.0.12 → trustgraph-flow-1.0.14}/trustgraph_flow.egg-info/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: trustgraph-flow
-Version: 1.0.12
+Version: 1.0.14
 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.0.12.tar.gz
+Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v1.0.14.tar.gz
 Author: trustgraph.ai
 Author-email: security@trustgraph.ai
 Classifier: Programming Language :: Python :: 3

{trustgraph-flow-1.0.12 → trustgraph-flow-1.0.14}/trustgraph_flow.egg-info/SOURCES.txt

@@ -49,6 +49,7 @@ scripts/text-completion-mistral
 scripts/text-completion-ollama
 scripts/text-completion-openai
 scripts/text-completion-tgi
+scripts/text-completion-vllm
 scripts/triples-query-cassandra
 scripts/triples-query-falkordb
 scripts/triples-query-memgraph
@@ -224,6 +225,9 @@ trustgraph/model/text_completion/openai/llm.py
 trustgraph/model/text_completion/tgi/__init__.py
 trustgraph/model/text_completion/tgi/__main__.py
 trustgraph/model/text_completion/tgi/llm.py
+trustgraph/model/text_completion/vllm/__init__.py
+trustgraph/model/text_completion/vllm/__main__.py
+trustgraph/model/text_completion/vllm/llm.py
 trustgraph/processing/__init__.py
 trustgraph/processing/__main__.py
 trustgraph/processing/processing.py

trustgraph-flow-1.0.12/trustgraph/flow_version.py

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