From 9fa67b5db757e8b3cd42bd721720f5644276a994 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 2 Dec 2024 17:22:40 +0000 Subject: [PATCH] Load --- test-api/test-agent2-api | 28 + test-api/test-dbpedia | 30 + test-api/test-encyclopedia | 30 + test-api/test-internet-search | 30 + test-api/test-triples-query-api | 5 +- .../trustgraph/api/gateway/agent.py | 30 + .../trustgraph/api/gateway/dbpedia.py | 29 + .../trustgraph/api/gateway/embeddings.py | 27 + .../trustgraph/api/gateway/encyclopedia.py | 29 + .../api/gateway/graph_embeddings_load.py | 60 ++ .../api/gateway/graph_embeddings_stream.py | 56 ++ .../trustgraph/api/gateway/internet_search.py | 29 + .../trustgraph/api/gateway/serialize.py | 2 +- .../trustgraph/api/gateway/service.py | 526 +++--------------- .../trustgraph/api/gateway/socket.py | 68 +++ .../trustgraph/api/gateway/triples_load.py | 59 ++ .../trustgraph/api/gateway/triples_stream.py | 56 ++ 17 files changed, 638 insertions(+), 456 deletions(-) create mode 100755 test-api/test-agent2-api create mode 100755 test-api/test-dbpedia create mode 100755 test-api/test-encyclopedia create mode 100755 test-api/test-internet-search create mode 100644 trustgraph-flow/trustgraph/api/gateway/agent.py create mode 100644 trustgraph-flow/trustgraph/api/gateway/dbpedia.py create mode 100644 trustgraph-flow/trustgraph/api/gateway/embeddings.py create mode 100644 trustgraph-flow/trustgraph/api/gateway/encyclopedia.py create mode 100644 trustgraph-flow/trustgraph/api/gateway/graph_embeddings_load.py create mode 100644 trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py create mode 100644 trustgraph-flow/trustgraph/api/gateway/internet_search.py create mode 100644 trustgraph-flow/trustgraph/api/gateway/socket.py create mode 100644 trustgraph-flow/trustgraph/api/gateway/triples_load.py create mode 100644 trustgraph-flow/trustgraph/api/gateway/triples_stream.py diff --git a/test-api/test-agent2-api b/test-api/test-agent2-api new file mode 100755 index 00000000..766b16c9 --- /dev/null +++ b/test-api/test-agent2-api @@ -0,0 +1,28 @@ +#!/usr/bin/env python3 + +import requests +import json +import sys + +url = "http://localhost:8088/api/v1/" + +############################################################################ + +input = { + "question": "What is 14 plus 12. Justify your answer.", +} + +resp = requests.post( + f"{url}agent", + json=input, +) + +resp = resp.json() + +if "error" in resp: + print(f"Error: {resp['error']}") + sys.exit(1) + +print(resp["answer"]) + + diff --git a/test-api/test-dbpedia b/test-api/test-dbpedia new file mode 100755 index 00000000..e361f533 --- /dev/null +++ b/test-api/test-dbpedia @@ -0,0 +1,30 @@ +#!/usr/bin/env python3 + +import requests +import json +import sys + +url = "http://localhost:8088/api/v1/" + +############################################################################ + +input = { + "term": "Cornwall", +} + +resp = requests.post( + f"{url}dbpedia", + json=input, +) + +resp = resp.json() + +if "error" in resp: + print(f"Error: {resp['error']}") + sys.exit(1) + +print(resp["text"]) + +sys.exit(0) +############################################################################ + diff --git a/test-api/test-encyclopedia b/test-api/test-encyclopedia new file mode 100755 index 00000000..ad4e5b36 --- /dev/null +++ b/test-api/test-encyclopedia @@ -0,0 +1,30 @@ +#!/usr/bin/env python3 + +import requests +import json +import sys + +url = "http://localhost:8088/api/v1/" + +############################################################################ + +input = { + "term": "Cornwall", +} + +resp = requests.post( + f"{url}encyclopedia", + json=input, +) + +resp = resp.json() + +if "error" in resp: + print(f"Error: {resp['error']}") + sys.exit(1) + +print(resp["text"]) + +sys.exit(0) +############################################################################ + diff --git a/test-api/test-internet-search b/test-api/test-internet-search new file mode 100755 index 00000000..8c854c77 --- /dev/null +++ b/test-api/test-internet-search @@ -0,0 +1,30 @@ +#!/usr/bin/env python3 + +import requests +import json +import sys + +url = "http://localhost:8088/api/v1/" + +############################################################################ + +input = { + "term": "Cornwall", +} + +resp = requests.post( + f"{url}internet-search", + json=input, +) + +resp = resp.json() + +if "error" in resp: + print(f"Error: {resp['error']}") + sys.exit(1) + +print(resp["text"]) + +sys.exit(0) +############################################################################ + diff --git a/test-api/test-triples-query-api b/test-api/test-triples-query-api index e2895a28..1aa8a0b1 100755 --- a/test-api/test-triples-query-api +++ b/test-api/test-triples-query-api @@ -9,7 +9,10 @@ url = "http://localhost:8088/api/v1/" ############################################################################ input = { - "p": "http://www.w3.org/2000/01/rdf-schema#label", + "p": { + "v": "http://www.w3.org/2000/01/rdf-schema#label", + "e": True, + }, "limit": 10 } diff --git a/trustgraph-flow/trustgraph/api/gateway/agent.py b/trustgraph-flow/trustgraph/api/gateway/agent.py new file mode 100644 index 00000000..28a1e185 --- /dev/null +++ b/trustgraph-flow/trustgraph/api/gateway/agent.py @@ -0,0 +1,30 @@ + +from ... schema import AgentRequest, AgentResponse +from ... schema import agent_request_queue +from ... schema import agent_response_queue + +from . endpoint import MultiResponseServiceEndpoint + +class AgentEndpoint(MultiResponseServiceEndpoint): + def __init__(self, pulsar_host, timeout): + + super(AgentEndpoint, self).__init__( + pulsar_host=pulsar_host, + request_queue=agent_request_queue, + response_queue=agent_response_queue, + request_schema=AgentRequest, + response_schema=AgentResponse, + endpoint_path="/api/v1/agent", + timeout=timeout, + ) + + def to_request(self, body): + return AgentRequest( + question=body["question"] + ) + + def from_response(self, message): + if message.answer: + return { "answer": message.answer }, True + else: + return {}, False diff --git a/trustgraph-flow/trustgraph/api/gateway/dbpedia.py b/trustgraph-flow/trustgraph/api/gateway/dbpedia.py new file mode 100644 index 00000000..0ccb3d6b --- /dev/null +++ b/trustgraph-flow/trustgraph/api/gateway/dbpedia.py @@ -0,0 +1,29 @@ + +from ... schema import LookupRequest, LookupResponse +from ... schema import dbpedia_lookup_request_queue +from ... schema import dbpedia_lookup_response_queue + +from . endpoint import ServiceEndpoint + +class DbpediaEndpoint(ServiceEndpoint): + def __init__(self, pulsar_host, timeout): + + super(DbpediaEndpoint, self).__init__( + pulsar_host=pulsar_host, + request_queue=dbpedia_lookup_request_queue, + response_queue=dbpedia_lookup_response_queue, + request_schema=LookupRequest, + response_schema=LookupResponse, + endpoint_path="/api/v1/dbpedia", + timeout=timeout, + ) + + def to_request(self, body): + return LookupRequest( + term=body["term"], + kind=body.get("kind", None), + ) + + def from_response(self, message): + return { "text": message.text } + diff --git a/trustgraph-flow/trustgraph/api/gateway/embeddings.py b/trustgraph-flow/trustgraph/api/gateway/embeddings.py new file mode 100644 index 00000000..b5fcc0a4 --- /dev/null +++ b/trustgraph-flow/trustgraph/api/gateway/embeddings.py @@ -0,0 +1,27 @@ + +from ... schema import EmbeddingsRequest, EmbeddingsResponse +from ... schema import embeddings_request_queue +from ... schema import embeddings_response_queue + +from . endpoint import ServiceEndpoint + +class EmbeddingsEndpoint(ServiceEndpoint): + def __init__(self, pulsar_host, timeout): + + super(EmbeddingsEndpoint, self).__init__( + pulsar_host=pulsar_host, + request_queue=embeddings_request_queue, + response_queue=embeddings_response_queue, + request_schema=EmbeddingsRequest, + response_schema=EmbeddingsResponse, + endpoint_path="/api/v1/embeddings", + timeout=timeout, + ) + + def to_request(self, body): + return EmbeddingsRequest( + text=body["text"] + ) + + def from_response(self, message): + return { "vectors": message.vectors } diff --git a/trustgraph-flow/trustgraph/api/gateway/encyclopedia.py b/trustgraph-flow/trustgraph/api/gateway/encyclopedia.py new file mode 100644 index 00000000..e379d7d4 --- /dev/null +++ b/trustgraph-flow/trustgraph/api/gateway/encyclopedia.py @@ -0,0 +1,29 @@ + +from ... schema import LookupRequest, LookupResponse +from ... schema import encyclopedia_lookup_request_queue +from ... schema import encyclopedia_lookup_response_queue + +from . endpoint import ServiceEndpoint + +class EncyclopediaEndpoint(ServiceEndpoint): + def __init__(self, pulsar_host, timeout): + + super(EncyclopediaEndpoint, self).__init__( + pulsar_host=pulsar_host, + request_queue=encyclopedia_lookup_request_queue, + response_queue=encyclopedia_lookup_response_queue, + request_schema=LookupRequest, + response_schema=LookupResponse, + endpoint_path="/api/v1/encyclopedia", + timeout=timeout, + ) + + def to_request(self, body): + return LookupRequest( + term=body["term"], + kind=body.get("kind", None), + ) + + def from_response(self, message): + return { "text": message.text } + diff --git a/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_load.py b/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_load.py new file mode 100644 index 00000000..3cc3f533 --- /dev/null +++ b/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_load.py @@ -0,0 +1,60 @@ + +import asyncio +from pulsar.schema import JsonSchema +import uuid +from aiohttp import WSMsgType + +from ... schema import Metadata +from ... schema import GraphEmbeddings +from ... schema import graph_embeddings_store_queue + +from . publisher import Publisher +from . socket import SocketEndpoint +from . serialize import to_subgraph, to_value + +class GraphEmbeddingsLoadEndpoint(SocketEndpoint): + + def __init__(self, pulsar_host, path="/api/v1/load/graph-embeddings"): + + super(GraphEmbeddingsLoadEndpoint, self).__init__( + endpoint_path=path + ) + + self.pulsar_host=pulsar_host + + self.publisher = Publisher( + self.pulsar_host, graph_embeddings_store_queue, + schema=JsonSchema(GraphEmbeddings) + ) + + async def start(self): + + self.task = asyncio.create_task( + self.publisher.run() + ) + + async def listener(self, ws, running): + + async for msg in ws: + # On error, finish + if msg.type == WSMsgType.ERROR: + break + else: + + data = msg.json() + + elt = GraphEmbeddings( + metadata=Metadata( + id=data["metadata"]["id"], + metadata=to_subgraph(data["metadata"]["metadata"]), + user=data["metadata"]["user"], + collection=data["metadata"]["collection"], + ), + entity=to_value(data["entity"]), + vectors=data["vectors"], + ) + + await self.publisher.send(None, elt) + + + running.stop() diff --git a/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py b/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py new file mode 100644 index 00000000..978684cf --- /dev/null +++ b/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py @@ -0,0 +1,56 @@ + +import asyncio +from pulsar.schema import JsonSchema +import uuid + +from ... schema import GraphEmbeddings +from ... schema import graph_embeddings_store_queue + +from . subscriber import Subscriber +from . socket import SocketEndpoint +from . serialize import serialize_graph_embeddings + +class GraphEmbeddingsStreamEndpoint(SocketEndpoint): + + def __init__(self, pulsar_host, path="/api/v1/stream/graph-embeddings"): + + super(GraphEmbeddingsStreamEndpoint, self).__init__( + endpoint_path=path + ) + + self.pulsar_host=pulsar_host + + self.subscriber = Subscriber( + self.pulsar_host, graph_embeddings_store_queue, + "api-gateway", "api-gateway", + schema=JsonSchema(GraphEmbeddings) + ) + + async def start(self): + + self.task = asyncio.create_task( + self.subscriber.run() + ) + + async def async_thread(self, ws, running): + + id = str(uuid.uuid4()) + + q = await self.subscriber.subscribe_all(id) + + while running.get(): + try: + resp = await asyncio.wait_for(q.get(), 0.5) + await ws.send_json(serialize_graph_embeddings(resp)) + + except TimeoutError: + continue + + except Exception as e: + print(f"Exception: {str(e)}", flush=True) + break + + await self.subscriber.unsubscribe_all(id) + + running.stop() + diff --git a/trustgraph-flow/trustgraph/api/gateway/internet_search.py b/trustgraph-flow/trustgraph/api/gateway/internet_search.py new file mode 100644 index 00000000..c84ed82a --- /dev/null +++ b/trustgraph-flow/trustgraph/api/gateway/internet_search.py @@ -0,0 +1,29 @@ + +from ... schema import LookupRequest, LookupResponse +from ... schema import internet_search_request_queue +from ... schema import internet_search_response_queue + +from . endpoint import ServiceEndpoint + +class InternetSearchEndpoint(ServiceEndpoint): + def __init__(self, pulsar_host, timeout): + + super(InternetSearchEndpoint, self).__init__( + pulsar_host=pulsar_host, + request_queue=internet_search_request_queue, + response_queue=internet_search_response_queue, + request_schema=LookupRequest, + response_schema=LookupResponse, + endpoint_path="/api/v1/internet-search", + timeout=timeout, + ) + + def to_request(self, body): + return LookupRequest( + term=body["term"], + kind=body.get("kind", None), + ) + + def from_response(self, message): + return { "text": message.text } + diff --git a/trustgraph-flow/trustgraph/api/gateway/serialize.py b/trustgraph-flow/trustgraph/api/gateway/serialize.py index 6bc77422..2b955645 100644 --- a/trustgraph-flow/trustgraph/api/gateway/serialize.py +++ b/trustgraph-flow/trustgraph/api/gateway/serialize.py @@ -52,6 +52,6 @@ def serialize_graph_embeddings(message): "collection": message.metadata.collection, }, "vectors": message.vectors, - "entity": message.entity, + "entity": serialize_value(message.entity), } diff --git a/trustgraph-flow/trustgraph/api/gateway/service.py b/trustgraph-flow/trustgraph/api/gateway/service.py index 135d7934..d2798a02 100755 --- a/trustgraph-flow/trustgraph/api/gateway/service.py +++ b/trustgraph-flow/trustgraph/api/gateway/service.py @@ -28,9 +28,6 @@ from prometheus_client import start_http_server from ... log_level import LogLevel -from trustgraph.clients.llm_client import LlmClient -from trustgraph.clients.prompt_client import PromptClient - from ... schema import Value, Metadata, Document, TextDocument, Triple from ... schema import Triples @@ -39,21 +36,7 @@ from ... schema import triples_store_queue from ... schema import GraphEmbeddings from ... schema import graph_embeddings_store_queue -from ... schema import AgentRequest, AgentResponse -from ... schema import agent_request_queue -from ... schema import agent_response_queue - -from ... schema import EmbeddingsRequest, EmbeddingsResponse -from ... schema import embeddings_request_queue -from ... schema import embeddings_response_queue - from ... schema import LookupRequest, LookupResponse -from ... schema import encyclopedia_lookup_request_queue -from ... schema import encyclopedia_lookup_response_queue -from ... schema import internet_search_request_queue -from ... schema import internet_search_response_queue -from ... schema import dbpedia_lookup_request_queue -from ... schema import dbpedia_lookup_response_queue from ... schema import document_ingest_queue, text_ingest_queue @@ -70,6 +53,15 @@ from . text_completion import TextCompletionEndpoint from . prompt import PromptEndpoint from . graph_rag import GraphRagEndpoint from . triples_query import TriplesQueryEndpoint +from . embeddings import EmbeddingsEndpoint +from . encyclopedia import EncyclopediaEndpoint +from . agent import AgentEndpoint +from . dbpedia import DbpediaEndpoint +from . internet_search import InternetSearchEndpoint +from . triples_stream import TriplesStreamEndpoint +from . graph_embeddings_stream import GraphEmbeddingsStreamEndpoint +from . triples_load import TriplesLoadEndpoint +from . graph_embeddings_load import GraphEmbeddingsLoadEndpoint logger = logging.getLogger("api") logger.setLevel(logging.INFO) @@ -79,73 +71,6 @@ default_timeout = 600 default_port = 8088 -class EmbeddingsEndpoint(ServiceEndpoint): - def __init__(self, pulsar_host, timeout): - - super(EmbeddingsEndpoint, self).__init__( - pulsar_host=pulsar_host, - request_queue=embeddings_request_queue, - response_queue=embeddings_response_queue, - request_schema=EmbeddingsRequest, - response_schema=EmbeddingsResponse, - endpoint_path="/api/v1/embeddings", - timeout=timeout, - ) - - def to_request(self, body): - return EmbeddingsRequest( - text=body["text"] - ) - - def from_response(self, message): - return { "vectors": message.vectors } - -class AgentEndpoint(MultiResponseServiceEndpoint): - def __init__(self, pulsar_host, timeout): - - super(AgentEndpoint, self).__init__( - pulsar_host=pulsar_host, - request_queue=agent_request_queue, - response_queue=agent_response_queue, - request_schema=AgentRequest, - response_schema=AgentResponse, - endpoint_path="/api/v1/agent", - timeout=timeout, - ) - - def to_request(self, body): - return AgentRequest( - question=body["question"] - ) - - def from_response(self, message): - if message.answer: - return { "answer": message.answer }, True - else: - return {}, False - -class EncyclopediaEndpoint(ServiceEndpoint): - def __init__(self, pulsar_host, timeout): - - super(EncyclopediaEndpoint, self).__init__( - pulsar_host=pulsar_host, - request_queue=encyclopedia_lookup_request_queue, - response_queue=encyclopedia_lookup_response_queue, - request_schema=LookupRequest, - response_schema=LookupResponse, - endpoint_path="/api/v1/encyclopedia", - timeout=timeout, - ) - - def to_request(self, body): - return LookupRequest( - term=body["term"], - kind=body.get("kind", None), - ) - - def from_response(self, message): - return { "text": message.text } - class Api: def __init__(self, **config): @@ -159,54 +84,59 @@ class Api: self.timeout = int(config.get("timeout", default_timeout)) self.pulsar_host = config.get("pulsar_host", default_pulsar_host) - self.text_completion = TextCompletionEndpoint( - pulsar_host=self.pulsar_host, timeout=self.timeout, - ) + self.endpoints = [ + TextCompletionEndpoint( + pulsar_host=self.pulsar_host, timeout=self.timeout, + ), + PromptEndpoint( + pulsar_host=self.pulsar_host, timeout=self.timeout, + ), + GraphRagEndpoint( + pulsar_host=self.pulsar_host, timeout=self.timeout, + ), + TriplesQueryEndpoint( + pulsar_host=self.pulsar_host, timeout=self.timeout, + ), + EmbeddingsEndpoint( + pulsar_host=self.pulsar_host, timeout=self.timeout, + ), + AgentEndpoint( + pulsar_host=self.pulsar_host, timeout=self.timeout, + ), + EncyclopediaEndpoint( + pulsar_host=self.pulsar_host, timeout=self.timeout, + ), + DbpediaEndpoint( + pulsar_host=self.pulsar_host, timeout=self.timeout, + ), + InternetSearchEndpoint( + pulsar_host=self.pulsar_host, timeout=self.timeout, + ), + TriplesStreamEndpoint( + pulsar_host=self.pulsar_host + ), + GraphEmbeddingsStreamEndpoint( + pulsar_host=self.pulsar_host + ), + TriplesLoadEndpoint( + pulsar_host=self.pulsar_host + ), + GraphEmbeddingsLoadEndpoint( + pulsar_host=self.pulsar_host + ), + ] - self.prompt = PromptEndpoint( - pulsar_host=self.pulsar_host, timeout=self.timeout, - ) - - self.graph_rag = GraphRagEndpoint( - pulsar_host=self.pulsar_host, timeout=self.timeout, - ) - - self.triples_query = TriplesQueryEndpoint( - pulsar_host=self.pulsar_host, timeout=self.timeout, - ) - - self.embeddings = EmbeddingsEndpoint( - pulsar_host=self.pulsar_host, timeout=self.timeout, - ) - - self.agent = AgentEndpoint( - pulsar_host=self.pulsar_host, timeout=self.timeout, - ) - - self.encyclopedia = EncyclopediaEndpoint( - pulsar_host=self.pulsar_host, timeout=self.timeout, - ) - - - - - self.triples_tap = Subscriber( - self.pulsar_host, triples_store_queue, - "api-gateway", "api-gateway", - schema=JsonSchema(Triples) - ) + # self.triples_tap = Subscriber( + # self.pulsar_host, triples_store_queue, + # "api-gateway", "api-gateway", + # schema=JsonSchema(Triples) + # ) self.triples_pub = Publisher( self.pulsar_host, triples_store_queue, schema=JsonSchema(Triples) ) - self.graph_embeddings_tap = Subscriber( - self.pulsar_host, graph_embeddings_store_queue, - "api-gateway", "api-gateway", - schema=JsonSchema(GraphEmbeddings) - ) - self.graph_embeddings_pub = Publisher( self.pulsar_host, graph_embeddings_store_queue, schema=JsonSchema(GraphEmbeddings) @@ -224,185 +154,27 @@ class Api: chunking_enabled=True, ) - - self.internet_search_out = Publisher( - self.pulsar_host, internet_search_request_queue, - schema=JsonSchema(LookupRequest) - ) - - self.internet_search_in = Subscriber( - self.pulsar_host, internet_search_response_queue, - "api-gateway", "api-gateway", - JsonSchema(LookupResponse) - ) - - self.dbpedia_lookup_out = Publisher( - self.pulsar_host, dbpedia_lookup_request_queue, - schema=JsonSchema(LookupRequest) - ) - - self.dbpedia_lookup_in = Subscriber( - self.pulsar_host, dbpedia_lookup_response_queue, - "api-gateway", "api-gateway", - JsonSchema(LookupResponse) - ) - - self.text_completion.add_routes(self.app) - self.prompt.add_routes(self.app) - self.graph_rag.add_routes(self.app) - self.triples_query.add_routes(self.app) - self.embeddings.add_routes(self.app) - self.agent.add_routes(self.app) - self.encyclopedia.add_routes(self.app) + for ep in self.endpoints: + ep.add_routes(self.app) self.app.add_routes([ -# web.post("/api/v1/triples-query", self.triples_query), -# web.post("/api/v1/internet-search", self.internet-search), -# web.post("/api/v1/dbpedia", self.dbpedia), + web.post("/api/v1/load/document", self.load_document), web.post("/api/v1/load/text", self.load_text), - web.get("/api/v1/ws", self.socket), - web.get("/api/v1/stream/triples", self.stream_triples), - web.get( - "/api/v1/stream/graph-embeddings", - self.stream_graph_embeddings - ), +# web.get("/api/v1/ws", self.socket), - web.get("/api/v1/load/triples", self.load_triples), - web.get( - "/api/v1/load/graph-embeddings", - self.load_graph_embeddings - ), +# web.get("/api/v1/stream/triples", self.stream_triples), + +# web.get("/api/v1/load/triples", self.load_triples), + +# web.get( +# "/api/v1/load/graph-embeddings", +# self.load_graph_embeddings +# ), ]) - async def encyclopedia(self, request): - - id = str(uuid.uuid4()) - - try: - - data = await request.json() - - q = await self.encyclopedia_lookup_in.subscribe(id) - - await self.encyclopedia_lookup_out.send( - id, - LookupRequest( - term=data["term"], - kind=data.get("kind", None), - ) - ) - - try: - resp = await asyncio.wait_for(q.get(), self.timeout) - except: - raise RuntimeError("Timeout waiting for response") - - if resp.error: - return web.json_response( - { "error": resp.error.message } - ) - - return web.json_response( - { "text": resp.text } - ) - - except Exception as e: - logging.error(f"Exception: {e}") - - return web.json_response( - { "error": str(e) } - ) - - finally: - await self.encyclopedia_lookup_in.unsubscribe(id) - - async def internet_search(self, request): - - id = str(uuid.uuid4()) - - try: - - data = await request.json() - - q = await self.internet_search_in.subscribe(id) - - await self.internet_search_out.send( - id, - LookupRequest( - term=data["term"], - kind=data.get("kind", None), - ) - ) - - try: - resp = await asyncio.wait_for(q.get(), self.timeout) - except: - raise RuntimeError("Timeout waiting for response") - - if resp.error: - return web.json_response( - { "error": resp.error.message } - ) - - return web.json_response( - { "text": resp.text } - ) - - except Exception as e: - logging.error(f"Exception: {e}") - - return web.json_response( - { "error": str(e) } - ) - - finally: - await self.internet_search_in.unsubscribe(id) - - async def dbpedia(self, request): - - id = str(uuid.uuid4()) - - try: - - data = await request.json() - - q = await self.dbpedia_lookup_in.subscribe(id) - - await self.dbpedia_lookup_out.send( - id, - LookupRequest( - term=data["term"], - kind=data.get("kind", None), - ) - ) - - try: - resp = await asyncio.wait_for(q.get(), self.timeout) - except: - raise RuntimeError("Timeout waiting for response") - - if resp.error: - return web.json_response( - { "error": resp.error.message } - ) - - return web.json_response( - { "text": resp.text } - ) - - except Exception as e: - logging.error(f"Exception: {e}") - - return web.json_response( - { "error": str(e) } - ) - - finally: - await self.dbpedia_lookup_in.unsubscribe(id) - async def load_document(self, request): try: @@ -521,159 +293,19 @@ class Api: except Exception as e: print(f"Exception: {str(e)}", flush=True) - async def stream_triples(self, request): - - id = str(uuid.uuid4()) - - q = await self.triples_tap.subscribe_all(id) - running = Running() - - ws = web.WebSocketResponse() - await ws.prepare(request) - - tsk = asyncio.create_task(self.stream( - q, - ws, - running, - serialize_triples, - )) - - async for msg in ws: - if msg.type == WSMsgType.ERROR: - break - else: - # Ignore incoming messages - pass - - running.stop() - - await self.triples_tap.unsubscribe_all(id) - await tsk - - return ws - - async def stream_graph_embeddings(self, request): - - id = str(uuid.uuid4()) - - q = await self.graph_embeddings_tap.subscribe_all(id) - running = Running() - - ws = web.WebSocketResponse() - await ws.prepare(request) - - tsk = asyncio.create_task(self.stream( - q, - ws, - running, - serialize_graph_embeddings, - )) - - async for msg in ws: - if msg.type == WSMsgType.ERROR: - break - else: - # Ignore incoming messages - pass - - running.stop() - - await self.graph_embeddings_tap.unsubscribe_all(id) - await tsk - - return ws - - async def load_triples(self, request): - - ws = web.WebSocketResponse() - await ws.prepare(request) - - async for msg in ws: - - try: - - if msg.type == WSMsgType.TEXT: - - data = msg.json() - - elt = Triples( - metadata=Metadata( - id=data["metadata"]["id"], - metadata=to_subgraph(data["metadata"]["metadata"]), - user=data["metadata"]["user"], - collection=data["metadata"]["collection"], - ), - triples=to_subgraph(data["triples"]), - ) - - await self.triples_pub.send(None, elt) - - elif msg.type == WSMsgType.ERROR: - break - - except Exception as e: - - print("Exception:", e) - - return ws - - async def load_graph_embeddings(self, request): - - ws = web.WebSocketResponse() - await ws.prepare(request) - - async for msg in ws: - - try: - - if msg.type == WSMsgType.TEXT: - - data = msg.json() - - elt = GraphEmbeddings( - metadata=Metadata( - id=data["metadata"]["id"], - metadata=to_subgraph(data["metadata"]["metadata"]), - user=data["metadata"]["user"], - collection=data["metadata"]["collection"], - ), - entity=to_value(data["entity"]), - vectors=data["vectors"], - ) - - await self.graph_embeddings_pub.send(None, elt) - - elif msg.type == WSMsgType.ERROR: - break - - except Exception as e: - - print("Exception:", e) - - return ws - async def app_factory(self): - await self.text_completion.start() - await self.prompt.start() - await self.graph_rag.start() - await self.triples_query.start() - await self.embeddings.start() - await self.agent.start() - await self.encyclopedia.start() + for ep in self.endpoints: + await ep.start() - self.triples_tap_task = asyncio.create_task( - self.triples_tap.run() - ) +# self.triples_tap_task = asyncio.create_task( +# self.triples_tap.run() +# ) self.triples_pub_task = asyncio.create_task( self.triples_pub.run() ) - self.graph_embeddings_tap_task = asyncio.create_task( - self.graph_embeddings_tap.run() - ) - self.graph_embeddings_pub_task = asyncio.create_task( self.graph_embeddings_pub.run() ) @@ -682,20 +314,6 @@ class Api: self.text_ingest_pub_task = asyncio.create_task(self.text_out.run()) - self.search_pub_task = asyncio.create_task( - self.internet_search_out.run() - ) - self.search_sub_task = asyncio.create_task( - self.internet_search_in.run() - ) - - self.dbpedia_pub_task = asyncio.create_task( - self.dbpedia_lookup_out.run() - ) - self.dbpedia_sub_task = asyncio.create_task( - self.dbpedia_lookup_in.run() - ) - return self.app def run(self): diff --git a/trustgraph-flow/trustgraph/api/gateway/socket.py b/trustgraph-flow/trustgraph/api/gateway/socket.py new file mode 100644 index 00000000..235bfd21 --- /dev/null +++ b/trustgraph-flow/trustgraph/api/gateway/socket.py @@ -0,0 +1,68 @@ + +import asyncio +from aiohttp import web, WSMsgType +import logging + +from . running import Running + +logger = logging.getLogger("socket") +logger.setLevel(logging.INFO) + +class SocketEndpoint: + + def __init__( + self, + endpoint_path="/api/v1/socket", + ): + + self.path = endpoint_path + + async def listener(self, ws, running): + + async for msg in ws: + # On error, finish + if msg.type == WSMsgType.ERROR: + break + else: + # Ignore incoming messages + pass + + running.stop() + + async def async_thread(self, ws, running): + + while running.get(): + try: + await asyncio.sleep(1) + + except TimeoutError: + continue + + except Exception as e: + print(f"Exception: {str(e)}", flush=True) + + async def handle(self, request): + + running = Running() + ws = web.WebSocketResponse() + await ws.prepare(request) + + task = asyncio.create_task(self.async_thread(ws, running)) + + await self.listener(ws, running) + + await task + + running.stop() + + return ws + + async def start(self): + pass + + def add_routes(self, app): + + app.add_routes([ + web.get(self.path, self.handle), + ]) + diff --git a/trustgraph-flow/trustgraph/api/gateway/triples_load.py b/trustgraph-flow/trustgraph/api/gateway/triples_load.py new file mode 100644 index 00000000..d835a363 --- /dev/null +++ b/trustgraph-flow/trustgraph/api/gateway/triples_load.py @@ -0,0 +1,59 @@ + +import asyncio +from pulsar.schema import JsonSchema +import uuid +from aiohttp import WSMsgType + +from ... schema import Metadata +from ... schema import Triples +from ... schema import triples_store_queue + +from . publisher import Publisher +from . socket import SocketEndpoint +from . serialize import to_subgraph + +class TriplesLoadEndpoint(SocketEndpoint): + + def __init__(self, pulsar_host, path="/api/v1/load/triples"): + + super(TriplesLoadEndpoint, self).__init__( + endpoint_path=path + ) + + self.pulsar_host=pulsar_host + + self.publisher = Publisher( + self.pulsar_host, triples_store_queue, + schema=JsonSchema(Triples) + ) + + async def start(self): + + self.task = asyncio.create_task( + self.publisher.run() + ) + + async def listener(self, ws, running): + + async for msg in ws: + # On error, finish + if msg.type == WSMsgType.ERROR: + break + else: + + data = msg.json() + + elt = Triples( + metadata=Metadata( + id=data["metadata"]["id"], + metadata=to_subgraph(data["metadata"]["metadata"]), + user=data["metadata"]["user"], + collection=data["metadata"]["collection"], + ), + triples=to_subgraph(data["triples"]), + ) + + await self.publisher.send(None, elt) + + + running.stop() diff --git a/trustgraph-flow/trustgraph/api/gateway/triples_stream.py b/trustgraph-flow/trustgraph/api/gateway/triples_stream.py new file mode 100644 index 00000000..e8b538a4 --- /dev/null +++ b/trustgraph-flow/trustgraph/api/gateway/triples_stream.py @@ -0,0 +1,56 @@ + +import asyncio +from pulsar.schema import JsonSchema +import uuid + +from ... schema import Triples +from ... schema import triples_store_queue + +from . subscriber import Subscriber +from . socket import SocketEndpoint +from . serialize import serialize_triples + +class TriplesStreamEndpoint(SocketEndpoint): + + def __init__(self, pulsar_host, path="/api/v1/stream/triples"): + + super(TriplesStreamEndpoint, self).__init__( + endpoint_path=path + ) + + self.pulsar_host=pulsar_host + + self.subscriber = Subscriber( + self.pulsar_host, triples_store_queue, + "api-gateway", "api-gateway", + schema=JsonSchema(Triples) + ) + + async def start(self): + + self.task = asyncio.create_task( + self.subscriber.run() + ) + + async def async_thread(self, ws, running): + + id = str(uuid.uuid4()) + + q = await self.subscriber.subscribe_all(id) + + while running.get(): + try: + resp = await asyncio.wait_for(q.get(), 0.5) + await ws.send_json(serialize_triples(resp)) + + except TimeoutError: + continue + + except Exception as e: + print(f"Exception: {str(e)}", flush=True) + break + + await self.subscriber.unsubscribe_all(id) + + running.stop() +