Working encyclopedia lookup

trustgraph-flow/scripts/wikipedia-lookup

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

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",
     ]
 )

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):

trustgraph-flow/trustgraph/external/wikipedia/__init__.py

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

trustgraph-flow/trustgraph/external/wikipedia/__main__.py

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

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():