diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index bc40d156..5fad114f 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/trustgraph-base/trustgraph/base/__init__.py @@ -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 diff --git a/trustgraph-base/trustgraph/base/embedings_client.py b/trustgraph-base/trustgraph/base/embedings_client.py new file mode 100644 index 00000000..4016c000 --- /dev/null +++ b/trustgraph-base/trustgraph/base/embedings_client.py @@ -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, + ) + diff --git a/trustgraph-base/trustgraph/base/prompt_client.py b/trustgraph-base/trustgraph/base/prompt_client.py new file mode 100644 index 00000000..393b3b5f --- /dev/null +++ b/trustgraph-base/trustgraph/base/prompt_client.py @@ -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, + ) + diff --git a/trustgraph-base/trustgraph/base/request_response_spec.py b/trustgraph-base/trustgraph/base/request_response_spec.py index 876a8349..16932baa 100644 --- a/trustgraph-base/trustgraph/base/request_response_spec.py +++ b/trustgraph-base/trustgraph/base/request_response_spec.py @@ -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 + diff --git a/trustgraph-base/trustgraph/base/text_completion_client.py b/trustgraph-base/trustgraph/base/text_completion_client.py new file mode 100644 index 00000000..2926ca4c --- /dev/null +++ b/trustgraph-base/trustgraph/base/text_completion_client.py @@ -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, + ) + diff --git a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py index 30d63dc7..a4ae35dc 100755 --- a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py +++ b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py @@ -49,4 +49,3 @@ def run(): Processor.launch(default_ident, __doc__) - diff --git a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py index cf8e2d8a..3189a79d 100755 --- a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py +++ b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py @@ -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 diff --git a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py index 45aae013..f95dadf9 100755 --- a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py @@ -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) diff --git a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py index 7f0cee15..dab6fb86 100755 --- a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py @@ -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: