130 lines
5.2 KiB
Python
130 lines
5.2 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Awaitable, Callable
|
|
|
|
from ..schemas import ExternalAdapterResult, ExternalEvidenceItem, ExternalQueryPlan
|
|
|
|
|
|
class MCPExternalAdapter:
|
|
adapter_name = "mcp"
|
|
supported_query_kinds = {"model", "company", "generic_web", "unknown"}
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
search_callable: Callable[[ExternalQueryPlan], Awaitable[dict[str, Any] | ExternalAdapterResult]] | None = None,
|
|
fetch_text_callable: Callable[[str], Awaitable[str]] | None = None,
|
|
chunk_text_callable: Callable[..., list[str]] | None = None,
|
|
score_chunk_callable: Callable[..., float] | None = None,
|
|
max_pages: int = 3,
|
|
) -> None:
|
|
self.search_callable = search_callable
|
|
self.fetch_text_callable = fetch_text_callable
|
|
self.chunk_text_callable = chunk_text_callable
|
|
self.score_chunk_callable = score_chunk_callable
|
|
self.max_pages = max(1, int(max_pages or 3))
|
|
|
|
async def execute(self, plan: ExternalQueryPlan) -> ExternalAdapterResult:
|
|
if not callable(self.search_callable):
|
|
return ExternalAdapterResult(
|
|
adapter_name=self.adapter_name,
|
|
status="unavailable",
|
|
evidence=[],
|
|
meta={"reason": "mcp_adapter_unavailable"},
|
|
)
|
|
result = await self.search_callable(plan)
|
|
if isinstance(result, ExternalAdapterResult):
|
|
validated = result
|
|
else:
|
|
payload = dict(result or {}) if isinstance(result, dict) else {}
|
|
payload.setdefault("adapter_name", self.adapter_name)
|
|
validated = ExternalAdapterResult.model_validate(payload)
|
|
evidence, fetched_pages = await self._hydrate_evidence_from_target_pages(
|
|
plan=plan,
|
|
evidence=list(validated.evidence or []),
|
|
)
|
|
meta = dict(validated.meta or {})
|
|
if fetched_pages:
|
|
meta["fetched_pages"] = int(fetched_pages)
|
|
return validated.model_copy(update={"evidence": evidence, "meta": meta})
|
|
|
|
def _page_fetch_ready(self) -> bool:
|
|
return all(
|
|
callable(item)
|
|
for item in (
|
|
self.fetch_text_callable,
|
|
self.chunk_text_callable,
|
|
self.score_chunk_callable,
|
|
)
|
|
)
|
|
|
|
async def _hydrate_evidence_from_target_pages(
|
|
self,
|
|
*,
|
|
plan: ExternalQueryPlan,
|
|
evidence: list[ExternalEvidenceItem],
|
|
) -> tuple[list[ExternalEvidenceItem], int]:
|
|
if not evidence or not self._page_fetch_ready():
|
|
return evidence, 0
|
|
hydrated: list[ExternalEvidenceItem] = []
|
|
fetch_attempts = 0
|
|
query = str((plan.search_queries[:1] or [plan.topic])[0] or "").strip() or str(plan.topic or "").strip()
|
|
for index, item in enumerate(evidence):
|
|
target_url = str(item.target_url or "").strip()
|
|
if index >= self.max_pages or (
|
|
not target_url.startswith("http://") and not target_url.startswith("https://")
|
|
):
|
|
hydrated.append(item)
|
|
continue
|
|
fetch_attempts += 1
|
|
try:
|
|
text = await self.fetch_text_callable(target_url)
|
|
except Exception:
|
|
hydrated.append(item)
|
|
continue
|
|
chunks = list(self.chunk_text_callable(text, chunk_size=900, overlap=120) or [])
|
|
if not chunks:
|
|
hydrated.append(item)
|
|
continue
|
|
|
|
best_chunk = ""
|
|
best_score = 0.0
|
|
title = str(item.title or "").strip()
|
|
for chunk in chunks[:6]:
|
|
score = float(
|
|
self.score_chunk_callable(
|
|
question=query or plan.topic or "",
|
|
title=title,
|
|
content=chunk,
|
|
url=target_url,
|
|
)
|
|
or 0.0
|
|
)
|
|
if score > best_score:
|
|
best_score = score
|
|
best_chunk = str(chunk or "").strip()
|
|
if not best_chunk:
|
|
hydrated.append(item)
|
|
continue
|
|
|
|
grounded_fields = [str(field or "").strip() for field in list(item.grounded_fields or []) if str(field or "").strip()]
|
|
if "page_content_snippet" not in grounded_fields:
|
|
grounded_fields.append("page_content_snippet")
|
|
raw_metadata = dict(item.raw_metadata or {})
|
|
original_snippet = str(item.snippet or "").strip()
|
|
if original_snippet and not str(raw_metadata.get("search_snippet") or "").strip():
|
|
raw_metadata["search_snippet"] = original_snippet[:1800]
|
|
raw_metadata["page_chunk_score"] = round(best_score, 4)
|
|
raw_metadata["target_page_fetched"] = True
|
|
hydrated.append(
|
|
item.model_copy(
|
|
update={
|
|
"snippet": best_chunk[:1800],
|
|
"freshness_score": round(max(float(item.freshness_score or 0.0), best_score), 4),
|
|
"grounded_fields": grounded_fields,
|
|
"raw_metadata": raw_metadata,
|
|
}
|
|
)
|
|
)
|
|
return hydrated, fetch_attempts
|