From a8b621f32efaedc23a478f4c4237751400ba57de Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 20 Nov 2024 17:55:01 +0000 Subject: [PATCH] RAG query works --- test-graph-rag-api | 31 ++++++++++++++ test-llm-api | 31 ++++++++++++++ test-prompt-api | 38 +++++++++++++++++ test-prompt2-api | 39 ++++++++++++++++++ trustgraph-flow/scripts/api-gateway | 64 ++++++++++++++++++++++++++++- 5 files changed, 202 insertions(+), 1 deletion(-) create mode 100755 test-graph-rag-api create mode 100755 test-llm-api create mode 100755 test-prompt-api create mode 100755 test-prompt2-api diff --git a/test-graph-rag-api b/test-graph-rag-api new file mode 100755 index 00000000..c329934c --- /dev/null +++ b/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) +############################################################################ + diff --git a/test-llm-api b/test-llm-api new file mode 100755 index 00000000..c33c6634 --- /dev/null +++ b/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"]) + +############################################################################ + diff --git a/test-prompt-api b/test-prompt-api new file mode 100755 index 00000000..1005bc90 --- /dev/null +++ b/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) +############################################################################ + diff --git a/test-prompt2-api b/test-prompt2-api new file mode 100755 index 00000000..f1b80c48 --- /dev/null +++ b/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) +############################################################################ + diff --git a/trustgraph-flow/scripts/api-gateway b/trustgraph-flow/scripts/api-gateway index 508404cc..eaf201d7 100755 --- a/trustgraph-flow/scripts/api-gateway +++ b/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):