269 lines
8.2 KiB
Python
269 lines
8.2 KiB
Python
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import sys
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
from devplacepy.services.deepsearch.embeddings import embed_texts, local_embed
|
|
from devplacepy.services.deepsearch.store import Chunk, VectorStore
|
|
from devplacepy.utils import generate_uid
|
|
|
|
from .chunking import chunk_text
|
|
from .crawl import content_hash, crawl, search_queries, url_hash
|
|
from .enhance import plan_queries
|
|
from .orchestrate import orchestrate
|
|
from .phases import (
|
|
PHASE_ANALYSIS,
|
|
PHASE_CRAWLING,
|
|
PHASE_INDEXING,
|
|
PHASE_LABELS,
|
|
PHASE_PLANNING,
|
|
PHASE_SEARCHING,
|
|
PHASE_SYNTHESIS,
|
|
TOTAL_PHASES,
|
|
phase_index,
|
|
)
|
|
|
|
CONTROL_FILE = "control.json"
|
|
PAUSE_POLL_SECONDS = 1.0
|
|
EMBED_BATCH = 64
|
|
FRAME_VERSION = 1
|
|
|
|
_first_frame_sent = False
|
|
|
|
|
|
def _emit(frame: dict) -> None:
|
|
global _first_frame_sent
|
|
if not _first_frame_sent:
|
|
frame = {"version": FRAME_VERSION, **frame}
|
|
_first_frame_sent = True
|
|
sys.stdout.write(json.dumps(frame, ensure_ascii=False) + "\n")
|
|
sys.stdout.flush()
|
|
|
|
|
|
def _stage(stage: str, message: str, phase: str) -> None:
|
|
_emit({"type": "stage", "stage": stage, "message": message})
|
|
_emit(
|
|
{
|
|
"type": "phase",
|
|
"phase": phase,
|
|
"index": phase_index(phase),
|
|
"total": TOTAL_PHASES,
|
|
"label": PHASE_LABELS.get(phase, phase),
|
|
"message": message,
|
|
}
|
|
)
|
|
|
|
|
|
def _read_control(output_dir: Path) -> str:
|
|
path = output_dir / CONTROL_FILE
|
|
if not path.is_file():
|
|
return "running"
|
|
try:
|
|
return str(json.loads(path.read_text(encoding="utf-8")).get("state", "running"))
|
|
except (ValueError, OSError):
|
|
return "running"
|
|
|
|
|
|
def _make_stop(output_dir: Path):
|
|
async def should_stop() -> bool:
|
|
while True:
|
|
state = _read_control(output_dir)
|
|
if state == "cancelled":
|
|
return True
|
|
if state == "paused":
|
|
_emit({"type": "stage", "stage": "paused", "message": "Paused"})
|
|
await asyncio.sleep(PAUSE_POLL_SECONDS)
|
|
continue
|
|
return False
|
|
|
|
return should_stop
|
|
|
|
|
|
async def _index_chunks(
|
|
store: VectorStore, pages: list, api_key: str, emit
|
|
) -> tuple[int, str]:
|
|
chunks: list[Chunk] = []
|
|
for page in pages:
|
|
for position, text in enumerate(chunk_text(page.text)):
|
|
chunks.append(
|
|
Chunk(
|
|
uid=generate_uid(),
|
|
text=text,
|
|
url=page.url,
|
|
title=page.title,
|
|
depth=page.depth,
|
|
source=page.source,
|
|
position=position,
|
|
)
|
|
)
|
|
total = len(chunks)
|
|
if not chunks:
|
|
emit({"type": "embed_done", "backend": "empty", "chunk_count": 0})
|
|
return 0, "empty"
|
|
total_batches = (total + EMBED_BATCH - 1) // EMBED_BATCH
|
|
backend = "gateway"
|
|
forced_local = False
|
|
done = 0
|
|
collected: list[list[float]] = []
|
|
for batch_no, start in enumerate(range(0, total, EMBED_BATCH), start=1):
|
|
batch = chunks[start : start + EMBED_BATCH]
|
|
emit(
|
|
{
|
|
"type": "embed_batch",
|
|
"batch": batch_no,
|
|
"total_batches": total_batches,
|
|
"backend": "local" if forced_local else backend,
|
|
"chunks": len(batch),
|
|
"done": done,
|
|
"total": total,
|
|
"message": f"Embedding batch {batch_no}/{total_batches}",
|
|
}
|
|
)
|
|
if forced_local:
|
|
result = local_embed([chunk.text for chunk in batch])
|
|
else:
|
|
result = await embed_texts([chunk.text for chunk in batch], api_key)
|
|
if not result.vectors or result.backend != "gateway":
|
|
forced_local = True
|
|
result = local_embed([chunk.text for chunk in batch])
|
|
backend = result.backend
|
|
collected.extend(result.vectors)
|
|
done += len(batch)
|
|
emit(
|
|
{
|
|
"type": "embed_batch",
|
|
"batch": batch_no,
|
|
"total_batches": total_batches,
|
|
"backend": backend,
|
|
"chunks": len(batch),
|
|
"done": done,
|
|
"total": total,
|
|
"message": f"Embedded batch {batch_no}/{total_batches} ({backend})",
|
|
}
|
|
)
|
|
if forced_local:
|
|
collected = local_embed([chunk.text for chunk in chunks]).vectors
|
|
backend = "local"
|
|
store.add(chunks, collected)
|
|
final_backend = "local" if forced_local else backend
|
|
emit({"type": "embed_done", "backend": final_backend, "chunk_count": total})
|
|
return total, final_backend
|
|
|
|
|
|
async def _run(payload: dict, output_dir: Path) -> dict:
|
|
query = payload.get("query", "")
|
|
max_pages = int(payload.get("max_pages", 12))
|
|
depth = int(payload.get("depth", 2))
|
|
api_key = payload.get("api_key", "")
|
|
collection = payload.get("collection", "")
|
|
cached_hashes = set(payload.get("cached_hashes", []))
|
|
should_stop = _make_stop(output_dir)
|
|
|
|
_stage("planning", "Planning research queries", PHASE_PLANNING)
|
|
queries = await plan_queries(query, api_key, _emit)
|
|
_emit({"type": "queries", "queries": queries})
|
|
|
|
_stage("searching", "Searching the web", PHASE_SEARCHING)
|
|
candidates = await search_queries(queries, _emit)
|
|
_emit({"type": "candidates", "count": len(candidates)})
|
|
|
|
_stage("crawling", "Crawling sources", PHASE_CRAWLING)
|
|
outcome = await crawl(
|
|
candidates,
|
|
max_pages,
|
|
_emit,
|
|
lambda url: url_hash(url) in cached_hashes,
|
|
should_stop,
|
|
query=query,
|
|
depth=depth,
|
|
)
|
|
|
|
new_cache = [
|
|
{
|
|
"url_hash": url_hash(page.url),
|
|
"url": page.url,
|
|
"title": page.title,
|
|
"content_hash": content_hash(page.text),
|
|
"status": page.status,
|
|
"byte_size": len(page.text),
|
|
}
|
|
for page in outcome.pages
|
|
]
|
|
|
|
store = VectorStore(collection)
|
|
_stage("indexing", "Indexing content", PHASE_INDEXING)
|
|
chunk_count, embed_backend = await _index_chunks(
|
|
store, outcome.pages, api_key, _emit
|
|
)
|
|
|
|
_stage("analysis", "Running research agents", PHASE_ANALYSIS)
|
|
result = await orchestrate(
|
|
query, outcome.pages, api_key, _emit, store=store, queries=queries
|
|
)
|
|
|
|
_stage("synthesis", "Compiling cited report", PHASE_SYNTHESIS)
|
|
|
|
sources = [
|
|
{"url": page.url, "title": page.title, "source": page.source}
|
|
for page in outcome.pages
|
|
]
|
|
report = {
|
|
"query": query,
|
|
"depth": depth,
|
|
"max_pages": max_pages,
|
|
"generated_at": datetime.now(timezone.utc).isoformat(),
|
|
"summary": result.summary,
|
|
"findings": result.findings,
|
|
"sources": sources,
|
|
"score": result.score,
|
|
"confidence": result.confidence,
|
|
"source_diversity": result.source_diversity,
|
|
"synthesis": result.synthesis,
|
|
"page_count": len(outcome.pages),
|
|
"chunk_count": chunk_count,
|
|
"embed_backend": embed_backend,
|
|
"collection": collection,
|
|
}
|
|
(output_dir / "report.json").write_text(
|
|
json.dumps(report, ensure_ascii=False), encoding="utf-8"
|
|
)
|
|
(output_dir / "url_cache.json").write_text(
|
|
json.dumps(new_cache, ensure_ascii=False), encoding="utf-8"
|
|
)
|
|
_emit(
|
|
{
|
|
"type": "report_ready",
|
|
"score": report["score"],
|
|
"confidence": report["confidence"],
|
|
"source_diversity": report["source_diversity"],
|
|
"synthesis": report["synthesis"],
|
|
"page_count": report["page_count"],
|
|
"chunk_count": report["chunk_count"],
|
|
}
|
|
)
|
|
return report
|
|
|
|
|
|
def main(argv: list) -> int:
|
|
if len(argv) != 3:
|
|
sys.stderr.write("usage: deepsearch.worker <payload_json> <output_dir>\n")
|
|
return 2
|
|
payload = json.loads(Path(argv[1]).read_text(encoding="utf-8"))
|
|
output_dir = Path(argv[2])
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
asyncio.run(_run(payload, output_dir))
|
|
except Exception as exc:
|
|
_emit({"type": "error", "message": str(exc)[:500]})
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main(sys.argv))
|