RAG query works

test-graph-rag-api

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+
+import requests
+import json
+import sys
+
+url = "http://localhost:8088/api/v1/"
+
+############################################################################
+
+input = {
+    "query": "Give me 10 facts",
+}
+
+resp = requests.post(
+    f"{url}graph-rag",
+    json=input,
+)
+
+resp = resp.json()
+
+print(resp)
+if "error" in resp:
+    print(f"Error: {resp['error']}")
+    sys.exit(1)
+
+print(resp["response"])
+
+sys.exit(0)
+############################################################################
+

test-llm-api

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+
+import requests
+import json
+import sys
+
+url = "http://localhost:8088/api/v1/"
+
+############################################################################
+
+input = {
+    "system": "Respond in French.  Use long word, form of numbers, no digits",
+#     "prompt": "Add 2 and 12"
+    "prompt": "Add 12 and 14, and then make a poem about llamas which incorporates that number.  Then write a joke about llamas"
+}
+
+resp = requests.post(
+    f"{url}text-completion",
+    json=input,
+)
+
+resp = resp.json()
+
+if "error" in resp:
+    print(f"Error: {resp['error']}")
+    sys.exit(1)
+
+print(resp["response"])
+
+############################################################################
+

test-prompt-api

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+
+import requests
+import json
+import sys
+
+url = "http://localhost:8088/api/v1/"
+
+############################################################################
+
+input = {
+    "id": "question",
+    "variables": {
+        "question": "Write a joke about llamas."
+    }
+}
+
+resp = requests.post(
+    f"{url}prompt",
+    json=input,
+)
+
+resp = resp.json()
+
+print(resp)
+if "error" in resp:
+    print(f"Error: {resp['error']}")
+    sys.exit(1)
+
+if "object" in resp:
+    print(f"Object: {resp['object']}")
+    sys.exit(1)
+
+print(resp["text"])
+
+sys.exit(0)
+############################################################################
+

test-prompt2-api

@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+
+import requests
+import json
+import sys
+
+url = "http://localhost:8088/api/v1/"
+
+############################################################################
+
+input = {
+    "id": "extract-definitions",
+    "variables": {
+        "text": "A cat is a large mammal."
+    }
+}
+
+resp = requests.post(
+    f"{url}prompt",
+    json=input,
+)
+
+resp = resp.json()
+
+print(resp)
+if "error" in resp:
+    print(f"Error: {resp['error']}")
+    sys.exit(1)
+
+if "object" in resp:
+    object = json.loads(resp["object"])
+    print(json.dumps(object, indent=4))
+    sys.exit(1)
+
+print(resp["text"])
+
+sys.exit(0)
+############################################################################
+

trustgraph-flow/scripts/api-gateway

@@ -23,6 +23,10 @@ from trustgraph.schema import PromptRequest, PromptResponse
 from trustgraph.schema import prompt_request_queue
 from trustgraph.schema import prompt_response_queue
 
+from trustgraph.schema import GraphRagQuery, GraphRagResponse
+from trustgraph.schema import graph_rag_request_queue
+from trustgraph.schema import graph_rag_response_queue
+
 logger = logging.getLogger("api")
 logger.setLevel(logging.INFO)
 
@@ -116,9 +120,21 @@ class Api:
             JsonSchema(PromptResponse)
         )
 
+        self.graph_rag_out = Publisher(
+            pulsar_host, graph_rag_request_queue,
+            schema=JsonSchema(GraphRagQuery)
+        )
+
+        self.graph_rag_in = Subscriber(
+            pulsar_host, graph_rag_response_queue,
+            "api-gateway", "api-gateway",
+            JsonSchema(GraphRagResponse)
+        )
+
         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),
         ])
 
     async def llm(self, request):
@@ -213,7 +229,50 @@ class Api:
             )
 
         finally:
-            await self.llm_in.unsubscribe(id)
+            await self.prompt_in.unsubscribe(id)
+
+    async def graph_rag(self, request):
+
+        id = str(uuid.uuid4())
+
+        try:
+
+            data = await request.json()
+
+            q = await self.graph_rag_in.subscribe(id)
+
+            await self.graph_rag_out.send(
+                id,
+                GraphRagQuery(
+                    query=data["query"],
+                    user=data.get("user", "trustgraph"),
+                    collection=data.get("collection", "default"),
+                )
+            )
+
+            try:
+                resp = await asyncio.wait_for(q.get(), TIME_OUT)
+            except:
+                raise RuntimeError("Timeout waiting for response")
+
+            if resp.error:
+                return web.json_response(
+                    { "error": resp.error.message }
+                )
+
+            return web.json_response(
+                { "response": resp.response }
+            )
+
+        except Exception as e:
+            logging.error(f"Exception: {e}")
+
+            return web.json_response(
+                { "error": str(e) }
+            )
+
+        finally:
+            await self.graph_rag_in.unsubscribe(id)
 
     async def app_factory(self):
 
@@ -223,6 +282,9 @@ class Api:
         self.prompt_pub_task = asyncio.create_task(self.prompt_in.run())
         self.prompt_sub_task = asyncio.create_task(self.prompt_out.run())
 
+        self.graph_rag_pub_task = asyncio.create_task(self.graph_rag_in.run())
+        self.graph_rag_sub_task = asyncio.create_task(self.graph_rag_out.run())
+
         return self.app
 
     def run(self):