From 6ced8dbba9fdec00da616228acea370c25b24e45 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Fri, 29 Nov 2024 22:56:04 +0000 Subject: [PATCH] Working encyclopedia lookup --- trustgraph-base/trustgraph/schema/lookup.py | 14 +++- trustgraph-flow/scripts/wikipedia-lookup | 6 ++ trustgraph-flow/setup.py | 1 + .../trustgraph/api/gateway/service.py | 67 ++++++++++++++++ .../trustgraph/external/wikipedia/__init__.py | 3 + .../trustgraph/external/wikipedia/__main__.py | 7 ++ .../trustgraph/external/wikipedia/service.py | 78 ++++++++++--------- 7 files changed, 135 insertions(+), 41 deletions(-) create mode 100755 trustgraph-flow/scripts/wikipedia-lookup create mode 100644 trustgraph-flow/trustgraph/external/wikipedia/__main__.py diff --git a/trustgraph-base/trustgraph/schema/lookup.py b/trustgraph-base/trustgraph/schema/lookup.py index f8aaf5fc..d0a0517c 100644 --- a/trustgraph-base/trustgraph/schema/lookup.py +++ b/trustgraph-base/trustgraph/schema/lookup.py @@ -1,5 +1,5 @@ -from pulsar.schema import Record, Bytes, String, Boolean, Integer, Array, Double +from pulsar.schema import Record, String from . types import Error, Value, Triple from . topic import topic @@ -15,14 +15,22 @@ class LookupRequest(Record): class LookupResponse(Record): text = String() + error = Error() -wikipedia_lookup_request_queue = topic( +encyclopedia_lookup_request_queue = topic( 'encyclopedia', kind='non-persistent', namespace='request' ) -wikipedia_lookup_response_queue = topic( +encyclopedia_lookup_response_queue = topic( 'encyclopedia', kind='non-persistent', namespace='response', ) +dbpedia_lookup_request_queue = topic( + 'dbpedia', kind='non-persistent', namespace='request' +) +dbpedia_lookup_response_queue = topic( + 'dbpedia', kind='non-persistent', namespace='response', +) + internet_search_request_queue = topic( 'internet-search', kind='non-persistent', namespace='request' ) diff --git a/trustgraph-flow/scripts/wikipedia-lookup b/trustgraph-flow/scripts/wikipedia-lookup new file mode 100755 index 00000000..a89b1009 --- /dev/null +++ b/trustgraph-flow/scripts/wikipedia-lookup @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from trustgraph.external.wikipedia import run + +run() + diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index 8e81e12c..65bb7326 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -106,5 +106,6 @@ setuptools.setup( "scripts/triples-query-neo4j", "scripts/triples-write-cassandra", "scripts/triples-write-neo4j", + "scripts/wikipedia-lookup", ] ) diff --git a/trustgraph-flow/trustgraph/api/gateway/service.py b/trustgraph-flow/trustgraph/api/gateway/service.py index 7b12e1a2..677d2303 100755 --- a/trustgraph-flow/trustgraph/api/gateway/service.py +++ b/trustgraph-flow/trustgraph/api/gateway/service.py @@ -63,6 +63,10 @@ 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 document_ingest_queue, text_ingest_queue logger = logging.getLogger("api") @@ -347,12 +351,24 @@ class Api: chunking_enabled=True, ) + self.encyclopedia_lookup_out = Publisher( + self.pulsar_host, encyclopedia_lookup_request_queue, + schema=JsonSchema(LookupRequest) + ) + + self.encyclopedia_lookup_in = Subscriber( + self.pulsar_host, encyclopedia_lookup_response_queue, + "api-gateway", "api-gateway", + JsonSchema(LookupResponse) + ) + self.app.add_routes([ web.post("/api/v1/text-completion", self.llm), web.post("/api/v1/prompt", self.prompt), web.post("/api/v1/graph-rag", self.graph_rag), web.post("/api/v1/triples-query", self.triples_query), web.post("/api/v1/agent", self.agent), + web.post("/api/v1/encyclopedia", self.encyclopedia), web.post("/api/v1/embeddings", self.embeddings), web.post("/api/v1/load/document", self.load_document), web.post("/api/v1/load/text", self.load_text), @@ -664,6 +680,50 @@ class Api: finally: await self.embeddings_in.unsubscribe(id) + async def encyclopedia(self, request): + + id = str(uuid.uuid4()) + + try: + + data = await request.json() + + print(data) + + 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 load_document(self, request): try: @@ -961,6 +1021,13 @@ class Api: self.text_ingest_pub_task = asyncio.create_task(self.text_out.run()) + self.encyclopedia_pub_task = asyncio.create_task( + self.encyclopedia_lookup_out.run() + ) + self.encyclopedia_sub_task = asyncio.create_task( + self.encyclopedia_lookup_in.run() + ) + return self.app def run(self): diff --git a/trustgraph-flow/trustgraph/external/wikipedia/__init__.py b/trustgraph-flow/trustgraph/external/wikipedia/__init__.py index e69de29b..ba844705 100644 --- a/trustgraph-flow/trustgraph/external/wikipedia/__init__.py +++ b/trustgraph-flow/trustgraph/external/wikipedia/__init__.py @@ -0,0 +1,3 @@ + +from . service import * + diff --git a/trustgraph-flow/trustgraph/external/wikipedia/__main__.py b/trustgraph-flow/trustgraph/external/wikipedia/__main__.py new file mode 100644 index 00000000..e9136855 --- /dev/null +++ b/trustgraph-flow/trustgraph/external/wikipedia/__main__.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 + +from . service import run + +if __name__ == '__main__': + run() + diff --git a/trustgraph-flow/trustgraph/external/wikipedia/service.py b/trustgraph-flow/trustgraph/external/wikipedia/service.py index 4b3b39c1..932e1213 100644 --- a/trustgraph-flow/trustgraph/external/wikipedia/service.py +++ b/trustgraph-flow/trustgraph/external/wikipedia/service.py @@ -1,23 +1,22 @@ """ -Embeddings service, applies an embeddings model selected from HuggingFace. -Input is text, output is embeddings vector. +Wikipedia lookup service. Fetchs an extract from the Wikipedia page +using the API. """ -from langchain_huggingface import HuggingFaceEmbeddings - -from trustgraph.schema import EmbeddingsRequest, EmbeddingsResponse, Error -from trustgraph.schema import embeddings_request_queue -from trustgraph.schema import embeddings_response_queue +from trustgraph.schema import LookupRequest, LookupResponse, Error +from trustgraph.schema import encyclopedia_lookup_request_queue +from trustgraph.schema import encyclopedia_lookup_response_queue from trustgraph.log_level import LogLevel from trustgraph.base import ConsumerProducer +import requests module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = embeddings_request_queue -default_output_queue = embeddings_response_queue +default_input_queue = encyclopedia_lookup_request_queue +default_output_queue = encyclopedia_lookup_response_queue default_subscriber = module -default_model="all-MiniLM-L6-v2" +default_url="https://en.wikipedia.org/" class Processor(ConsumerProducer): @@ -26,19 +25,19 @@ class Processor(ConsumerProducer): input_queue = params.get("input_queue", default_input_queue) output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) - model = params.get("model", default_model) + url = params.get("url", default_url) super(Processor, self).__init__( **params | { "input_queue": input_queue, "output_queue": output_queue, "subscriber": subscriber, - "input_schema": EmbeddingsRequest, - "output_schema": EmbeddingsResponse, + "input_schema": LookupRequest, + "output_schema": LookupResponse, } ) - self.embeddings = HuggingFaceEmbeddings(model_name=model) + self.url = url def handle(self, msg): @@ -47,37 +46,40 @@ class Processor(ConsumerProducer): # Sender-produced ID id = msg.properties()["id"] - print(f"Handling input {id}...", flush=True) + print(f"Handling {v.kind} / {v.term}...", flush=True) try: - text = v.text - embeds = self.embeddings.embed_documents([text]) + url = f"{self.url}/api/rest_v1/page/summary/{v.term}" - print("Send response...", flush=True) - r = EmbeddingsResponse(vectors=embeds, error=None) - self.producer.send(r, properties={"id": id}) + resp = Result = requests.get(url).json() + resp = resp["extract"] - print("Done.", flush=True) - - - except Exception as e: - - print(f"Exception: {e}") - - print("Send error response...", flush=True) - - r = EmbeddingsResponse( - error=Error( - type = "llm-error", - message = str(e), - ), - response=None, + r = LookupResponse( + error=None, + text=resp ) self.producer.send(r, properties={"id": id}) self.consumer.acknowledge(msg) + + return + + except Exception as e: + + r = LookupResponse( + error=Error( + type = "lookup-error", + message = str(e), + ), + text=None, + ) + self.producer.send(r, properties={"id": id}) + + self.consumer.acknowledge(msg) + + return @staticmethod @@ -89,9 +91,9 @@ class Processor(ConsumerProducer): ) parser.add_argument( - '-m', '--model', - default="all-MiniLM-L6-v2", - help=f'LLM model (default: all-MiniLM-L6-v2)' + '-u', '--url', + default=default_url, + help=f'LLM model (default: {default_url})' ) def run():