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

Request/response spec working, implemented in extract-defs and vertexai

Browse source
This commit is contained in:
Cyber MacGeddon 2025-04-18 18:35:17 +01:00
parent 36013212e0
commit 1047cd2fa1
22 changed files with 378 additions and 370 deletions
Download patch file Download diff file Expand all files Collapse all files

22
trustgraph-flow/trustgraph/chunking/recursive/chunker.py
View file

@ -10,7 +10,7 @@ from prometheus_client import Histogram
from ... schema import TextDocument, Chunk, Metadata
from ... schema import text_ingest_queue, chunk_ingest_queue
from ... log_level import LogLevel
from ... base import FlowProcessor
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
default_ident = "chunker"
@ -41,15 +41,19 @@ class Processor(FlowProcessor):
is_separator_regex=False,
)
self.register_consumer(
name = "input",
schema = TextDocument,
handler = self.on_message,
self.register_specification(
ConsumerSpec(
name = "input",
schema = TextDocument,
handler = self.on_message,
)
)
self.register_producer(
name = "output",
schema = Chunk,
self.register_specification(
ProducerSpec(
name = "output",
schema = Chunk,
)
)
print("Chunker initialised", flush=True)
@ -76,7 +80,7 @@ class Processor(FlowProcessor):
id=consumer.id, flow=consumer.flow
).observe(len(chunk.page_content))
await flow.producer["output"].send(r)
await flow("output").send(r)
print("Done.", flush=True)

20
trustgraph-flow/trustgraph/chunking/token/chunker.py
View file

@ -40,15 +40,19 @@ class Processor(FlowProcessor):
chunk_overlap=chunk_overlap,
)
self.register_consumer(
name = "input",
schema = TextDocument,
handler = self.on_message,
self.register_specification(
ConsumerSpec(
name = "input",
schema = TextDocument,
handler = self.on_message,
)
)
self.register_producer(
name = "output",
schema = Chunk,
self.register_specification(
ProducerSpec(
name = "output",
schema = Chunk,
)
)
print("Chunker initialised", flush=True)
@ -75,7 +79,7 @@ class Processor(FlowProcessor):
id=consumer.id, flow=consumer.flow
).observe(len(chunk.page_content))
await flow.producer["output"].send(r)
await flow("output").send(r)
print("Done.", flush=True)

29
trustgraph-flow/trustgraph/config/service/service.py
View file

@ -14,6 +14,7 @@ from trustgraph.base import AsyncProcessor, Consumer, Producer
from . config import Configuration
from ... base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
from ... base import Consumer, Producer
default_ident = "config-svc"
@ -42,29 +43,33 @@ class Processor(AsyncProcessor):
}
)
self.request_metrics = ConsumerMetrics(id + "-request")
self.response_metrics = ProducerMetrics(id + "-response")
self.push_metrics = ProducerMetrics(id + "-push")
request_metrics = ConsumerMetrics(id + "-request")
response_metrics = ProducerMetrics(id + "-response")
push_metrics = ProducerMetrics(id + "-push")
self.push_pub = self.publish(
queue = push_queue,
self.push_pub = Producer(
client = self.client,
topic = push_queue,
schema = ConfigPush,
metrics = self.push_metrics,
metrics = push_metrics,
)
self.response_pub = self.publish(
queue = response_queue,
self.response_pub = Producer(
client = self.client,
topic = response_queue,
schema = ConfigResponse,
metrics = self.response_metrics,
metrics = response_metrics,
)
self.subs = self.subscribe(
self.subs = Consumer(
taskgroup = self.taskgroup,
client = self.client,
flow = None,
queue = request_queue,
topic = request_queue,
subscriber = id,
schema = request_schema,
handler = self.on_message,
metrics = self.request_metrics,
metrics = request_metrics,
)
self.config = Configuration(self.push)

22
trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py
View file

@ -11,7 +11,7 @@ from langchain_community.document_loaders import PyPDFLoader
from ... schema import Document, TextDocument, Metadata
from ... schema import document_ingest_queue, text_ingest_queue
from ... log_level import LogLevel
from ... base import FlowProcessor
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
default_ident = "pdf-decoder"
@ -27,15 +27,19 @@ class Processor(FlowProcessor):
}
)
self.register_consumer(
name = "input",
schema = Document,
handler = self.on_message,
self.register_specification(
ConsumerSpec(
name = "input",
schema = Document,
handler = self.on_message,
)
)
self.register_producer(
name = "output",
schema = TextDocument,
self.register_specification(
ProducerSpec(
name = "output",
schema = TextDocument,
)
)
print("PDF inited", flush=True)
@ -67,7 +71,7 @@ class Processor(FlowProcessor):
text=page.page_content.encode("utf-8"),
)
await flow.producer["output"].send(r)
await flow("output").send(r)
print("Done.", flush=True)

38
trustgraph-flow/trustgraph/extract/kg/definitions/extract.py
View file

@ -6,7 +6,6 @@ entity/context definitions for embedding.
"""
import json
import asyncio
import urllib.parse
from pulsar.schema import JsonSchema
import uuid
@ -18,7 +17,7 @@ from .... log_level import LogLevel
from .... clients.prompt_client import PromptClient
from .... rdf import TRUSTGRAPH_ENTITIES, DEFINITION, RDF_LABEL, SUBJECT_OF
from .... base import FlowProcessor, SubscriberSpec, ConsumerSpec
from .... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
from .... base import ProducerSpec
DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True)
@ -48,16 +47,11 @@ class Processor(FlowProcessor):
)
self.register_specification(
ProducerSpec(
name = "prompt-request",
schema = PromptRequest
)
)
self.register_specification(
SubscriberSpec(
name = "prompt-response",
schema = PromptResponse,
RequestResponseSpec(
request_name = "prompt-request",
request_schema = PromptRequest,
response_name = "prompt-response",
response_schema = PromptResponse,
)
)
@ -116,26 +110,16 @@ class Processor(FlowProcessor):
try:
q = await flow.consumer["prompt-response"].subscribe(id)
try:
await flow.producer["prompt-request"].send(
resp = await flow("prompt-request").request(
PromptRequest(
id="extract-definitions",
terms={
"text": json.dumps(chunk)
},
),
properties={"id": id}
)
)
# FIXME: hard-coded?
resp = await asyncio.wait_for(
q.get(),
timeout=600
)
print("Response", resp, flush=True)
if resp.error is not None:
@ -150,8 +134,6 @@ class Processor(FlowProcessor):
except Exception as e:
print("Prompt exception:", e, flush=True)
raise e
finally:
await flow.consumer["prompt-response"].unsubscribe(id)
triples = []
entities = []
@ -201,7 +183,7 @@ class Processor(FlowProcessor):
entities.append(ec)
await self.emit_triples(
flow.producer["triples"],
flow("triples"),
Metadata(
id=v.metadata.id,
metadata=[],
@ -212,7 +194,7 @@ class Processor(FlowProcessor):
)
await self.emit_ecs(
flow.producer["entity-contexts"],
flow("entity-contexts"),
Metadata(
id=v.metadata.id,
metadata=[],

36
trustgraph-flow/trustgraph/model/prompt/template/service.py
View file

@ -13,7 +13,7 @@ from .... schema import PromptRequest, PromptResponse, Error
from .... schema import TextCompletionRequest, TextCompletionResponse
from .... base import FlowProcessor
from .... base import ProducerSpec, SubscriberSpec, ConsumerSpec
from .... base import ProducerSpec, ConsumerSpec, RequestResponseSpec
from . prompt_manager import PromptConfiguration, Prompt, PromptManager
@ -43,16 +43,11 @@ class Processor(FlowProcessor):
)
self.register_specification(
ProducerSpec(
name = "text-completion-request",
schema = TextCompletionRequest
)
)
self.register_specification(
SubscriberSpec(
name = "text-completion-response",
schema = TextCompletionResponse,
RequestResponseSpec(
request_name = "text-completion-request",
request_schema = TextCompletionRequest,
response_name = "text-completion-response",
response_schema = TextCompletionResponse,
)
)
@ -139,24 +134,15 @@ class Processor(FlowProcessor):
print(f"Handling kind {kind}...", flush=True)
q = await flow.consumer["text-completion-response"].subscribe(id)
async def llm(system, prompt):
print(system, flush=True)
print(prompt, flush=True)
await flow.producer["text-completion-request"].send(
resp = await flow("text-completion-request").request(
TextCompletionRequest(
system=system, prompt=prompt
),
properties={"id": id}
)
# FIXME: hard-coded?
resp = await asyncio.wait_for(
q.get(),
timeout=600
)
try:
@ -170,8 +156,6 @@ class Processor(FlowProcessor):
except Exception as e:
print("Invocation exception:", e, flush=True)
raise e
finally:
await flow.consumer["text-completion-response"].unsubscribe(id)
print(resp, flush=True)
@ -185,7 +169,7 @@ class Processor(FlowProcessor):
error=None,
)
await flow.response.send(r, properties={"id": id})
await flow("response").send(r, properties={"id": id})
return
@ -200,7 +184,7 @@ class Processor(FlowProcessor):
error=None,
)
await flow.response.send(r, properties={"id": id})
await flow("response").send(r, properties={"id": id})
return
@ -218,7 +202,7 @@ class Processor(FlowProcessor):
response=None,
)
await flow.response.send(r, properties={"id": id})
await flow("response").send(r, properties={"id": id})
except Exception as e: