apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-05-18 03:45:12 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Lock in subscriber names

Browse source
This commit is contained in:
Cyber MacGeddon 2024-07-23 23:02:09 +01:00
parent de3525371c
commit d75d2923fe
7 changed files with 36 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

5
trustgraph/embeddings/vectorize/vectorize.py
View file

@ -34,7 +34,10 @@ class Processor(ConsumerProducer):
}
)
self.embeddings = EmbeddingsClient(pulsar_host=self.pulsar_host)
self.embeddings = EmbeddingsClient(
pulsar_host=self.pulsar_host,
subscriber=module + "emb",
)
def emit(self, source, chunk, vectors):

10
trustgraph/embeddings_client.py
View file

@ -17,14 +17,14 @@ DEBUG=_pulsar.LoggerLevel.Debug
class EmbeddingsClient:
def __init__(
self, log_level=ERROR, client_id=None,
self, log_level=ERROR, subscriber=None,
pulsar_host="pulsar://pulsar:6650",
):
self.client = None
if client_id == None:
client_id = str(uuid.uuid4())
if subscriber == None:
subscriber = str(uuid.uuid4())
self.client = pulsar.Client(
pulsar_host,
@ -38,7 +38,7 @@ class EmbeddingsClient:
)
self.consumer = self.client.subscribe(
embeddings_response_queue, client_id,
embeddings_response_queue, subscriber,
schema=JsonSchema(EmbeddingsResponse),
)
@ -68,7 +68,7 @@ class EmbeddingsClient:
def __del__(self):
if hasattr(self, "consumer"):
self.consumer.unsubscribe()
# self.consumer.unsubscribe()
self.consumer.close()
if hasattr(self, "producer"):

11
trustgraph/graph_rag.py
View file

@ -19,6 +19,7 @@ class GraphRag:
entity_limit=50,
triple_limit=30,
max_subgraph_size=3000,
module="test",
):
self.verbose=verbose
@ -31,7 +32,10 @@ class GraphRag:
self.graph = TrustGraph(graph_hosts)
self.embeddings = EmbeddingsClient(pulsar_host=pulsar_host)
self.embeddings = EmbeddingsClient(
pulsar_host=pulsar_host,
subscriber=module + "-emb",
)
self.vecstore = TripleVectors(vector_store)
@ -41,7 +45,10 @@ class GraphRag:
self.label_cache = {}
self.llm = LlmClient(pulsar_host=pulsar_host)
self.llm = LlmClient(
pulsar_host=pulsar_host,
subscriber=module + "-llm",
)
if self.verbose:
print("Initialised", flush=True)

10
trustgraph/graph_rag_client.py
View file

@ -18,12 +18,12 @@ DEBUG=_pulsar.LoggerLevel.Debug
class GraphRagClient:
def __init__(
self, log_level=ERROR, client_id=None,
self, log_level=ERROR, subscriber=None,
pulsar_host="pulsar://pulsar:6650",
):
if client_id == None:
client_id = str(uuid.uuid4())
if subscriber == None:
subscriber = str(uuid.uuid4())
self.client = pulsar.Client(
pulsar_host,
@ -37,7 +37,7 @@ class GraphRagClient:
)
self.consumer = self.client.subscribe(
graph_rag_response_queue, client_id,
graph_rag_response_queue, subscriber,
schema=JsonSchema(GraphRagResponse),
)
@ -67,7 +67,7 @@ class GraphRagClient:
def __del__(self):
if hasattr(self, "consumer"):
self.consumer.unsubscribe()
# self.consumer.unsubscribe()
self.consumer.close()
if hasattr(self, "producer"):

5
trustgraph/kg/extract_definitions/extract.py
View file

@ -41,7 +41,10 @@ class Processor(ConsumerProducer):
}
)
self.llm = LlmClient(pulsar_host=self.pulsar_host)
self.llm = LlmClient(
pulsar_host=self.pulsar_host,
subscriber = module + "-llm",
)
def to_uri(self, text):

5
trustgraph/kg/extract_relationships/extract.py
View file

@ -61,7 +61,10 @@ class Processor(ConsumerProducer):
"vector_schema": GraphEmbeddings.__name__,
})
self.llm = LlmClient(pulsar_host=self.pulsar_host)
self.llm = LlmClient(
pulsar_host = self.pulsar_host,
subscriber = module + "-llm",
)
def to_uri(self, text):

10
trustgraph/llm_client.py
View file

@ -19,12 +19,12 @@ DEBUG=_pulsar.LoggerLevel.Debug
class LlmClient:
def __init__(
self, log_level=ERROR, client_id=None,
self, log_level=ERROR, subscriber=None,
pulsar_host="pulsar://pulsar:6650",
):
if client_id == None:
client_id = str(uuid.uuid4())
if subscriber == None:
subscriber = str(uuid.uuid4())
self.client = pulsar.Client(
pulsar_host,
@ -38,7 +38,7 @@ class LlmClient:
)
self.consumer = self.client.subscribe(
text_completion_response_queue, client_id,
text_completion_response_queue, subscriber,
schema=JsonSchema(TextCompletionResponse),
)
@ -69,7 +69,7 @@ class LlmClient:
def __del__(self):
if hasattr(self, "consumer"):
self.consumer.unsubscribe()
# self.consumer.unsubscribe()
self.consumer.close()
if hasattr(self, "producer"):