Skills › Scientific Agent Skills › open-notebook › source_ingestion.py
source_ingestion.py
144 lines · skills/open-notebook/scripts/source_ingestion.py at 5f2bc03
1"""Source ingestion helpers; CLI lists sources without creating demo records.
3Requires requests and a running Open Notebook backend with a processing worker.
4Set OPEN_NOTEBOOK_URL and, when enabled, OPEN_NOTEBOOK_PASSWORD.
5"""
7import argparse
8import json
9import math
10from pathlib import Path
11import time
13from _common import record_path, request_json
16def _source_fields(notebook_id, source_type, process_async, embed):
17 return {
18 "type": source_type,
19 "notebooks": json.dumps([notebook_id]),
20 "async_processing": str(process_async).lower(),
21 "embed": str(embed).lower(),
22 }
25def add_url_source(notebook_id, url, process_async=True, embed=False):
26 data = _source_fields(notebook_id, "link", process_async, embed)
27 data["url"] = url
28 return request_json("POST", "/sources", data=data)
31def add_text_source(notebook_id, title, text, process_async=False, embed=False):
32 data = _source_fields(notebook_id, "text", process_async, embed)
33 data.update(title=title, content=text)
34 return request_json("POST", "/sources", data=data)
37def upload_file_source(notebook_id, file_path, process_async=True, embed=False):
38 path = Path(file_path)
39 with path.open("rb") as handle:
40 return request_json("POST", "/sources",
41 data=_source_fields(notebook_id, "upload", process_async, embed),
42 files={"file": (path.name, handle)})
45def wait_for_processing(source_id, poll_interval=5, timeout=300):
46 """Return completed status; raise on failure, unknown legacy state or timeout.
48 Status=None with no command can describe synchronous/legacy sources. Confirm
49 full_text exists before treating those as complete. Embeddings need a separate
50 embedded_chunks check if vector retrieval is required.
51 """
52 if not all(math.isfinite(value) and value > 0 for value in (timeout, poll_interval)):
53 raise ValueError("timeout and poll_interval must be finite and positive")
54 deadline = time.monotonic() + timeout
55 path = f"/sources/{record_path(source_id)}"
56 while True:
57 remaining = deadline - time.monotonic()
58 if remaining <= 0:
59 raise TimeoutError(f"Source {source_id} processing timed out")
60 status = request_json("GET", path + "/status", timeout=min(30, remaining))
61 state = status.get("status")
62 if state == "completed":
63 return status
64 if state in {"failed", "error", "canceled", "cancelled"}:
65 raise RuntimeError(f"Source {source_id} processing {state}: {status.get('message', '')}")
66 if state is None and not status.get("command_id"):
67 source = request_json("GET", path, timeout=max(0.001, min(30, deadline - time.monotonic())))
68 if source.get("full_text"):
69 return status
70 raise RuntimeError(f"Source {source_id} has no command and no extracted text")
71 remaining = deadline - time.monotonic()
72 if remaining <= 0:
73 raise TimeoutError(f"Source {source_id} processing timed out")
74 time.sleep(min(poll_interval, remaining))
77def list_sources(notebook_id=None, limit=20, offset=0):
78 """Return one page. /sources supports 1..100 rows and zero-based offsets."""
79 if not 1 <= limit <= 100 or offset < 0:
80 raise ValueError("limit must be 1..100 and offset must be nonnegative")
81 params = {"limit": limit, "offset": offset, "sort_by": "created", "sort_order": "asc"}
82 if notebook_id:
83 params["notebook_id"] = notebook_id
84 return request_json("GET", "/sources", params=params)
87def iter_sources(notebook_id=None, page_size=100, max_pages=1000):
88 """Paginate a stable collection, raising on malformed/repeated pages or the cap.
90 Writes during iteration may shift pages. Discard partial results after an error;
91 an empty or short page is required to establish completion within max_pages.
92 """
93 if not isinstance(max_pages, int) or isinstance(max_pages, bool) or max_pages < 1:
94 raise ValueError("max_pages must be a positive integer")
95 offset = 0
96 seen_ids = set()
97 for _ in range(max_pages):
98 page = list_sources(notebook_id, page_size, offset)
99 if not isinstance(page, list) or len(page) > page_size:
100 raise RuntimeError("Invalid source page: expected a list within the requested limit")
101 page_ids = []
102 for source in page:
103 source_id = source.get("id") if isinstance(source, dict) else None
104 if not isinstance(source_id, str) or not source_id:
105 raise RuntimeError("Invalid source page: each record must have a nonempty string id")
106 page_ids.append(source_id)
107 if len(set(page_ids)) != len(page_ids) or seen_ids.intersection(page_ids):
108 raise RuntimeError("Repeated source IDs: collection changed or pagination did not advance")
109 seen_ids.update(page_ids)
110 yield from page
111 if len(page) < page_size:
112 return
113 offset += len(page)
114 raise RuntimeError("Source pagination reached max_pages; completeness is unknown")
117def get_source_insights(source_id):
118 return request_json("GET", f"/sources/{record_path(source_id)}/insights")
121def retry_failed_source(source_id):
122 """Retry once after inspecting failure; upstream retry always requests embedding."""
123 path = f"/sources/{record_path(source_id)}"
124 status = request_json("GET", path + "/status")
125 if status.get("status") not in {"failed", "error"}:
126 raise ValueError("Retry only a confirmed failed source; inspect active/unknown jobs first")
127 return request_json("POST", path + "/retry")
130def delete_source(source_id):
131 return request_json("DELETE", f"/sources/{record_path(source_id)}")
134def main():
135 parser = argparse.ArgumentParser(description=__doc__)
136 parser.add_argument("--notebook-id")
137 parser.add_argument("--limit", type=int, default=20)
138 parser.add_argument("--offset", type=int, default=0)
139 args = parser.parse_args()
140 print(json.dumps(list_sources(args.notebook_id, args.limit, args.offset), indent=2))
143if __name__ == "__main__":
144 main()