|
| 1 | +// Copyright (c) 2026 vectorless developers |
| 2 | +// SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +//! PyO3 streaming query wrapper. |
| 5 | +//! |
| 6 | +//! Bridges Rust's `mpsc::Receiver<RetrieveEvent>` to a Python async iterator, |
| 7 | +//! yielding real-time retrieval progress events as dicts. |
| 8 | +
|
| 9 | +use pyo3::exceptions::PyStopAsyncIteration; |
| 10 | +use pyo3::prelude::*; |
| 11 | +use pyo3::types::PyDict; |
| 12 | +use pyo3_async_runtimes::tokio::future_into_py; |
| 13 | +use std::sync::Arc; |
| 14 | +use tokio::sync::{mpsc, Mutex}; |
| 15 | + |
| 16 | +use ::vectorless::retrieval::{RetrieveEvent, SufficiencyLevel}; |
| 17 | + |
| 18 | +/// Convert a `RetrieveEvent` into a Python dict with a `"type"` key. |
| 19 | +fn event_to_dict(event: RetrieveEvent, py: Python<'_>) -> PyObject { |
| 20 | + let dict = PyDict::new(py); |
| 21 | + match event { |
| 22 | + RetrieveEvent::Started { query, strategy } => { |
| 23 | + dict.set_item("type", "started").unwrap(); |
| 24 | + dict.set_item("query", query).unwrap(); |
| 25 | + dict.set_item("strategy", strategy).unwrap(); |
| 26 | + } |
| 27 | + RetrieveEvent::StageCompleted { stage, elapsed_ms } => { |
| 28 | + dict.set_item("type", "stage_completed").unwrap(); |
| 29 | + dict.set_item("stage", stage).unwrap(); |
| 30 | + dict.set_item("elapsed_ms", elapsed_ms).unwrap(); |
| 31 | + } |
| 32 | + RetrieveEvent::NodeVisited { |
| 33 | + node_id, |
| 34 | + title, |
| 35 | + score, |
| 36 | + } => { |
| 37 | + dict.set_item("type", "node_visited").unwrap(); |
| 38 | + dict.set_item("node_id", node_id).unwrap(); |
| 39 | + dict.set_item("title", title).unwrap(); |
| 40 | + dict.set_item("score", score).unwrap(); |
| 41 | + } |
| 42 | + RetrieveEvent::ContentFound { |
| 43 | + node_id, |
| 44 | + title, |
| 45 | + preview, |
| 46 | + score, |
| 47 | + } => { |
| 48 | + dict.set_item("type", "content_found").unwrap(); |
| 49 | + dict.set_item("node_id", node_id).unwrap(); |
| 50 | + dict.set_item("title", title).unwrap(); |
| 51 | + dict.set_item("preview", preview).unwrap(); |
| 52 | + dict.set_item("score", score).unwrap(); |
| 53 | + } |
| 54 | + RetrieveEvent::Backtracking { from, to, reason } => { |
| 55 | + dict.set_item("type", "backtracking").unwrap(); |
| 56 | + dict.set_item("from", from).unwrap(); |
| 57 | + dict.set_item("to", to).unwrap(); |
| 58 | + dict.set_item("reason", reason).unwrap(); |
| 59 | + } |
| 60 | + RetrieveEvent::SufficiencyCheck { level, tokens } => { |
| 61 | + let level_str = match level { |
| 62 | + SufficiencyLevel::Sufficient => "sufficient", |
| 63 | + SufficiencyLevel::PartialSufficient => "partial_sufficient", |
| 64 | + SufficiencyLevel::Insufficient => "insufficient", |
| 65 | + }; |
| 66 | + dict.set_item("type", "sufficiency_check").unwrap(); |
| 67 | + dict.set_item("level", level_str).unwrap(); |
| 68 | + dict.set_item("tokens", tokens).unwrap(); |
| 69 | + } |
| 70 | + RetrieveEvent::Completed { response } => { |
| 71 | + dict.set_item("type", "completed").unwrap(); |
| 72 | + dict.set_item("confidence", response.confidence).unwrap(); |
| 73 | + dict.set_item("is_sufficient", response.is_sufficient).unwrap(); |
| 74 | + dict.set_item("strategy_used", response.strategy_used).unwrap(); |
| 75 | + dict.set_item("tokens_used", response.tokens_used).unwrap(); |
| 76 | + dict.set_item("content", response.content).unwrap(); |
| 77 | + |
| 78 | + let results: Vec<PyObject> = response |
| 79 | + .results |
| 80 | + .into_iter() |
| 81 | + .map(|r| { |
| 82 | + let rd = PyDict::new(py); |
| 83 | + rd.set_item("node_id", &r.node_id).unwrap(); |
| 84 | + rd.set_item("title", &r.title).unwrap(); |
| 85 | + rd.set_item("content", &r.content).unwrap(); |
| 86 | + rd.set_item("score", r.score).unwrap(); |
| 87 | + rd.set_item("depth", r.depth).unwrap(); |
| 88 | + rd.into() |
| 89 | + }) |
| 90 | + .collect(); |
| 91 | + dict.set_item("results", results).unwrap(); |
| 92 | + } |
| 93 | + RetrieveEvent::Error { message } => { |
| 94 | + dict.set_item("type", "error").unwrap(); |
| 95 | + dict.set_item("message", message).unwrap(); |
| 96 | + } |
| 97 | + } |
| 98 | + dict.into() |
| 99 | +} |
| 100 | + |
| 101 | +/// Python-facing async iterator over streaming retrieval events. |
| 102 | +/// |
| 103 | +/// Usage:: |
| 104 | +/// |
| 105 | +/// stream = await engine.query_stream(ctx) |
| 106 | +/// async for event in stream: |
| 107 | +/// print(event["type"]) |
| 108 | +#[pyclass(name = "StreamingQuery")] |
| 109 | +pub struct PyStreamingQuery { |
| 110 | + rx: Arc<Mutex<Option<mpsc::Receiver<RetrieveEvent>>>>, |
| 111 | +} |
| 112 | + |
| 113 | +impl PyStreamingQuery { |
| 114 | + pub fn new(rx: mpsc::Receiver<RetrieveEvent>) -> Self { |
| 115 | + Self { |
| 116 | + rx: Arc::new(Mutex::new(Some(rx))), |
| 117 | + } |
| 118 | + } |
| 119 | +} |
| 120 | + |
| 121 | +#[pymethods] |
| 122 | +impl PyStreamingQuery { |
| 123 | + fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { |
| 124 | + slf |
| 125 | + } |
| 126 | + |
| 127 | + fn __anext__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> { |
| 128 | + let rx = Arc::clone(&self.rx); |
| 129 | + future_into_py(py, async move { |
| 130 | + let mut guard = rx.lock().await; |
| 131 | + match guard.as_mut() { |
| 132 | + None => Err(PyStopAsyncIteration::new_err("stream exhausted")), |
| 133 | + Some(receiver) => match receiver.recv().await { |
| 134 | + Some(event) => { |
| 135 | + let is_terminal = matches!( |
| 136 | + &event, |
| 137 | + RetrieveEvent::Completed { .. } | RetrieveEvent::Error { .. } |
| 138 | + ); |
| 139 | + if is_terminal { |
| 140 | + *guard = None; |
| 141 | + } |
| 142 | + // Convert to Python dict — safe because future_into_py |
| 143 | + // ensures we're on a thread that can acquire the GIL. |
| 144 | + let obj = Python::with_gil(|py| event_to_dict(event, py)); |
| 145 | + Ok(obj) |
| 146 | + } |
| 147 | + None => { |
| 148 | + *guard = None; |
| 149 | + Err(PyStopAsyncIteration::new_err("stream closed")) |
| 150 | + } |
| 151 | + }, |
| 152 | + } |
| 153 | + }) |
| 154 | + } |
| 155 | + |
| 156 | + fn __repr__(&self) -> String { |
| 157 | + "StreamingQuery(...)".to_string() |
| 158 | + } |
| 159 | +} |
0 commit comments