Working encyclopedia lookup

This commit is contained in:
Cyber MacGeddon 2024-11-29 22:56:04 +00:00
parent a4dd1c8fa3
commit 6ced8dbba9
7 changed files with 135 additions and 41 deletions

View file

@ -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 . types import Error, Value, Triple
from . topic import topic from . topic import topic
@ -15,14 +15,22 @@ class LookupRequest(Record):
class LookupResponse(Record): class LookupResponse(Record):
text = String() text = String()
error = Error()
wikipedia_lookup_request_queue = topic( encyclopedia_lookup_request_queue = topic(
'encyclopedia', kind='non-persistent', namespace='request' 'encyclopedia', kind='non-persistent', namespace='request'
) )
wikipedia_lookup_response_queue = topic( encyclopedia_lookup_response_queue = topic(
'encyclopedia', kind='non-persistent', namespace='response', '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_request_queue = topic(
'internet-search', kind='non-persistent', namespace='request' 'internet-search', kind='non-persistent', namespace='request'
) )

View file

@ -0,0 +1,6 @@
#!/usr/bin/env python3
from trustgraph.external.wikipedia import run
run()

View file

@ -106,5 +106,6 @@ setuptools.setup(
"scripts/triples-query-neo4j", "scripts/triples-query-neo4j",
"scripts/triples-write-cassandra", "scripts/triples-write-cassandra",
"scripts/triples-write-neo4j", "scripts/triples-write-neo4j",
"scripts/wikipedia-lookup",
] ]
) )

View file

@ -63,6 +63,10 @@ from ... schema import EmbeddingsRequest, EmbeddingsResponse
from ... schema import embeddings_request_queue from ... schema import embeddings_request_queue
from ... schema import embeddings_response_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 from ... schema import document_ingest_queue, text_ingest_queue
logger = logging.getLogger("api") logger = logging.getLogger("api")
@ -347,12 +351,24 @@ class Api:
chunking_enabled=True, 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([ self.app.add_routes([
web.post("/api/v1/text-completion", self.llm), web.post("/api/v1/text-completion", self.llm),
web.post("/api/v1/prompt", self.prompt), web.post("/api/v1/prompt", self.prompt),
web.post("/api/v1/graph-rag", self.graph_rag), web.post("/api/v1/graph-rag", self.graph_rag),
web.post("/api/v1/triples-query", self.triples_query), web.post("/api/v1/triples-query", self.triples_query),
web.post("/api/v1/agent", self.agent), 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/embeddings", self.embeddings),
web.post("/api/v1/load/document", self.load_document), web.post("/api/v1/load/document", self.load_document),
web.post("/api/v1/load/text", self.load_text), web.post("/api/v1/load/text", self.load_text),
@ -664,6 +680,50 @@ class Api:
finally: finally:
await self.embeddings_in.unsubscribe(id) 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): async def load_document(self, request):
try: try:
@ -961,6 +1021,13 @@ class Api:
self.text_ingest_pub_task = asyncio.create_task(self.text_out.run()) 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 return self.app
def run(self): def run(self):

View file

@ -0,0 +1,3 @@
from . service import *

View file

@ -0,0 +1,7 @@
#!/usr/bin/env python3
from . service import run
if __name__ == '__main__':
run()

View file

@ -1,23 +1,22 @@
""" """
Embeddings service, applies an embeddings model selected from HuggingFace. Wikipedia lookup service. Fetchs an extract from the Wikipedia page
Input is text, output is embeddings vector. using the API.
""" """
from langchain_huggingface import HuggingFaceEmbeddings from trustgraph.schema import LookupRequest, LookupResponse, Error
from trustgraph.schema import encyclopedia_lookup_request_queue
from trustgraph.schema import EmbeddingsRequest, EmbeddingsResponse, Error from trustgraph.schema import encyclopedia_lookup_response_queue
from trustgraph.schema import embeddings_request_queue
from trustgraph.schema import embeddings_response_queue
from trustgraph.log_level import LogLevel from trustgraph.log_level import LogLevel
from trustgraph.base import ConsumerProducer from trustgraph.base import ConsumerProducer
import requests
module = ".".join(__name__.split(".")[1:-1]) module = ".".join(__name__.split(".")[1:-1])
default_input_queue = embeddings_request_queue default_input_queue = encyclopedia_lookup_request_queue
default_output_queue = embeddings_response_queue default_output_queue = encyclopedia_lookup_response_queue
default_subscriber = module default_subscriber = module
default_model="all-MiniLM-L6-v2" default_url="https://en.wikipedia.org/"
class Processor(ConsumerProducer): class Processor(ConsumerProducer):
@ -26,19 +25,19 @@ class Processor(ConsumerProducer):
input_queue = params.get("input_queue", default_input_queue) input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue) output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber) subscriber = params.get("subscriber", default_subscriber)
model = params.get("model", default_model) url = params.get("url", default_url)
super(Processor, self).__init__( super(Processor, self).__init__(
**params | { **params | {
"input_queue": input_queue, "input_queue": input_queue,
"output_queue": output_queue, "output_queue": output_queue,
"subscriber": subscriber, "subscriber": subscriber,
"input_schema": EmbeddingsRequest, "input_schema": LookupRequest,
"output_schema": EmbeddingsResponse, "output_schema": LookupResponse,
} }
) )
self.embeddings = HuggingFaceEmbeddings(model_name=model) self.url = url
def handle(self, msg): def handle(self, msg):
@ -47,38 +46,41 @@ class Processor(ConsumerProducer):
# Sender-produced ID # Sender-produced ID
id = msg.properties()["id"] id = msg.properties()["id"]
print(f"Handling input {id}...", flush=True) print(f"Handling {v.kind} / {v.term}...", flush=True)
try: try:
text = v.text url = f"{self.url}/api/rest_v1/page/summary/{v.term}"
embeds = self.embeddings.embed_documents([text])
print("Send response...", flush=True) resp = Result = requests.get(url).json()
r = EmbeddingsResponse(vectors=embeds, error=None) resp = resp["extract"]
self.producer.send(r, properties={"id": id})
print("Done.", flush=True) r = LookupResponse(
error=None,
text=resp
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,
) )
self.producer.send(r, properties={"id": id}) self.producer.send(r, properties={"id": id})
self.consumer.acknowledge(msg) 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 @staticmethod
def add_args(parser): def add_args(parser):
@ -89,9 +91,9 @@ class Processor(ConsumerProducer):
) )
parser.add_argument( parser.add_argument(
'-m', '--model', '-u', '--url',
default="all-MiniLM-L6-v2", default=default_url,
help=f'LLM model (default: all-MiniLM-L6-v2)' help=f'LLM model (default: {default_url})'
) )
def run(): def run():