Edge-RPC-RAG

A local Retrieval-Augmented Generation engine for VRAM-constrained edge hardware (built and tested against an NVIDIA GTX 1050, 4GB VRAM), with priority-based task preemption and streaming structured output.

---

Overview

Edge-RPC-RAG combines a priority-preemptible gRPC inference queue with a tolerant partial-JSON parser on the frontend, so structured UI output can render incrementally while a response is still streaming, and so a higher-priority query can interrupt and evict a lower-priority one mid-generation, rather than waiting for it to finish.

The demo UI shown here uses a medical-triage example workload (patient vitals, history queries) to illustrate the preemption behavior; the underlying engine is domain-agnostic.

---

Preemption behavior

A routine, longer-running query (Level 2) is dispatched and begins streaming tokens into the UI.

While it's still generating, a critical query (Level 0) arrives from a concurrent session:

Interrupt signal:

Eviction shown in primary UI:

The gRPC engine detects the Level 0 signal, aborts the lower-priority generation, evicts it from VRAM, and begins servicing the critical query with no queue wait.

---

Architecture

┌────────────────────────────────────────────────────────┐
│                   Next.js 14 control UI                 │
│         (partial-JSON repair, reactive UI cards)        │
└───────────────────────────┬────────────────────────────┘
                            │ WebSockets (JSON patches)
                            ▼
┌────────────────────────────────────────────────────────┐
│               FastAPI WebSocket bridge                  │
│         (translates sockets to streaming gRPC)          │
└───────────────────────────┬────────────────────────────┘
                            │ gRPC / Protocol Buffers (port 50051)
                            ▼
┌────────────────────────────────────────────────────────┐
│               Python core engine                        │
│      (priority preemption queue, hybrid retriever)      │
└───────────────────────────┬────────────────────────────┘
                            │ CUDA
                            ▼
┌────────────────────────────────────────────────────────┐
│           NVIDIA GTX 1050 (4GB VRAM)                     │
└───────────────────────────┴────────────────────────────┘

---

System priority levels

Level	Name	Behavior on arrival
0	Critical	Stops any running 1/2 task immediately, evicts its VRAM KV cache, and streams the critical request right away.
1	High	Placed at the head of the pending queue. Does not evict an already-running lower-priority task mid-stream.
2	Routine	Default execution level. Subject to eviction if a 0 arrives.

The three components that implement this:

• Priority preemption queue: three levels as above, backed by an async gRPC engine that checks for higher-priority arrivals between generation steps.
• Streaming partial-JSON repair: a tolerant frontend parser (partialParseJSON) that repairs unclosed strings/objects/arrays on each incoming token, so structured UI components render before generation finishes rather than only after.
• Async gRPC pipeline: keeps the CUDA inference layer decoupled from the presentation layer via protobuf messages and a WebSocket translation bridge.

---

Repository structure

Edge-RPC-RAG/
├── api_bridge.py          # FastAPI WebSocket bridge
├── server.py              # Core gRPC server
├── retriever.py           # Hybrid retrieval pipeline
├── ingest.py              # Index creation & data ingest
├── templates.py           # GBNF schemas and prompt templates
├── docs/screenshots/      # Demonstration screenshots
├── edge-rag-ui/           # Next.js 14 web interface
├── proto/                 # Protocol Buffer specifications
├── tests/                 # Integration, preemption, & load tests
└── client/                # Standalone streaming client utilities

Tech stack

Layer	Technology	Role
Frontend	Next.js 14, React, Tailwind CSS	Dashboard, live JSON repair, dynamic UI card rendering
Bridge	FastAPI, WebSockets, Uvicorn	Async proxy translating browser sockets to gRPC
Engine	Python, gRPC, Protobuf	Preemption signaling, priority queue, hybrid search
Target	NVIDIA GTX 1050 (CUDA)	VRAM-constrained hardware target

---

Setup

Prerequisites

• Python 3.10+
• Node.js 18+
• CUDA Toolkit installed and configured

Python environment

git clone https://github.com/mustafa-5493/Edge-RPC-RAG.git
cd Edge-RPC-RAG

python -m venv .venv
.\.venv\Scripts\Activate.ps1
pip install -r requirements.txt

Next.js control panel

cd edge-rag-ui
npm install
cd ..

---

Running the full stack

Three separate terminals:

Terminal 1 — core gRPC engine

.\.venv\Scripts\Activate.ps1
python server.py

Expected output: [*] Edge-RAG gRPC Server listening on port 50051...

Terminal 2 — FastAPI WebSocket bridge

.\.venv\Scripts\Activate.ps1
uvicorn api_bridge:app --host 0.0.0.0 --port 8000

Expected output: Uvicorn running on http://0.0.0.0:8000

Terminal 3 — Next.js UI

cd edge-rag-ui
npm run dev

Navigate to http://localhost:3000.

---

Testing and verification

With Terminals 1–3 running, open a fourth terminal to simulate workloads and verify preemption behavior:

1. Dispatch a routine stream (Priority Level 2):

.\.venv\Scripts\Activate.ps1
python -m tests.test_client

2. Send a critical preemption interrupt (Priority Level 0):

.\.venv\Scripts\Activate.ps1
python -m tests.test_preemption

3. Run the full test suite:

.\.venv\Scripts\Activate.ps1
pytest

---

License

Distributed under the MIT License. See LICENSE for details.