agent_search_gateway.result_writer

  1"""Compact JSONL persistence for search command results."""
  2
  3import json
  4from collections.abc import Iterable, Mapping
  5from contextlib import suppress
  6from datetime import date
  7from pathlib import Path
  8from typing import Literal
  9
 10from .errors import InputFailure
 11from .models import PaperRecord, SearchRecord
 12from .request_ids import ResultKind, result_filename
 13from .url_normalization import normalize_url
 14
 15PaperResultKind = Literal["paper", "llm"]
 16
 17
 18def _web_payload(record: SearchRecord) -> dict[str, object]:
 19    abstract = record.abstract.strip()
 20    if not abstract:
 21        raise ValueError("search result abstract must be non-empty")
 22    return {"url": str(record.url), "abstract": abstract}
 23
 24
 25def _serialize_record(record: SearchRecord) -> str:
 26    return json.dumps(
 27        _web_payload(record),
 28        ensure_ascii=False,
 29        separators=(",", ":"),
 30    )
 31
 32
 33def _normalized_url_or_error(value: object, label: str) -> str:
 34    if not isinstance(value, str):
 35        raise ValueError(f"paper result {label} must be a normalized HTTP(S) URL")
 36    try:
 37        normalized = normalize_url(value)
 38    except InputFailure as exc:
 39        raise ValueError(f"paper result {label} must be a normalized HTTP(S) URL") from exc
 40    if normalized != value:
 41        raise ValueError(f"paper result {label} must already be normalized")
 42    return str(normalized)
 43
 44
 45def _optional_normalized_url_or_error(value: object, label: str) -> str | None:
 46    if value is None:
 47        return None
 48    return _normalized_url_or_error(value, label)
 49
 50
 51def _string_tuple(value: object, label: str, *, unique: bool = False) -> list[str]:
 52    if not isinstance(value, tuple) or any(not isinstance(item, str) for item in value):
 53        raise ValueError(f"paper result {label} must be a tuple of strings")
 54    items = list(value)
 55    if unique and len(items) != len(set(items)):
 56        raise ValueError(f"paper result {label} must not contain duplicates")
 57    return items
 58
 59
 60def _date_value(value: object, label: str) -> str | None:
 61    if value is None:
 62        return None
 63    if not isinstance(value, date):
 64        raise ValueError(f"paper result {label} must be a date or null")
 65    return value.isoformat()
 66
 67
 68def _citation_counts(value: object) -> dict[str, int]:
 69    if not isinstance(value, Mapping):
 70        raise ValueError("paper result citation_counts must be a mapping")
 71    result: dict[str, int] = {}
 72    for source, count in value.items():
 73        if not isinstance(source, str) or not source:
 74            raise ValueError("paper result citation source must be a non-empty string")
 75        if isinstance(count, bool) or not isinstance(count, int) or count < 0:
 76            raise ValueError("paper result citation count must be a non-negative integer")
 77        result[source] = count
 78    return result
 79
 80
 81def _paper_payload(record: PaperRecord) -> dict[str, object]:
 82    title = record.title.strip() if isinstance(record.title, str) else ""
 83    if not title:
 84        raise ValueError("paper result title must be non-empty")
 85    abstract = record.abstract.strip() if isinstance(record.abstract, str) else ""
 86    venue = record.venue.strip() if isinstance(record.venue, str) else ""
 87    oa_status = record.oa_status.strip() if isinstance(record.oa_status, str) else ""
 88    license_value = record.license.strip() if isinstance(record.license, str) else ""
 89    if record.is_open_access is not None and not isinstance(record.is_open_access, bool):
 90        raise ValueError("paper result is_open_access must be boolean or null")
 91
 92    identifiers = record.identifiers
 93    identifier_payload = {
 94        "doi": identifiers.doi,
 95        "arxiv_id": identifiers.arxiv_id,
 96        "semantic_scholar_id": identifiers.semantic_scholar_id,
 97        "openalex_id": identifiers.openalex_id,
 98        "dblp_key": identifiers.dblp_key,
 99        "core_id": identifiers.core_id,
100    }
101    if any(not isinstance(value, str) for value in identifier_payload.values()):
102        raise ValueError("paper result identifiers must contain strings")
103
104    return {
105        "title": title,
106        "authors": _string_tuple(record.authors, "authors"),
107        "abstract": abstract,
108        "identifiers": identifier_payload,
109        "published_date": _date_value(record.published_date, "published_date"),
110        "updated_date": _date_value(record.updated_date, "updated_date"),
111        "url": _normalized_url_or_error(record.url, "url"),
112        "pdf_url": _optional_normalized_url_or_error(record.pdf_url, "pdf_url"),
113        "venue": venue,
114        "topics": _string_tuple(record.topics, "topics"),
115        "citation_counts": _citation_counts(record.citation_counts),
116        "is_open_access": record.is_open_access,
117        "oa_status": oa_status,
118        "license": license_value,
119        "sources": _string_tuple(record.sources, "sources", unique=True),
120    }
121
122
123def _serialize_paper_record(record: PaperRecord) -> str:
124    return json.dumps(
125        _paper_payload(record),
126        ensure_ascii=False,
127        separators=(",", ":"),
128    )
129
130
131def _serialize_mixed_web(record: SearchRecord) -> str:
132    return json.dumps(
133        {"type": "web", **_web_payload(record)},
134        ensure_ascii=False,
135        separators=(",", ":"),
136    )
137
138
139def _serialize_mixed_paper(record: PaperRecord) -> str:
140    return json.dumps(
141        {"type": "paper", **_paper_payload(record)},
142        ensure_ascii=False,
143        separators=(",", ":"),
144    )
145
146
147class ResultWriter:
148    def __init__(self, results_dir: Path) -> None:
149        self._results_dir = results_dir
150
151    def write_results(
152        self,
153        kind: ResultKind,
154        records: Iterable[SearchRecord],
155        *,
156        request_id: str,
157    ) -> Path:
158        filename = result_filename(kind, request_id)
159        serialized = tuple(_serialize_record(record) for record in records)
160        return self._write_serialized(filename, serialized)
161
162    def write_paper_results(
163        self,
164        kind: PaperResultKind,
165        records: Iterable[PaperRecord],
166        *,
167        request_id: str,
168    ) -> Path:
169        if kind not in {"paper", "llm"}:
170            raise ValueError(f"invalid paper result kind: {kind}")
171        filename = result_filename(kind, request_id)
172        serialized = tuple(_serialize_paper_record(record) for record in records)
173        return self._write_serialized(filename, serialized)
174
175    def write_mixed_results(
176        self,
177        web_records: Iterable[SearchRecord],
178        paper_records: Iterable[PaperRecord],
179        *,
180        request_id: str,
181    ) -> Path:
182        filename = result_filename("llm", request_id)
183        serialized_web = tuple(_serialize_mixed_web(record) for record in web_records)
184        serialized_paper = tuple(_serialize_mixed_paper(record) for record in paper_records)
185        return self._write_serialized(filename, (*serialized_web, *serialized_paper))
186
187    def _write_serialized(self, filename: str, serialized: Iterable[str]) -> Path:
188        self._results_dir.mkdir(parents=True, exist_ok=True)
189        target = self._results_dir / filename
190        created = False
191        try:
192            with target.open("x", encoding="utf-8", newline="\n") as handle:
193                created = True
194                for line in serialized:
195                    handle.write(line)
196                    handle.write("\n")
197            return target.resolve()
198        except BaseException:
199            if created:
200                with suppress(OSError):
201                    target.unlink()
202            raise