Spaces:
Runtime error
Runtime error
Download src/pipeline.py from MuhammadAzfar/ResearchMind-AI: direct link, hf CLI and curl.
- Browser
- Download file 11.7 kB
-
https://huggingface.co/spaces/MuhammadAzfar/ResearchMind-AI/resolve/main/src/pipeline.py
- Command line
-
hf download hf://spaces/MuhammadAzfar/ResearchMind-AI/src/pipeline.py
-
curl -L -o pipeline.py https://huggingface.co/spaces/MuhammadAzfar/ResearchMind-AI/resolve/main/src/pipeline.py
11.7 kB
| """Orchestrates the agents: Planner -> Researcher -> Retriever -> Writer -> Checker -> Reviser -> Translator.""" | |
| import re | |
| import time | |
| import types | |
| from src import agents, config, llm | |
| from src.documents import load_documents | |
| from src.evidence import EvidenceError, EvidenceStore | |
| from src.search import search_topic | |
| DEFAULT_LLM = types.SimpleNamespace(complete=llm.complete, stream=llm.stream_completion) | |
| _ICONS = {"run": "β³", "ok": "β ", "warn": "β οΈ", "fail": "β"} | |
| class Trace: | |
| """The live 'agent activity' log shown in the UI.""" | |
| def __init__(self): | |
| self.steps = [] | |
| def start(self, name: str, detail: str = "") -> int: | |
| self.steps.append({"name": name, "status": "run", "detail": detail, "items": []}) | |
| return len(self.steps) - 1 | |
| def finish(self, idx: int, detail: str = None, status: str = "ok", items: list = None): | |
| step = self.steps[idx] | |
| step["status"] = status | |
| if detail is not None: | |
| step["detail"] = detail | |
| if items: | |
| step["items"] = items | |
| def render(self) -> str: | |
| lines = ["### π€ Agent activity", ""] | |
| for s in self.steps: | |
| lines.append(f"{_ICONS[s['status']]} **{s['name']}** β {s['detail']}".rstrip(" β")) | |
| for item in s["items"]: | |
| lines.append(f" - {item}") | |
| return "\n".join(lines) | |
| def _event(trace, report=None, done=False, error=False, result=None): | |
| return {"trace": trace.render(), "report": report, "done": done, "error": error, "result": result} | |
| def sources_markdown(sources: list) -> str: | |
| lines = [] | |
| for s in sources: | |
| if s.get("kind") == "document": | |
| lines.append(f"{s['n']}. π **{s['title']}** (uploaded document)") | |
| else: | |
| title = re.sub(r"[\[\]]", "", s["title"]) or s["url"] | |
| lines.append(f"{s['n']}. [{title}]({s['url']})") | |
| return "\n\n---\n\n## Sources\n\n" + "\n".join(lines) | |
| def verification_markdown(final: dict, first: dict, revised: bool) -> str: | |
| if not final or not final["checked"]: | |
| return "\n\n> π **Verification:** the automatic fact-check was unavailable for this report. Please verify key facts yourself." | |
| text = f"{final['supported']} of {final['checked']} cited claims are fully supported by their sources" | |
| extras = [] | |
| if final["partial"]: | |
| extras.append(f"{final['partial']} partially supported") | |
| if final["unsupported"]: | |
| extras.append(f"{final['unsupported']} unsupported") | |
| if extras: | |
| text += " (" + ", ".join(extras) + ")" | |
| if revised and first and first["checked"]: | |
| text += f". The first draft scored {first['supported']}/{first['checked']}; the Reviser agent corrected the rest" | |
| return f"\n\n> π **Verification:** {text}. This is an automated check β verify critical facts yourself." | |
| def _check(report, store, complete): | |
| claims, uncited = agents.extract_claims(report) | |
| results = agents.check_claims(claims, store, complete) if claims else {} | |
| return claims, results, agents.summarize_check(claims, results, uncited) | |
| def run_research( | |
| topic: str, | |
| mode: str = "Standard", | |
| language: str = "English", | |
| depth: str = config.DEFAULT_DEPTH, | |
| use_web: bool = True, | |
| document_paths=None, | |
| *, | |
| searcher=None, | |
| embedder=None, | |
| llm_api=None, | |
| ): | |
| """Generator. Yields {"trace", "report", "done", "error", "result"} events as the agents work. | |
| The final event has done=True and a `result` dict with the report, sources and fact-check metrics. | |
| `searcher`, `embedder` and `llm_api` can be replaced for testing. | |
| """ | |
| started = time.time() | |
| searcher = searcher or search_topic | |
| llm_api = llm_api or DEFAULT_LLM | |
| cfg = config.DEPTH_PRESETS.get(depth, config.DEPTH_PRESETS[config.DEFAULT_DEPTH]) | |
| trace = Trace() | |
| topic = (topic or "").strip() | |
| if len(topic) < 2: | |
| yield _event(trace, report="### β οΈ Please enter a research topic.", error=True) | |
| return | |
| # ---- 0. uploaded documents --------------------------------------------------------- | |
| documents = [] | |
| if document_paths: | |
| step = trace.start("Documents", "reading uploaded files...") | |
| yield _event(trace) | |
| documents, warnings = load_documents(document_paths) | |
| trace.finish( | |
| step, | |
| f"{len(documents)} file(s) loaded" if documents else "no readable documents", | |
| "ok" if documents else "warn", | |
| items=warnings, | |
| ) | |
| if not use_web and not documents: | |
| yield _event(trace, report="### β οΈ Web search is off and no readable documents were uploaded, so there is nothing to research.", error=True) | |
| return | |
| # ---- 1. Planner -------------------------------------------------------------------- | |
| step = trace.start("Planner", "splitting the topic into research questions...") | |
| yield _event(trace) | |
| try: | |
| questions = agents.plan(topic, cfg["sub_questions"], llm_api.complete) | |
| trace.finish(step, f"{len(questions)} research question(s)", items=[f"{i}. {q}" for i, q in enumerate(questions, 1)]) | |
| except llm.LLMError as e: | |
| questions = [topic] | |
| trace.finish(step, f"planner unavailable ({e}); using the topic as a single question", "warn") | |
| yield _event(trace) | |
| # ---- 2. Researcher ----------------------------------------------------------------- | |
| sources, errors = list(documents), [] | |
| if use_web: | |
| step = trace.start("Researcher", f"searching the web ({len(questions)} parallel searches)...") | |
| yield _event(trace) | |
| web_sources, errors = agents.research_web(questions, cfg, searcher) | |
| sources = web_sources + sources | |
| if web_sources: | |
| trace.finish(step, f"{len(web_sources)} web source(s) found", items=[f"note: {e}" for e in errors]) | |
| else: | |
| trace.finish(step, "no web results", "warn", items=errors) | |
| yield _event(trace) | |
| if not sources: | |
| message = "### β οΈ " + (errors[0] if use_web and errors else "No sources were found for this topic.") | |
| yield _event(trace, report=message, error=True) | |
| return | |
| # ---- 3. Index + retrieve ----------------------------------------------------------- | |
| step = trace.start("Retriever", "chunking, embedding and indexing sources...") | |
| yield _event(trace) | |
| try: | |
| store = EvidenceStore(sources, embedder) | |
| selected, numbered = store.gather(questions, cfg["chunks_per_question"]) | |
| except EvidenceError as e: | |
| trace.finish(step, str(e), "fail") | |
| yield _event(trace, report=f"### β οΈ {e}", error=True) | |
| return | |
| except Exception as e: | |
| print(f"[DEBUG ERROR] Retrieval failed: {e}") | |
| trace.finish(step, f"retrieval failed: {e}", "fail") | |
| yield _event(trace, report=f"### β Retrieval failed: `{e}`", error=True) | |
| return | |
| trace.finish(step, f"{len(store.records)} passages indexed, {len(selected)} selected from {len(numbered)} source(s)") | |
| yield _event(trace) | |
| # ---- 4. Writer --------------------------------------------------------------------- | |
| step = trace.start("Writer", "drafting the cited report...") | |
| draft = "" | |
| try: | |
| for delta in agents.write_stream(topic, mode, questions, selected, numbered, llm_api.stream): | |
| draft += delta | |
| shown = agents.clean_citations(draft, len(numbered)) | |
| yield _event(trace, report="*Draft β fact-checking will follow...*\n\n" + shown) | |
| except llm.LLMError as e: | |
| trace.finish(step, str(e), "fail") | |
| yield _event(trace, report=f"### β {e}", error=True) | |
| return | |
| report = agents.clean_citations(draft, len(numbered)).strip() | |
| if not report: | |
| trace.finish(step, "empty answer", "fail") | |
| yield _event(trace, report="### β οΈ The model returned an empty answer. Please try again.", error=True) | |
| return | |
| trace.finish(step, f"draft written ({len(report.split())} words)") | |
| yield _event(trace, report="*Draft β fact-checking in progress...*\n\n" + report) | |
| # ---- 5. Checker (+ Reviser loop) --------------------------------------------------- | |
| first_metrics = final_metrics = None | |
| revised_flag = False | |
| step = trace.start("Checker", "verifying every cited claim against its sources...") | |
| yield _event(trace, report="*Draft β fact-checking in progress...*\n\n" + report) | |
| try: | |
| claims, results, first_metrics = _check(report, store, llm_api.complete) | |
| final_metrics = first_metrics | |
| if first_metrics["checked"]: | |
| m = first_metrics | |
| trace.finish(step, f"{m['supported']}/{m['checked']} claims supported, {m['partial']} partial, {m['unsupported']} unsupported") | |
| else: | |
| trace.finish(step, "no claims could be verified", "warn") | |
| yield _event(trace, report="*Draft β fact-checking in progress...*\n\n" + report) | |
| for _ in range(config.MAX_REVISION_ROUNDS): | |
| flagged = final_metrics["partial"] + final_metrics["unsupported"] | |
| if not flagged: | |
| break | |
| step = trace.start("Reviser", f"rewriting {flagged} flagged claim(s)...") | |
| yield _event(trace, report="*Draft β revising flagged claims...*\n\n" + report) | |
| candidate = agents.revise(report, claims, results, store, llm_api.complete) | |
| if candidate == report: | |
| trace.finish(step, "no usable revision", "warn") | |
| break | |
| claims2, results2, metrics2 = _check(candidate, store, llm_api.complete) | |
| old_rate = final_metrics["support_rate"] or 0 | |
| new_rate = metrics2["support_rate"] or 0 | |
| if metrics2["checked"] and new_rate >= old_rate: | |
| report, claims, results, final_metrics, revised_flag = candidate, claims2, results2, metrics2, True | |
| trace.finish(step, f"revised; now {metrics2['supported']}/{metrics2['checked']} claims supported") | |
| else: | |
| trace.finish(step, "revision did not improve the report; kept the original draft", "warn") | |
| yield _event(trace, report=report) | |
| except llm.LLMError as e: | |
| trace.finish(step, f"unavailable ({e}); report delivered unverified", "warn") | |
| yield _event(trace, report=report) | |
| # ---- 6. Translator ----------------------------------------------------------------- | |
| english_report = report | |
| if language != "English": | |
| step = trace.start("Translator", f"translating the verified report to {language}...") | |
| yield _event(trace, report=report) | |
| try: | |
| report = agents.translate(report, language, llm_api.complete) | |
| trace.finish(step, f"translated to {language}") | |
| except llm.LLMError as e: | |
| trace.finish(step, f"translation failed ({e}); showing English", "warn") | |
| # ---- 7. Assemble ------------------------------------------------------------------- | |
| final_md = report + sources_markdown(numbered) + verification_markdown(final_metrics, first_metrics, revised_flag) | |
| elapsed = round(time.time() - started, 1) | |
| result = { | |
| "topic": topic, | |
| "report": final_md, | |
| "report_english": english_report, | |
| "language": language, | |
| "questions": questions, | |
| "sources": numbered, | |
| "metrics": { | |
| "first_pass": first_metrics, | |
| "final": final_metrics, | |
| "revised": revised_flag, | |
| "sources": len(numbered), | |
| "passages_indexed": len(store.records), | |
| "seconds": elapsed, | |
| }, | |
| } | |
| yield _event(trace, report=final_md, done=True, result=result) | |