Added agent API

This commit is contained in:
Cyber MacGeddon 2024-11-20 18:59:28 +00:00
parent e76cdc0f12
commit f96cf2a682
3 changed files with 101 additions and 4 deletions

28
test-agent-api Executable file
View file

@ -0,0 +1,28 @@
#!/usr/bin/env python3
import requests
import json
import sys
url = "http://localhost:8088/api/v1/"
############################################################################
input = {
"question": "What is the highest risk aspect of running a space shuttle program? Provide 5 detailed reasons to justify our answer.",
}
resp = requests.post(
f"{url}agent",
json=input,
)
resp = resp.json()
if "error" in resp:
print(f"Error: {resp['error']}")
sys.exit(1)
print(resp["answer"])

View file

@ -30,5 +30,6 @@ if "error" in resp:
print(resp["response"]) print(resp["response"])
sys.exit(0) sys.exit(0)
############################################################################ ############################################################################

View file

@ -31,6 +31,10 @@ from trustgraph.schema import TriplesQueryRequest, TriplesQueryResponse, Value
from trustgraph.schema import triples_request_queue from trustgraph.schema import triples_request_queue
from trustgraph.schema import triples_response_queue from trustgraph.schema import triples_response_queue
from trustgraph.schema import AgentRequest, AgentResponse
from trustgraph.schema import agent_request_queue
from trustgraph.schema import agent_response_queue
logger = logging.getLogger("api") logger = logging.getLogger("api")
logger.setLevel(logging.INFO) logger.setLevel(logging.INFO)
@ -54,7 +58,7 @@ class Publisher:
while True: while True:
id, item = await self.q.get() id, item = await self.q.get()
await producer.send(item, { "id": id }) await producer.send(item, { "id": id })
print("message out") # print("message out")
async def send(self, id, msg): async def send(self, id, msg):
await self.q.put((id, msg)) await self.q.put((id, msg))
@ -80,7 +84,7 @@ class Subscriber:
) as consumer: ) as consumer:
while True: while True:
msg = await consumer.receive() msg = await consumer.receive()
print("message in") # print("message in", self.topic)
id = msg.properties()["id"] id = msg.properties()["id"]
value = msg.value() value = msg.value()
if id in self.q: if id in self.q:
@ -146,11 +150,23 @@ class Api:
JsonSchema(TriplesQueryResponse) JsonSchema(TriplesQueryResponse)
) )
self.agent_out = Publisher(
pulsar_host, agent_request_queue,
schema=JsonSchema(AgentRequest)
)
self.agent_in = Subscriber(
pulsar_host, agent_response_queue,
"api-gateway", "api-gateway",
JsonSchema(AgentResponse)
)
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),
]) ])
async def llm(self, request): async def llm(self, request):
@ -298,8 +314,6 @@ class Api:
data = await request.json() data = await request.json()
print(data)
q = await self.triples_query_in.subscribe(id) q = await self.triples_query_in.subscribe(id)
if "s" in data: if "s" in data:
@ -380,6 +394,57 @@ class Api:
finally: finally:
await self.graph_rag_in.unsubscribe(id) await self.graph_rag_in.unsubscribe(id)
async def agent(self, request):
id = str(uuid.uuid4())
try:
data = await request.json()
q = await self.agent_in.subscribe(id)
await self.agent_out.send(
id,
AgentRequest(
question=data["question"],
)
)
while True:
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 }
)
if resp.answer: break
if resp.thought: print("thought:", resp.thought)
if resp.observation: print("observation:", resp.observation)
if resp.answer:
return web.json_response(
{ "answer": resp.answer }
)
# Can't happen, ook at the logic
raise RuntimeError("Strange state")
except Exception as e:
logging.error(f"Exception: {e}")
return web.json_response(
{ "error": str(e) }
)
finally:
await self.agent_in.unsubscribe(id)
async def app_factory(self): async def app_factory(self):
self.llm_pub_task = asyncio.create_task(self.llm_in.run()) self.llm_pub_task = asyncio.create_task(self.llm_in.run())
@ -398,6 +463,9 @@ class Api:
self.triples_query_out.run() self.triples_query_out.run()
) )
self.agent_pub_task = asyncio.create_task(self.agent_in.run())
self.agent_sub_task = asyncio.create_task(self.agent_out.run())
return self.app return self.app
def run(self): def run(self):