Spaces:
Sleeping
Sleeping
File size: 3,654 Bytes
30f6ce2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """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"
@lru_cache(maxsize=1)
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
|