research/types.py

"""Typed research results shared by exploration tools and rewards."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any


@dataclass
class ResearchChunk:
    """A ranked passage returned by a research tool."""

    source: str
    tool: str
    title: str
    url: str
    text: str
    score: float = 0.0
    rank: int = 0
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def snippet(self) -> str:
        return self.text


@dataclass
class ResearchResult:
    """Structured response from a research tool."""

    tool: str
    query: str
    chunks: list[ResearchChunk] = field(default_factory=list)
    error: str = ""
    raw_count: int = 0

    @property
    def ok(self) -> bool:
        return not self.error and bool(self.chunks)

    @property
    def text(self) -> str:
        return "\n\n".join(chunk.text for chunk in self.chunks)

    @property
    def sources(self) -> set[str]:
        return {chunk.source for chunk in self.chunks}

    def render(self) -> str:
        """Render compact context for the agent."""
        if self.error:
            return f"{self.tool} error: {self.error}"
        if not self.chunks:
            return f"No results for {self.tool}: {self.query}"

        parts = []
        for chunk in self.chunks:
            url = f"\nURL: {chunk.url}" if chunk.url else ""
            parts.append(
                f"[{chunk.rank}] {chunk.source}: {chunk.title}{url}\n{chunk.text}"
            )
        return "\n\n---\n\n".join(parts)