Spaces:
Sleeping
Sleeping
| """Client for the WaterScopeAI Hugging Face Space /agent_run endpoint.""" | |
| from __future__ import annotations | |
| import concurrent.futures | |
| import os | |
| import re | |
| from functools import lru_cache | |
| from typing import Any | |
| from gradio_client import Client | |
| SPACE_ID = os.getenv("WATERSCOPE_SPACE_ID", "EnvironmentalAI/WaterScopeAI") | |
| TIMEOUT_SEC = int(os.getenv("WATERSCOPE_TIMEOUT_SEC", "300")) | |
| _SPACE_EXECUTOR = concurrent.futures.ThreadPoolExecutor( | |
| max_workers=2, | |
| thread_name_prefix="waterscope-space", | |
| ) | |
| MODELS = ("base", "dpo", "da_dpo", "tulucore") | |
| APPLICATIONS = ( | |
| "general_qa", | |
| "document_qa", | |
| "paper_synthesis", | |
| "method_recommendation", | |
| "research_gaps", | |
| "claim_check", | |
| ) | |
| _SYNTHESIS = re.compile(r"\b(summarize|summary|overview|synthesize|main finding)\b", re.I) | |
| _METHODS = re.compile(r"\b(method|approach|model|algorithm|recommend)\b", re.I) | |
| _GAPS = re.compile(r"\b(gap|limitation|future work|open question|unclear)\b", re.I) | |
| _CLAIM = re.compile(r"\b(does the paper|is it true|claim|support|evidence that)\b", re.I) | |
| def infer_application(question: str, *, explicit: str | None = None) -> str: | |
| if explicit and explicit != "auto": | |
| if explicit not in APPLICATIONS: | |
| raise ValueError(f"Unknown application {explicit!r}") | |
| return explicit | |
| text = question.lower() | |
| if _SYNTHESIS.search(text): | |
| return "paper_synthesis" | |
| if _GAPS.search(text): | |
| return "research_gaps" | |
| if _CLAIM.search(text): | |
| return "claim_check" | |
| if _METHODS.search(text): | |
| return "method_recommendation" | |
| return "document_qa" | |
| def _client() -> Client: | |
| token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN") | |
| kwargs: dict[str, Any] = {"httpx_kwargs": {"timeout": TIMEOUT_SEC}} | |
| if token: | |
| kwargs["hf_token"] = token | |
| return Client(SPACE_ID, **kwargs) | |
| def run_agent( | |
| *, | |
| context: str, | |
| question: str, | |
| model: str | None = None, | |
| application: str = "document_qa", | |
| max_new_tokens: int | None = None, | |
| ) -> tuple[str, str]: | |
| """Call /agent_run. Returns (answer, diagnostics).""" | |
| selected_model = model or os.getenv("WATERSCOPE_MODEL", "dpo") | |
| if selected_model not in MODELS: | |
| raise ValueError(f"Unknown model {selected_model!r}. Choose from: {MODELS}") | |
| app_type = infer_application(question, explicit=application if application != "auto" else None) | |
| tokens = max_new_tokens or int(os.getenv("WATERSCOPE_MAX_TOKENS", "1024")) | |
| def _predict() -> tuple[str, str]: | |
| response = _client().predict( | |
| selected_model, | |
| app_type, | |
| context, | |
| question, | |
| tokens, | |
| api_name="/agent_run", | |
| ) | |
| answer = response[0] if response else "" | |
| diagnostics = response[1] if len(response) > 1 else "" | |
| return answer, diagnostics | |
| # Run in a worker thread so Gradio asyncio does not clash with MCP's loop. | |
| answer, diagnostics = _SPACE_EXECUTOR.submit(_predict).result(timeout=TIMEOUT_SEC) | |
| return answer, diagnostics | |
| def agent_answer( | |
| context: str, | |
| question: str, | |
| *, | |
| application: str = "auto", | |
| model: str | None = None, | |
| max_new_tokens: int | None = None, | |
| ) -> tuple[str, str, str]: | |
| """Returns (answer, application_used, diagnostics).""" | |
| app_type = infer_application(question, explicit=application if application != "auto" else None) | |
| answer, diagnostics = run_agent( | |
| context=context, | |
| question=question, | |
| model=model, | |
| application=app_type, | |
| max_new_tokens=max_new_tokens, | |
| ) | |
| return answer, app_type, diagnostics | |