Spec clients

This commit is contained in:
Cyber MacGeddon 2025-04-19 13:38:48 +01:00
parent c5cd70dacb
commit 93bb2adb9e
9 changed files with 140 additions and 56 deletions

View file

@ -14,5 +14,7 @@ from . subscriber_spec import SubscriberSpec
from . request_response_spec import RequestResponseSpec
from . llm_service import LlmService, LlmResult
from . embeddings_service import EmbeddingsService
from . embeddings_client import EmbeddingsClientSpec
from . text_completion_client import TextCompletionsClientSpec
from . prompt_client import PromptClientSpec

View file

@ -0,0 +1,26 @@
from . request_response_spec import RequestResponse, RequestResponseSpec
from .. schema import EmbeddingsReqeust, EmbeddingsResponse
class EmbeddingsClient(RequestResponse):
async def embed(self, text, timeout=30):
resp = await self.request(
EmbeddingsRequest(
text = text
),
timeout=timeout
)
return resp.vectors
class EmbeddingsClientSpec(RequestResponseSpec):
def __init__(
self, request_name, response_name,
):
super(EmbeddingsRequestResponseSpec, self).__init__(
request_name = request_name,
request_schema = EmbeddingsRequest,
response_name = response_name,
response_schema = EmbeddingsResponse,
impl = EmbeddingsClient,
)

View file

@ -0,0 +1,52 @@
import json
from . request_response_spec import RequestResponse, RequestResponseSpec
from .. schema import EmbeddingsReqeust, EmbeddingsResponse
class PromptClient(RequestResponse):
async def prompt(self, id, variables, timeout=600):
resp = await self.request(
PromptRequest(
id = id,
terms = {
k: json.dumps(v)
for k, v in variables.items()
}
),
timeout=timeout
)
if resp.error:
raise RuntimeError(resp.error.message)
if resp.text: return resp.text
return json.loads(resp.object)
async def extract_definitions(self, text, timeout=600):
return await self.prompt(
id = "extract-definitions",
variables = { "text": text }
)
async def extract_relationships(self, text, timeout=600):
return await self.prompt(
id = "extract-relationships",
variables = { "text": text }
)
class PromptClientSpec(RequestResponseSpec):
def __init__(
self, request_name, response_name,
):
super(EmbeddingsRequestResponseSpec, self).__init__(
request_name = request_name,
request_schema = EmbeddingsRequest,
response_name = response_name,
response_schema = EmbeddingsResponse,
impl = TextCompletionClient,
)

View file

@ -75,12 +75,14 @@ class RequestResponse(Subscriber):
# - we receive on the response topic as a subscriber
class RequestResponseSpec(Spec):
def __init__(
self, request_name, request_schema, response_name, response_schema
self, request_name, request_schema, response_name,
response_schema, impl=RequestResponse
):
self.request_name = request_name
self.request_schema = request_schema
self.response_name = response_name
self.response_schema = response_schema
self.impl = impl
def add(self, flow, processor, definition):
@ -88,7 +90,7 @@ class RequestResponseSpec(Spec):
flow.id, f"{flow.name}-{self.response_name}"
)
rr = RequestResponse(
rr = self.impl(
client = processor.client,
subscription = flow.id,
consumer_name = flow.id,
@ -102,3 +104,4 @@ class RequestResponseSpec(Spec):
flow.consumer[self.request_name] = rr

View file

@ -0,0 +1,30 @@
from . request_response_spec import RequestResponse, RequestResponseSpec
from .. schema import EmbeddingsReqeust, EmbeddingsResponse
class TextCompletionClient(RequestResponse):
async def text_completion(self, system, prompt, timeout=600):
resp = await self.request(
TextCompletionRequest(
system = system, prompt = prompt
),
timeout=timeout
)
if resp.error:
raise RuntimeError(resp.error.message)
return resp.response
class TextCompletionClientSpec(RequestResponseSpec):
def __init__(
self, request_name, response_name,
):
super(EmbeddingsRequestResponseSpec, self).__init__(
request_name = request_name,
request_schema = EmbeddingsRequest,
response_name = response_name,
response_schema = EmbeddingsResponse,
impl = TextCompletionClient,
)

View file

@ -49,4 +49,3 @@ def run():
Processor.launch(default_ident, __doc__)

View file

@ -8,7 +8,7 @@ Output is entity plus embedding.
from ... schema import EntityContexts, EntityEmbeddings, GraphEmbeddings
from ... schema import EmbeddingsRequest, EmbeddingsResponse
from ... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
from ... base import FlowProcessor, EmbeddingsClientSpec, ConsumerSpec
from ... base import ProducerSpec
default_ident = "graph-embeddings"
@ -34,11 +34,9 @@ class Processor(FlowProcessor):
)
self.register_specification(
RequestResponseSpec(
EmbeddingsClientSpec(
request_name = "embeddings-request",
request_schema = EmbeddingsRequest,
response_name = "embeddings-response",
response_schema = EmbeddingsResponse,
)
)
@ -60,10 +58,8 @@ class Processor(FlowProcessor):
for entity in v.entities:
resp = await flow("embeddings-request").request(
EmbeddingsRequest(
text = entity.context
)
resp = await flow("embeddings-request").embed(
text = entity.context
)
vectors = resp.vectors

View file

@ -13,8 +13,8 @@ from .... schema import EntityContext, EntityContexts
from .... schema import PromptRequest, PromptResponse
from .... rdf import TRUSTGRAPH_ENTITIES, DEFINITION, RDF_LABEL, SUBJECT_OF
from .... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
from .... base import ProducerSpec
from .... base import FlowProcessor, ConsumerSpec, ProducerSpec
from .... base import PromptClientSpec
DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True)
RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
@ -43,11 +43,9 @@ class Processor(FlowProcessor):
)
self.register_specification(
RequestResponseSpec(
PromptClientSpec(
request_name = "prompt-request",
request_schema = PromptRequest,
response_name = "prompt-response",
response_schema = PromptResponse,
)
)
@ -102,24 +100,14 @@ class Processor(FlowProcessor):
try:
resp = await flow("prompt-request").request(
PromptRequest(
id="extract-definitions",
terms={
"text": json.dumps(chunk)
},
)
defs = await flow("prompt-request").extract_definitions(
text = chunk
)
print("Response", resp, flush=True)
if resp.error is not None:
print("Error:", resp.error.message, flush=True)
raise RuntimeError(resp.error.message)
print("Response", defs, flush=True)
if resp.object is None:
raise RuntimeError("Expecting object in prompt response")
defs = json.loads(resp.object)
if type(defs) != list:
raise RuntimeError("Expecting array in prompt response")
except Exception as e:
print("Prompt exception:", e, flush=True)

View file

@ -13,8 +13,8 @@ from .... schema import Metadata, Value
from .... schema import PromptRequest, PromptResponse
from .... rdf import RDF_LABEL, TRUSTGRAPH_ENTITIES, SUBJECT_OF
from .... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
from .... base import ProducerSpec
from .... base import FlowProcessor, ConsumerSpec, ProducerSpec
from .... base import PromptClientSpec
RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True)
@ -42,11 +42,9 @@ class Processor(FlowProcessor):
)
self.register_specification(
RequestResponseSpec(
PromptClientSpec(
request_name = "prompt-request",
request_schema = PromptRequest,
response_name = "prompt-response",
response_schema = PromptResponse,
)
)
@ -86,24 +84,14 @@ class Processor(FlowProcessor):
try:
resp = await flow("prompt-request").request(
PromptRequest(
id="extract-relationships",
terms={
"text": json.dumps(chunk)
},
)
rels = await flow("prompt-request").extract_relationships(
text = chunk
)
print("Response", resp, flush=True)
if resp.error is not None:
print("Error:", resp.error.message, flush=True)
raise RuntimeError(resp.error.message)
print("Response", rels, flush=True)
if resp.object is None:
raise RuntimeError("Expecting object in prompt response")
rels = json.loads(resp.object)
if type(rels) != list:
raise RuntimeError("Expecting array in prompt response")
except Exception as e:
print("Prompt exception:", e, flush=True)
@ -118,9 +106,9 @@ class Processor(FlowProcessor):
for rel in rels:
s = rel["s"]
p = rel["p"]
o = rel["o"]
s = rel["subject"]
p = rel["predicate"]
o = rel["object"]
if s == "": continue
if p == "": continue
@ -136,7 +124,7 @@ class Processor(FlowProcessor):
p_uri = self.to_uri(p)
p_value = Value(value=str(p_uri), is_uri=True)
if rel["o_entity"]:
if rel["object-entity"]:
o_uri = self.to_uri(o)
o_value = Value(value=str(o_uri), is_uri=True)
else: