agent_search_gateway.url_store

 1"""In-memory five-field URL state machine."""
 2
 3from dataclasses import replace
 4
 5from .models import URLRecord
 6from .url_normalization import NormalizedURL
 7
 8
 9def _first_non_empty(existing: str, candidate: str) -> str:
10    return existing if existing else candidate
11
12
13class URLStore:
14    def __init__(self) -> None:
15        self._records: dict[NormalizedURL, URLRecord] = {}
16
17    def admit(
18        self,
19        url: NormalizedURL,
20        abstract: str,
21        *,
22        raw_content: str = "",
23        content: str = "",
24    ) -> URLRecord:
25        normalized_abstract = abstract.strip()
26        if not normalized_abstract:
27            raise ValueError("URL records require a non-empty abstract")
28
29        existing = self._records.get(url)
30        if existing is None:
31            record = URLRecord(
32                url=url,
33                raw_content=raw_content,
34                content=content,
35                abstract=normalized_abstract,
36            )
37        else:
38            record = replace(
39                existing,
40                raw_content=_first_non_empty(existing.raw_content, raw_content),
41                content=_first_non_empty(existing.content, content),
42                abstract=_first_non_empty(existing.abstract, normalized_abstract),
43            )
44        self._records[url] = record
45        return record
46
47    def merge_body(
48        self,
49        url: NormalizedURL,
50        *,
51        raw_content: str = "",
52        content: str = "",
53    ) -> URLRecord:
54        existing = self._require(url)
55        record = replace(
56            existing,
57            raw_content=_first_non_empty(existing.raw_content, raw_content),
58            content=_first_non_empty(existing.content, content),
59        )
60        self._records[url] = record
61        return record
62
63    def mark_unavailable(self, url: NormalizedURL) -> URLRecord:
64        existing = self._require(url)
65        if not existing.available:
66            return existing
67        record = replace(existing, available=False)
68        self._records[url] = record
69        return record
70
71    def get(self, url: NormalizedURL) -> URLRecord | None:
72        return self._records.get(url)
73
74    def _require(self, url: NormalizedURL) -> URLRecord:
75        record = self._records.get(url)
76        if record is None:
77            raise KeyError(str(url))
78        return record