diff --git a/trustgraph/chunker/recursive/chunker.py b/trustgraph/chunker/recursive/chunker.py index b116eca5..5418f951 100755 --- a/trustgraph/chunker/recursive/chunker.py +++ b/trustgraph/chunker/recursive/chunker.py @@ -17,25 +17,22 @@ default_subscriber = 'chunker-recursive' class Processor(ConsumerProducer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - output_queue=default_output_queue, - subscriber=default_subscriber, - log_level=LogLevel.INFO, - chunk_size=2000, - chunk_overlap=100, - ): + def __init__(self, **params): + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) + chunk_size = params.get("chunk_size", 2000) + chunk_overlap = params.get("chunk_overlap", 100) + super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - output_queue=output_queue, - subscriber=subscriber, - input_schema=TextDocument, - output_schema=Chunk, + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": TextDocument, + "output_schema": Chunk, + } ) self.text_splitter = RecursiveCharacterTextSplitter( diff --git a/trustgraph/decoder/pdf/pdf_decoder.py b/trustgraph/decoder/pdf/pdf_decoder.py index e87f2905..1d86c366 100755 --- a/trustgraph/decoder/pdf/pdf_decoder.py +++ b/trustgraph/decoder/pdf/pdf_decoder.py @@ -18,23 +18,20 @@ default_subscriber = 'pdf-decoder' class Processor(ConsumerProducer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - output_queue=default_output_queue, - subscriber=default_subscriber, - log_level=LogLevel.INFO, - ): + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - output_queue=output_queue, - subscriber=subscriber, - input_schema=Document, - output_schema=TextDocument, + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": Document, + "output_schema": TextDocument, + } ) print("PDF inited") diff --git a/trustgraph/embeddings/hf/hf.py b/trustgraph/embeddings/hf/hf.py index 82273310..e5a5728f 100755 --- a/trustgraph/embeddings/hf/hf.py +++ b/trustgraph/embeddings/hf/hf.py @@ -17,24 +17,21 @@ default_model="all-MiniLM-L6-v2" class Processor(ConsumerProducer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - output_queue=default_output_queue, - subscriber=default_subscriber, - log_level=LogLevel.INFO, - model=default_model, - ): + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) + model = params.get("model", default_model) super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - output_queue=output_queue, - subscriber=subscriber, - input_schema=EmbeddingsRequest, - output_schema=EmbeddingsResponse, + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": EmbeddingsRequest, + "output_schema": EmbeddingsResponse, + } ) self.embeddings = HuggingFaceEmbeddings(model_name=model) diff --git a/trustgraph/embeddings/ollama/processor.py b/trustgraph/embeddings/ollama/processor.py index 5c36f86a..38c177b1 100755 --- a/trustgraph/embeddings/ollama/processor.py +++ b/trustgraph/embeddings/ollama/processor.py @@ -17,25 +17,20 @@ default_ollama = 'http://localhost:11434' class Processor(ConsumerProducer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - output_queue=default_output_queue, - subscriber=default_subscriber, - log_level=LogLevel.INFO, - model=default_model, - ollama=default_ollama, - ): + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - output_queue=output_queue, - subscriber=subscriber, - input_schema=EmbeddingsRequest, - output_schema=EmbeddingsResponse, + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": EmbeddingsRequest, + "output_schema": EmbeddingsResponse, + } ) self.embeddings = OllamaEmbeddings(base_url=ollama, model=model) diff --git a/trustgraph/embeddings/vectorize/vectorize.py b/trustgraph/embeddings/vectorize/vectorize.py index e779c6ed..2a92827e 100755 --- a/trustgraph/embeddings/vectorize/vectorize.py +++ b/trustgraph/embeddings/vectorize/vectorize.py @@ -15,26 +15,23 @@ default_subscriber = 'embeddings-vectorizer' class Processor(ConsumerProducer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - output_queue=default_output_queue, - subscriber=default_subscriber, - log_level=LogLevel.INFO, - ): + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - output_queue=output_queue, - subscriber=subscriber, - input_schema=Chunk, - output_schema=VectorsChunk, + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": Chunk, + "output_schema": VectorsChunk, + } ) - self.embeddings = EmbeddingsClient(pulsar_host=pulsar_host) + self.embeddings = EmbeddingsClient(pulsar_host=self.pulsar_host) def emit(self, source, chunk, vectors): diff --git a/trustgraph/graph/cassandra_write/write.py b/trustgraph/graph/cassandra_write/write.py index df317c2a..fe5864a0 100755 --- a/trustgraph/graph/cassandra_write/write.py +++ b/trustgraph/graph/cassandra_write/write.py @@ -20,27 +20,22 @@ default_graph_host='localhost' class Processor(Consumer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - subscriber=default_subscriber, - graph_host=default_graph_host, - log_level=LogLevel.INFO, - ): + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + subscriber = params.get("subscriber", default_subscriber) + graph_host = params.get("graph_host", default_graph_host) super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - subscriber=subscriber, - input_schema=Triple, + **params | { + "input_queue": input_queue, + "subscriber": subscriber, + "input_schema": Triple, + } ) self.tg = TrustGraph([graph_host]) - self.count = 0 - def handle(self, msg): v = msg.value() @@ -51,11 +46,6 @@ class Processor(Consumer): v.o.value ) - self.count += 1 - - if (self.count % 1000) == 0: - print(self.count, "...", flush=True) - @staticmethod def add_args(parser): diff --git a/trustgraph/kg/extract_definitions/extract.py b/trustgraph/kg/extract_definitions/extract.py index 42cbacc9..9d08470d 100755 --- a/trustgraph/kg/extract_definitions/extract.py +++ b/trustgraph/kg/extract_definitions/extract.py @@ -22,23 +22,20 @@ default_subscriber = 'kg-extract-definitions' class Processor(ConsumerProducer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - output_queue=default_output_queue, - subscriber=default_subscriber, - log_level=LogLevel.INFO, - ): + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - output_queue=output_queue, - subscriber=subscriber, - input_schema=VectorsChunk, - output_schema=Triple, + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": VectorsChunk, + "output_schema": Triple, + } ) self.llm = LlmClient(pulsar_host=pulsar_host) diff --git a/trustgraph/kg/extract_relationships/extract.py b/trustgraph/kg/extract_relationships/extract.py index adc0b71a..bf3b94c7 100755 --- a/trustgraph/kg/extract_relationships/extract.py +++ b/trustgraph/kg/extract_relationships/extract.py @@ -7,6 +7,7 @@ graph edges. import urllib.parse import json +import os from pulsar.schema import JsonSchema from ... schema import VectorsChunk, Triple, VectorsAssociation, Source, Value @@ -25,24 +26,21 @@ default_vector_queue='vectors-load' class Processor(ConsumerProducer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - vector_queue=default_vector_queue, - output_queue=default_output_queue, - subscriber=default_subscriber, - log_level=LogLevel.INFO, - ): + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + vector_queue = params.get("vector_queue", default_vector_queue) + subscriber = params.get("subscriber", default_subscriber) super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - output_queue=output_queue, - subscriber=subscriber, - input_schema=VectorsChunk, - output_schema=Triple, + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": VectorsChunk, + "output_schema": Triple, + } ) self.vec_prod = self.client.create_producer( @@ -50,7 +48,17 @@ class Processor(ConsumerProducer): schema=JsonSchema(VectorsAssociation), ) - self.llm = LlmClient(pulsar_host=pulsar_host) + __class__.pubsub_metric.info({ + "input_queue": input_queue, + "output_queue": output_queue, + "vector_queue": vector_queue, + "subscriber": subscriber, + "input_schema": VectorsChunk.__name__, + "output_schema": Triple.__name__, + "vector_schema": VectorsAssociation.__name__, + }) + + self.llm = LlmClient(pulsar_host=self.pulsar_host) def to_uri(self, text): diff --git a/trustgraph/llm/azure_text/llm.py b/trustgraph/llm/azure_text/llm.py index 12a7931e..bccbf04f 100755 --- a/trustgraph/llm/azure_text/llm.py +++ b/trustgraph/llm/azure_text/llm.py @@ -17,25 +17,22 @@ default_subscriber = 'llm-azure-text' class Processor(ConsumerProducer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - output_queue=default_output_queue, - subscriber=default_subscriber, - log_level=LogLevel.INFO, - endpoint=None, - token=None, - ): + def __init__(self, **params): + + 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") super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - output_queue=output_queue, - subscriber=subscriber, - input_schema=TextCompletionRequest, - output_schema=TextCompletionResponse, + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": TextCompletionRequest, + "output_schema": TextCompletionResponse, + } ) self.endpoint = endpoint diff --git a/trustgraph/llm/claude_text/llm.py b/trustgraph/llm/claude_text/llm.py index af516b71..230d540c 100755 --- a/trustgraph/llm/claude_text/llm.py +++ b/trustgraph/llm/claude_text/llm.py @@ -15,27 +15,25 @@ default_output_queue = 'llm-complete-text-response' default_subscriber = 'llm-claude-text' default_model = 'claude-3-5-sonnet-20240620' -class Processor: +class Processor(ConsumerProducer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - output_queue=default_output_queue, - subscriber=default_subscriber, - log_level=LogLevel.INFO, - model=default_model, - api_key="", - ): + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + 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") super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - output_queue=output_queue, - subscriber=subscriber, - input_schema=TextCompletionRequest, - output_schema=TextCompletionResponse, + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": TextCompletionRequest, + "output_schema": TextCompletionResponse, + "model": model, + } ) self.model = model diff --git a/trustgraph/llm/vertexai_text/llm.py b/trustgraph/llm/vertexai_text/llm.py index c6c6c81d..30896fb5 100755 --- a/trustgraph/llm/vertexai_text/llm.py +++ b/trustgraph/llm/vertexai_text/llm.py @@ -31,26 +31,23 @@ default_subscriber = 'llm-vertexai-text' class Processor(ConsumerProducer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - output_queue=default_output_queue, - subscriber=default_subscriber, - log_level=LogLevel.INFO, - region="us-west1", - model="gemini-1.0-pro-001", - private_key=None, - ): + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) + region = params.get("region", "us-west1") + model = params.get("model", "gemini-1.0-pro-001") + private_key = params.get("private_key") super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - output_queue=output_queue, - subscriber=subscriber, - input_schema=TextCompletionRequest, - output_schema=TextCompletionResponse, + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": TextCompletionRequest, + "output_schema": TextCompletionResponse, + } ) self.parameters = { diff --git a/trustgraph/rag/graph/rag.py b/trustgraph/rag/graph/rag.py index 3ea322c1..3c1d2b8a 100755 --- a/trustgraph/rag/graph/rag.py +++ b/trustgraph/rag/graph/rag.py @@ -17,32 +17,32 @@ default_vector_store = 'http://localhost:19530' class Processor(ConsumerProducer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - output_queue=default_output_queue, - subscriber=default_subscriber, - log_level=LogLevel.INFO, - graph_hosts=default_graph_hosts, - vector_store=default_vector_store, - entity_limit=50, - triple_limit=30, - max_subgraph_size=3000, - ): + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) + graph_hosts = params.get("graph_hosts", default_graph_hosts) + vector_store = params.get("vector_store", default_vector_store) + entity_limit = params.get("entity_limit", 50) + triple_limit = params.get("triple_limit", 30) + max_subgraph_size = params.get("max_subgraph_size", 3000) super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - output_queue=output_queue, - subscriber=subscriber, - input_schema=GraphRagQuery, - output_schema=GraphRagResponse, + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": GraphRagQuery, + "output_schema": GraphRagResponse, + "entity_limit": entity_limit, + "triple_limit": triple_limit, + "max_subgraph_size": max_subgraph_size, + } ) self.rag = GraphRag( - pulsar_host=pulsar_host, + pulsar_host=self.pulsar_host, graph_hosts=graph_hosts.split(","), vector_store=vector_store, verbose=True, diff --git a/trustgraph/vector/milvus_write/write.py b/trustgraph/vector/milvus_write/write.py index cf2a5f76..8ad56055 100755 --- a/trustgraph/vector/milvus_write/write.py +++ b/trustgraph/vector/milvus_write/write.py @@ -14,21 +14,19 @@ default_store_uri = 'http://localhost:19530' class Processor(Consumer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - subscriber=default_subscriber, - store_uri=default_store_uri, - log_level=LogLevel.INFO, - ): + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + subscriber = params.get("subscriber", default_subscriber) + store_uri = params.get("store_uri", default_store_uri) super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - subscriber=subscriber, - input_schema=VectorsAssociation, + **params | { + "input_queue": input_queue, + "subscriber": subscriber, + "input_schema": VectorsAssociation, + "store_uri": store_uri, + } ) self.vecstore = TripleVectors(store_uri) @@ -40,6 +38,7 @@ class Processor(Consumer): if v.entity.value != "": for vec in v.vectors: self.vecstore.insert(vec, v.entity.value) + @staticmethod def add_args(parser):