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

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