Source code for aidputils.agents.tools.experimental.websearch.tavily

from __future__ import annotations
from typing import Any

from aidputils.agents.tools.experimental.websearch.base import (
    BaseWebSearchHandler,
    _as_str,
)

from aidputils.agents.tools.experimental.websearch.web_search_models import NormalizedWebResult



class TavilyWebSearchHandler(BaseWebSearchHandler):
    """
    Tavily search parser (engine-native JSON).

    Typical Tavily response shape:
      {
        "answer": "...",                # optional
        "query": "...",                 # optional
        "follow_up_questions": [...],   # optional
        "results": [
          {
            "title": "...",
            "url": "...",
            "content": "...",
            "score": 0.73,
            "raw_content": "..."        # optional (if include_raw_content)
          },
          ...
        ]
      }

    We normalize Tavily items into NormalizedWebResult where:
      - snippet <- content
      - source <- "TAVILY"
    """

    def __init__(self):
        super().__init__("TAVILY")



    def get_endpoint(self, conf: dict) -> str:
        return "https://api.tavily.com/search"




    def parse(self, data: Any) -> list[NormalizedWebResult]:
        if not isinstance(data, dict):
            return []

        results = data.get("results")
        if not isinstance(results, list):
            return self._fallback(data)

        out: list[NormalizedWebResult] = []
        for r in results:
            if not isinstance(r, dict):
                continue
            out.append(
                NormalizedWebResult(
                    title=_as_str(r.get("title")),
                    url=_as_str(r.get("url")),
                    snippet=_as_str(r.get("content") or r.get("snippet") or r.get("description")),
                    source=_as_str(r.get("source") or self.engine),
                    published_time=_as_str(r.get("published_time") or r.get("publishedTime")),
                )
            )
        return out