trustgraph-flow 0.12.5__tar.gz → 0.13.1__tar.gz

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/PKG-INFO

@@ -1,12 +1,12 @@
 Metadata-Version: 2.1
 Name: trustgraph-flow
-Version: 0.12.5
+Version: 0.13.1
 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
 Author: trustgraph.ai
 Author-email: security@trustgraph.ai
 License: UNKNOWN
-Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.12.5.tar.gz
+Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.13.1.tar.gz
 Platform: UNKNOWN
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/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>=0.12,<0.13",
+        "trustgraph-base>=0.13,<0.14",
         "urllib3",
         "rdflib",
         "pymilvus",

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/chunking/recursive/chunker.py

@@ -63,16 +63,8 @@ class Processor(ConsumerProducer):
 
         for ix, chunk in enumerate(texts):
 
-            id = v.metadata.id + "-c" + str(ix)
-
             r = Chunk(
-                metadata=Metadata(
-                    source=v.metadata.source,
-                    id=id,
-                    title=v.metadata.title,
-                    user=v.metadata.user,
-                    collection=v.metadata.collection,
-                ),
+                metadata=v.metadata,
                 chunk=chunk.page_content.encode("utf-8"),
             )
 

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/chunking/token/chunker.py

@@ -62,16 +62,8 @@ class Processor(ConsumerProducer):
 
         for ix, chunk in enumerate(texts):
 
-            id = v.metadata.id + "-c" + str(ix)
-
             r = Chunk(
-                metadata=Metadata(
-                    source=v.metadata.source,
-                    id=id,
-                    title=v.metadata.title,
-                    user=v.metadata.user,
-                    collection=v.metadata.collection,
-                ),
+                metadata=v.metadata,
                 chunk=chunk.page_content.encode("utf-8"),
             )
 

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/decoding/pdf/pdf_decoder.py

@@ -59,15 +59,8 @@ class Processor(ConsumerProducer):
 
                 for ix, page in enumerate(pages):
 
-                    id = v.metadata.id + "-p" + str(ix)
                     r = TextDocument(
-                        metadata=Metadata(
-                            source=v.metadata.source,
-                            title=v.metadata.title,
-                            id=id,
-                            user=v.metadata.user,
-                            collection=v.metadata.collection,
-                        ),
+                        metadata=v.metadata,
                         text=page.page_content.encode("utf-8"),
                     )
 

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/extract/kg/definitions/extract.py

@@ -7,16 +7,18 @@ get entity definitions which are output as graph edges.
 import urllib.parse
 import json
 
-from .... schema import ChunkEmbeddings, Triple, Metadata, Value
+from .... schema import ChunkEmbeddings, Triple, Triples, Metadata, Value
 from .... schema import chunk_embeddings_ingest_queue, triples_store_queue
 from .... schema import prompt_request_queue
 from .... schema import prompt_response_queue
 from .... log_level import LogLevel
 from .... clients.prompt_client import PromptClient
-from .... rdf import TRUSTGRAPH_ENTITIES, DEFINITION
+from .... rdf import TRUSTGRAPH_ENTITIES, DEFINITION, RDF_LABEL, SUBJECT_OF
 from .... base import ConsumerProducer
 
 DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True)
+RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
+SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True)
 
 module = ".".join(__name__.split(".")[1:-1])
 
@@ -44,7 +46,7 @@ class Processor(ConsumerProducer):
                 "output_queue": output_queue,
                 "subscriber": subscriber,
                 "input_schema": ChunkEmbeddings,
-                "output_schema": Triple,
+                "output_schema": Triples,
                 "prompt_request_queue": pr_request_queue,
                 "prompt_response_queue": pr_response_queue,
             }
@@ -69,9 +71,12 @@ class Processor(ConsumerProducer):
 
         return self.prompt.request_definitions(chunk)
 
-    def emit_edge(self, metadata, s, p, o):
+    def emit_edges(self, metadata, triples):
 
-        t = Triple(metadata=metadata, s=s, p=p, o=o)
+        t = Triples(
+            metadata=metadata,
+            triples=triples,
+        )
         self.producer.send(t)
 
     def handle(self, msg):
@@ -85,6 +90,13 @@ class Processor(ConsumerProducer):
 
             defs = self.get_definitions(chunk)
 
+            triples = []
+
+            # FIXME: Putting metadata into triples store is duplicated in
+            # relationships extractor too
+            for t in v.metadata.metadata:
+                triples.append(t)
+
             for defn in defs:
 
                 s = defn.name
@@ -101,7 +113,31 @@ class Processor(ConsumerProducer):
                 s_value = Value(value=str(s_uri), is_uri=True)
                 o_value = Value(value=str(o), is_uri=False)
 
-                self.emit_edge(v.metadata, s_value, DEFINITION_VALUE, o_value)
+                triples.append(Triple(
+                    s=s_value,
+                    p=RDF_LABEL_VALUE,
+                    o=Value(value=s, is_uri=False),
+                ))
+
+                triples.append(Triple(
+                    s=s_value, p=DEFINITION_VALUE, o=o_value
+                ))
+
+                triples.append(Triple(
+                    s=s_value,
+                    p=SUBJECT_OF_VALUE,
+                    o=Value(value=v.metadata.id, is_uri=True)
+                ))
+
+            self.emit_edges(
+                Metadata(
+                    id=v.metadata.id,
+                    metadata=[],
+                    user=v.metadata.user,
+                    collection=v.metadata.collection,
+                ),
+                triples
+            )
 
         except Exception as e:
             print("Exception: ", e, flush=True)

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/extract/kg/relationships/extract.py

@@ -9,7 +9,7 @@ import urllib.parse
 import os
 from pulsar.schema import JsonSchema
 
-from .... schema import ChunkEmbeddings, Triple, GraphEmbeddings
+from .... schema import ChunkEmbeddings, Triple, Triples, GraphEmbeddings
 from .... schema import Metadata, Value
 from .... schema import chunk_embeddings_ingest_queue, triples_store_queue
 from .... schema import graph_embeddings_store_queue
@@ -17,10 +17,11 @@ from .... schema import prompt_request_queue
 from .... schema import prompt_response_queue
 from .... log_level import LogLevel
 from .... clients.prompt_client import PromptClient
-from .... rdf import RDF_LABEL, TRUSTGRAPH_ENTITIES
+from .... rdf import RDF_LABEL, TRUSTGRAPH_ENTITIES, SUBJECT_OF
 from .... base import ConsumerProducer
 
 RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
+SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True)
 
 module = ".".join(__name__.split(".")[1:-1])
 
@@ -50,7 +51,7 @@ class Processor(ConsumerProducer):
                 "output_queue": output_queue,
                 "subscriber": subscriber,
                 "input_schema": ChunkEmbeddings,
-                "output_schema": Triple,
+                "output_schema": Triples,
                 "prompt_request_queue": pr_request_queue,
                 "prompt_response_queue": pr_response_queue,
             }
@@ -69,7 +70,7 @@ class Processor(ConsumerProducer):
             "prompt_response_queue": pr_response_queue,
             "subscriber": subscriber,
             "input_schema": ChunkEmbeddings.__name__,
-            "output_schema": Triple.__name__,
+            "output_schema": Triples.__name__,
             "vector_schema": GraphEmbeddings.__name__,
         })
 
@@ -92,9 +93,12 @@ class Processor(ConsumerProducer):
 
         return self.prompt.request_relationships(chunk)
 
-    def emit_edge(self, metadata, s, p, o):
+    def emit_edges(self, metadata, triples):
 
-        t = Triple(metadata=metadata, s=s, p=p, o=o)
+        t = Triples(
+            metadata=metadata,
+            triples=triples,
+        )
         self.producer.send(t)
 
     def emit_vec(self, metadata, ent, vec):
@@ -113,6 +117,13 @@ class Processor(ConsumerProducer):
 
             rels = self.get_relationships(chunk)
 
+            triples = []
+
+            # FIXME: Putting metadata into triples store is duplicated in
+            # relationships extractor too
+            for t in v.metadata.metadata:
+                triples.append(t)
+
             for rel in rels:
 
                 s = rel.s
@@ -139,43 +150,65 @@ class Processor(ConsumerProducer):
                 else:
                     o_value = Value(value=str(o), is_uri=False)
 
-                self.emit_edge(
-                    v.metadata,
-                    s_value,
-                    p_value,
-                    o_value
-                )
+                triples.append(Triple(
+                    s=s_value,
+                    p=p_value,
+                    o=o_value
+                ))
 
                 # Label for s
-                self.emit_edge(
-                    v.metadata,
-                    s_value,
-                    RDF_LABEL_VALUE,
-                    Value(value=str(s), is_uri=False)
-                )
+                triples.append(Triple(
+                    s=s_value,
+                    p=RDF_LABEL_VALUE,
+                    o=Value(value=str(s), is_uri=False)
+                ))
 
                 # Label for p
-                self.emit_edge(
-                    v.metadata,
-                    p_value,
-                    RDF_LABEL_VALUE,
-                    Value(value=str(p), is_uri=False)
-                )
+                triples.append(Triple(
+                    s=p_value,
+                    p=RDF_LABEL_VALUE,
+                    o=Value(value=str(p), is_uri=False)
+                ))
 
                 if rel.o_entity:
                     # Label for o
-                    self.emit_edge(
-                        v.metadata,
-                        o_value,
-                        RDF_LABEL_VALUE,
-                        Value(value=str(o), is_uri=False)
-                    )
+                    triples.append(Triple(
+                        s=o_value,
+                        p=RDF_LABEL_VALUE,
+                        o=Value(value=str(o), is_uri=False)
+                    ))
+
+                # 'Subject of' for s
+                triples.append(Triple(
+                    s=s_value,
+                    p=SUBJECT_OF_VALUE,
+                    o=Value(value=v.metadata.id, is_uri=True)
+                ))
+
+                if rel.o_entity:
+                    # 'Subject of' for o
+                    triples.append(Triple(
+                        s=o_value,
+                        p=RDF_LABEL_VALUE,
+                        o=Value(value=v.metadata.id, is_uri=True)
+                    ))
 
                 self.emit_vec(v.metadata, s_value, v.vectors)
                 self.emit_vec(v.metadata, p_value, v.vectors)
+
                 if rel.o_entity:
                     self.emit_vec(v.metadata, o_value, v.vectors)
 
+            self.emit_edges(
+                Metadata(
+                    id=v.metadata.id,
+                    metadata=[],
+                    user=v.metadata.user,
+                    collection=v.metadata.collection,
+                ),
+                triples
+            )
+
         except Exception as e:
             print("Exception: ", e, flush=True)
 

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/extract/kg/topics/extract.py

@@ -7,7 +7,7 @@ get entity definitions which are output as graph edges.
 import urllib.parse
 import json
 
-from .... schema import ChunkEmbeddings, Triple, Metadata, Value
+from .... schema import ChunkEmbeddings, Triple, Triples, Metadata, Value
 from .... schema import chunk_embeddings_ingest_queue, triples_store_queue
 from .... schema import prompt_request_queue
 from .... schema import prompt_response_queue
@@ -44,7 +44,7 @@ class Processor(ConsumerProducer):
                 "output_queue": output_queue,
                 "subscriber": subscriber,
                 "input_schema": ChunkEmbeddings,
-                "output_schema": Triple,
+                "output_schema": Triples,
                 "prompt_request_queue": pr_request_queue,
                 "prompt_response_queue": pr_response_queue,
             }
@@ -71,7 +71,10 @@ class Processor(ConsumerProducer):
 
     def emit_edge(self, metadata, s, p, o):
 
-        t = Triple(metadata=metadata, s=s, p=p, o=o)
+        t = Triples(
+            metadata=metadata,
+            triples=[Triple(s=s, p=p, o=o)],
+        )
         self.producer.send(t)
 
     def handle(self, msg):

trustgraph-flow-0.13.1/trustgraph/flow_version.py

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

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/model/text_completion/azure/llm.py

@@ -7,6 +7,7 @@ serverless endpoint service.  Input is prompt, output is response.
 import requests
 import json
 from prometheus_client import Histogram
+import os
 
 from .... schema import TextCompletionRequest, TextCompletionResponse, Error
 from .... schema import text_completion_request_queue
@@ -23,6 +24,8 @@ default_subscriber = module
 default_temperature = 0.0
 default_max_output = 4192
 default_model = "AzureAI"
+default_endpoint = os.getenv("AZURE_ENDPOINT")
+default_token = os.getenv("AZURE_TOKEN")
 
 class Processor(ConsumerProducer):
 
@@ -31,12 +34,18 @@ class Processor(ConsumerProducer):
         input_queue = params.get("input_queue", default_input_queue)
         output_queue = params.get("output_queue", default_output_queue)
         subscriber = params.get("subscriber", default_subscriber)
-        endpoint = params.get("endpoint")
-        token = params.get("token")
+        endpoint = params.get("endpoint", default_endpoint)
+        token = params.get("token", default_token)
         temperature = params.get("temperature", default_temperature)
         max_output = params.get("max_output", default_max_output)
         model = default_model
 
+        if endpoint is None:
+            raise RuntimeError("Azure endpoint not specified")
+
+        if token is None:
+            raise RuntimeError("Azure token not specified")
+
         super(Processor, self).__init__(
             **params | {
                 "input_queue": input_queue,
@@ -199,11 +208,13 @@ class Processor(ConsumerProducer):
 
         parser.add_argument(
             '-e', '--endpoint',
+            default=default_endpoint,
             help=f'LLM model endpoint'
         )
 
         parser.add_argument(
             '-k', '--token',
+            default=default_token,
             help=f'LLM model token'
         )
 

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/model/text_completion/azure_openai/llm.py

@@ -8,6 +8,7 @@ import requests
 import json
 from prometheus_client import Histogram
 from openai import AzureOpenAI
+import os
 
 from .... schema import TextCompletionRequest, TextCompletionResponse, Error
 from .... schema import text_completion_request_queue
@@ -24,6 +25,8 @@ default_subscriber = module
 default_temperature = 0.0
 default_max_output = 4192
 default_api = "2024-02-15-preview"
+default_endpoint = os.getenv("AZURE_ENDPOINT")
+default_token = os.getenv("AZURE_TOKEN")
 
 class Processor(ConsumerProducer):
 
@@ -32,13 +35,19 @@ class Processor(ConsumerProducer):
         input_queue = params.get("input_queue", default_input_queue)
         output_queue = params.get("output_queue", default_output_queue)
         subscriber = params.get("subscriber", default_subscriber)
-        endpoint = params.get("endpoint")
-        token = params.get("token")
+        endpoint = params.get("endpoint", default_endpoint)
+        token = params.get("token", default_token)
         temperature = params.get("temperature", default_temperature)
         max_output = params.get("max_output", default_max_output)
         model = params.get("model")
         api = params.get("api_version", default_api)
 
+        if endpoint is None:
+            raise RuntimeError("Azure endpoint not specified")
+
+        if token is None:
+            raise RuntimeError("Azure token not specified")
+
         super(Processor, self).__init__(
             **params | {
                 "input_queue": input_queue,

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/model/text_completion/claude/llm.py

@@ -6,6 +6,7 @@ Input is prompt, output is response.
 
 import anthropic
 from prometheus_client import Histogram
+import os
 
 from .... schema import TextCompletionRequest, TextCompletionResponse, Error
 from .... schema import text_completion_request_queue
@@ -22,6 +23,7 @@ default_subscriber = module
 default_model = 'claude-3-5-sonnet-20240620'
 default_temperature = 0.0
 default_max_output = 8192
+default_api_key = os.getenv("CLAUDE_KEY")
 
 class Processor(ConsumerProducer):
 
@@ -31,10 +33,13 @@ class Processor(ConsumerProducer):
         output_queue = params.get("output_queue", default_output_queue)
         subscriber = params.get("subscriber", default_subscriber)
         model = params.get("model", default_model)
-        api_key = params.get("api_key")
+        api_key = params.get("api_key", default_api_key)
         temperature = params.get("temperature", default_temperature)
         max_output = params.get("max_output", default_max_output)
 
+        if api_key is None:
+            raise RuntimeError("Claude API key not specified")
+
         super(Processor, self).__init__(
             **params | {
                 "input_queue": input_queue,
@@ -175,6 +180,7 @@ class Processor(ConsumerProducer):
 
         parser.add_argument(
             '-k', '--api-key',
+            default=default_api_key,
             help=f'Claude API key'
         )
 

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/model/text_completion/cohere/llm.py

@@ -6,6 +6,7 @@ Input is prompt, output is response.
 
 import cohere
 from prometheus_client import Histogram
+import os
 
 from .... schema import TextCompletionRequest, TextCompletionResponse, Error
 from .... schema import text_completion_request_queue
@@ -21,6 +22,7 @@ default_output_queue = text_completion_response_queue
 default_subscriber = module
 default_model = 'c4ai-aya-23-8b'
 default_temperature = 0.0
+default_api_key = os.getenv("COHERE_KEY")
 
 class Processor(ConsumerProducer):
 
@@ -30,9 +32,12 @@ class Processor(ConsumerProducer):
         output_queue = params.get("output_queue", default_output_queue)
         subscriber = params.get("subscriber", default_subscriber)
         model = params.get("model", default_model)
-        api_key = params.get("api_key")
+        api_key = params.get("api_key", default_api_key)
         temperature = params.get("temperature", default_temperature)
 
+        if api_key is None:
+            raise RuntimeError("Cohere API key not specified")
+
         super(Processor, self).__init__(
             **params | {
                 "input_queue": input_queue,
@@ -162,6 +167,7 @@ class Processor(ConsumerProducer):
 
         parser.add_argument(
             '-k', '--api-key',
+            default=default_api_key,
             help=f'Cohere API key'
         )
 

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/model/text_completion/googleaistudio/llm.py

@@ -7,6 +7,7 @@ Input is prompt, output is response.
 import google.generativeai as genai
 from google.generativeai.types import HarmCategory, HarmBlockThreshold
 from prometheus_client import Histogram
+import os
 
 from .... schema import TextCompletionRequest, TextCompletionResponse, Error
 from .... schema import text_completion_request_queue
@@ -23,6 +24,7 @@ default_subscriber = module
 default_model = 'gemini-1.5-flash-002'
 default_temperature = 0.0
 default_max_output = 8192
+default_api_key = os.getenv("GOOGLE_AI_STUDIO_KEY")
 
 class Processor(ConsumerProducer):
 
@@ -32,10 +34,13 @@ class Processor(ConsumerProducer):
         output_queue = params.get("output_queue", default_output_queue)
         subscriber = params.get("subscriber", default_subscriber)
         model = params.get("model", default_model)
-        api_key = params.get("api_key")
+        api_key = params.get("api_key", default_api_key)
         temperature = params.get("temperature", default_temperature)
         max_output = params.get("max_output", default_max_output)
 
+        if api_key is None:
+            raise RuntimeError("Google AI Studio API key not specified")
+
         super(Processor, self).__init__(
             **params | {
                 "input_queue": input_queue,
@@ -190,6 +195,7 @@ class Processor(ConsumerProducer):
 
         parser.add_argument(
             '-k', '--api-key',
+            default=default_api_key,
             help=f'GoogleAIStudio API key'
         )
 
@@ -211,4 +217,4 @@ def run():
 
     Processor.start(module, __doc__)
 
-    
+    

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/model/text_completion/llamafile/llm.py

@@ -20,7 +20,7 @@ default_input_queue = text_completion_request_queue
 default_output_queue = text_completion_response_queue
 default_subscriber = module
 default_model = 'LLaMA_CPP'
-default_llamafile = 'http://localhost:8080/v1'
+default_llamafile = os.getenv("LLAMAFILE_URL", "http://localhost:8080/v1")
 default_temperature = 0.0
 default_max_output = 4096
 

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/model/text_completion/ollama/llm.py

@@ -6,6 +6,7 @@ Input is prompt, output is response.
 
 from ollama import Client
 from prometheus_client import Histogram, Info
+import os
 
 from .... schema import TextCompletionRequest, TextCompletionResponse, Error
 from .... schema import text_completion_request_queue
@@ -19,8 +20,8 @@ module = ".".join(__name__.split(".")[1:-1])
 default_input_queue = text_completion_request_queue
 default_output_queue = text_completion_response_queue
 default_subscriber = module
-default_model = 'gemma2'
-default_ollama = 'http://localhost:11434'
+default_model = 'gemma2:9b'
+default_ollama = os.getenv("OLLAMA_HOST", 'http://localhost:11434')
 
 class Processor(ConsumerProducer):
 
@@ -152,7 +153,7 @@ class Processor(ConsumerProducer):
         parser.add_argument(
             '-m', '--model',
             default="gemma2",
-            help=f'LLM model (default: gemma2)'
+            help=f'LLM model (default: {default_model})'
         )
 
         parser.add_argument(

{trustgraph-flow-0.12.5 → trustgraph-flow-0.13.1}/trustgraph/model/text_completion/openai/llm.py

@@ -6,6 +6,7 @@ Input is prompt, output is response.
 
 from openai import OpenAI
 from prometheus_client import Histogram
+import os
 
 from .... schema import TextCompletionRequest, TextCompletionResponse, Error
 from .... schema import text_completion_request_queue
@@ -22,6 +23,7 @@ default_subscriber = module
 default_model = 'gpt-3.5-turbo'
 default_temperature = 0.0
 default_max_output = 4096
+default_api_key = os.getenv("OPENAI_KEY")
 
 class Processor(ConsumerProducer):
 
@@ -31,10 +33,13 @@ class Processor(ConsumerProducer):
         output_queue = params.get("output_queue", default_output_queue)
         subscriber = params.get("subscriber", default_subscriber)
         model = params.get("model", default_model)
-        api_key = params.get("api_key")
+        api_key = params.get("api_key", default_api_key)
         temperature = params.get("temperature", default_temperature)
         max_output = params.get("max_output", default_max_output)
 
+        if api_key is None:
+            raise RuntimeError("OpenAI API key not specified")
+
         super(Processor, self).__init__(
             **params | {
                 "input_queue": input_queue,
@@ -185,6 +190,7 @@ class Processor(ConsumerProducer):
 
         parser.add_argument(
             '-k', '--api-key',
+            default=default_api_key,
             help=f'OpenAI API key'
         )