File size: 1,254 Bytes
61411b5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
from __future__ import annotations

import logging
import os
from typing import Any, Dict, List, Optional

from tavily import TavilyClient

logger = logging.getLogger(__name__)


class TavilyWebSearchTool:
    def __init__(self, api_key: Optional[str] = None) -> None:
        api_key = api_key or os.getenv("TAVILY_API_KEY", "")
        if not api_key:
            raise ValueError("Missing TAVILY_API_KEY.")
        self._client = TavilyClient(api_key=api_key)

    def search(self, query: str, *, max_results: int = 5) -> Dict[str, Any]:
        logger.info("Tavily search: %s", query)
        res = self._client.search(query=query, max_results=max_results)
        return res

    @staticmethod
    def summarize(search_result: Dict[str, Any]) -> str:
        results: List[Dict[str, Any]] = search_result.get("results", []) or []
        lines = []
        for r in results[:8]:
            title = r.get("title") or ""
            url = r.get("url") or ""
            content = (r.get("content") or "").strip()
            if len(content) > 400:
                content = content[:400] + "..."
            lines.append(f"- {title} ({url})\n  {content}")
        return "\n".join(lines).strip() or "No results."