devplacepy/worker.py at 0a554ebc32ed500531f145ce19df92797442ade6

 # 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
 CONTROL_FILE = "control.json"
 PAUSE_POLL_SECONDS = 1.0
 EMBED_BATCH = 64
 def _emit(frame: dict) -> None:
     sys.stdout.write(json.dumps(frame, ensure_ascii=False) + "\n")
     sys.stdout.flush()
 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
 ) -> 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,
                 )
             )
     if not chunks:
         return 0, "empty"
     backend = "gateway"
     for start in range(0, len(chunks), EMBED_BATCH):
         batch = chunks[start : start + EMBED_BATCH]
         result = await embed_texts([chunk.text for chunk in batch], api_key)
         if not result.vectors:
             result = local_embed([chunk.text for chunk in batch])
         backend = result.backend
         store.add(batch, result.vectors)
     return len(chunks), 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)
     _emit({"type": "stage", "stage": "planning", "message": "Planning research queries"})
     queries = await plan_queries(query, api_key)
     _emit({"type": "queries", "queries": queries})
     _emit({"type": "stage", "stage": "searching", "message": "Searching the web"})
     candidates = await search_queries(queries, _emit)
     _emit({"type": "candidates", "count": len(candidates)})
     _emit({"type": "stage", "stage": "crawling", "message": "Crawling sources"})
     outcome = await crawl(
         candidates,
         max_pages,
         _emit,
         lambda url: url_hash(url) in cached_hashes,
         should_stop,
     )
     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)
     _emit({"type": "stage", "stage": "indexing", "message": "Indexing content"})
     chunk_count, embed_backend = await _index_chunks(store, outcome.pages, api_key)
     _emit({"type": "stage", "stage": "analysis", "message": "Running research agents"})
     result = await orchestrate(query, outcome.pages, api_key, _emit)
     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,
         "gaps": result.gaps,
         "sources": sources,
         "score": result.score,
         "confidence": result.confidence,
         "source_diversity": result.source_diversity,
         "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"],
             "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))