diff --git a/docs/extras/integrations/retrievers/tavily.ipynb b/docs/extras/integrations/retrievers/tavily.ipynb new file mode 100644 index 00000000000..9c87e35ed73 --- /dev/null +++ b/docs/extras/integrations/retrievers/tavily.ipynb @@ -0,0 +1,79 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Tavily Search API\n", + "\n", + "[Tavily's Search API](https://tavily.com) is a search engine built specifically for AI agents (LLMs), delivering real-time, accurate, and factual results at speed.\n", + "\n", + "## Usage\n", + "\n", + "For a full list of allowed arguments, see [the official documentation](https://app.tavily.com/documentation/python). You can also pass any param to the SDK via a `kwargs` dictionary." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# %pip install tavily-python" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[Document(page_content='Nintendo Designer (s) Hidemaro Fujibayashi (director) Eiji Aonuma (producer/group manager) Release date (s) United States of America: • March 3, 2017 Japan: • March 3, 2017 Australia / New Zealand: • March 2, 2017 Belgium: • March 3, 2017 Hong Kong: • Feburary 1, 2018 South Korea: • February 1, 2018 The UK / Ireland: • March 3, 2017 Content ratings', metadata={'title': 'The Legend of Zelda: Breath of the Wild - Zelda Wiki', 'source': 'https://zelda.fandom.com/wiki/The_Legend_of_Zelda:_Breath_of_the_Wild', 'score': 0.96994, 'images': None}),\n", + " Document(page_content='02/01/23 Nintendo Switch Online member exclusive: Save on two digital games Read more 09/13/22 Out of the Shadows … the Legend of Zelda: Tears of the Kingdom Launches for Nintendo Switch on May...', metadata={'title': 'The Legend of Zelda™: Breath of the Wild - Nintendo', 'source': 'https://www.nintendo.com/store/products/the-legend-of-zelda-breath-of-the-wild-switch/', 'score': 0.94346, 'images': None}),\n", + " Document(page_content='Now we finally have a concrete release date of May 12, 2023. The date was announced alongside this brief (and mysterious) new trailer that also confirmed its title: The Legend of Zelda: Tears...', metadata={'title': 'The Legend of Zelda: Tears of the Kingdom: Release Date, Gameplay ... - IGN', 'source': 'https://www.ign.com/articles/the-legend-of-zelda-breath-of-the-wild-2-release-date-gameplay-news-rumors', 'score': 0.94145, 'images': None}),\n", + " Document(page_content='It was eventually released on March 3, 2017, as a launch game for the Switch and the final Nintendo game for the Wii U. It received widespread acclaim and won numerous Game of the Year accolades. Critics praised its open-ended gameplay, open-world design, and attention to detail, though some criticized its technical performance.', metadata={'title': 'The Legend of Zelda: Breath of the Wild - Wikipedia', 'source': 'https://en.wikipedia.org/wiki/The_Legend_of_Zelda:_Breath_of_the_Wild', 'score': 0.92102, 'images': None})]" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import os\n", + "from langchain.retrievers.tavily_search_api import TavilySearchAPIRetriever\n", + "\n", + "os.environ[\"TAVILY_API_KEY\"] = \"YOUR_API_KEY\"\n", + "\n", + "retriever = TavilySearchAPIRetriever(k=4)\n", + "\n", + "retriever.invoke(\"what year was breath of the wild released?\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.5" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/libs/langchain/langchain/retrievers/__init__.py b/libs/langchain/langchain/retrievers/__init__.py index 372093dc10f..33a259b64ec 100644 --- a/libs/langchain/langchain/retrievers/__init__.py +++ b/libs/langchain/langchain/retrievers/__init__.py @@ -49,6 +49,7 @@ from langchain.retrievers.re_phraser import RePhraseQueryRetriever from langchain.retrievers.remote_retriever import RemoteLangChainRetriever from langchain.retrievers.self_query.base import SelfQueryRetriever from langchain.retrievers.svm import SVMRetriever +from langchain.retrievers.tavily_search_api import TavilySearchAPIRetriever from langchain.retrievers.tfidf import TFIDFRetriever from langchain.retrievers.time_weighted_retriever import ( TimeWeightedVectorStoreRetriever, @@ -82,6 +83,7 @@ __all__ = [ "RemoteLangChainRetriever", "SVMRetriever", "SelfQueryRetriever", + "TavilySearchAPIRetriever", "TFIDFRetriever", "BM25Retriever", "TimeWeightedVectorStoreRetriever", diff --git a/libs/langchain/langchain/retrievers/tavily_search_api.py b/libs/langchain/langchain/retrievers/tavily_search_api.py new file mode 100644 index 00000000000..1385e9ddd22 --- /dev/null +++ b/libs/langchain/langchain/retrievers/tavily_search_api.py @@ -0,0 +1,82 @@ +import os +from enum import Enum +from typing import Any, Dict, List, Optional + +from langchain.callbacks.manager import CallbackManagerForRetrieverRun +from langchain.schema import Document +from langchain.schema.retriever import BaseRetriever + + +class SearchDepth(Enum): + BASIC = "basic" + ADVANCED = "advanced" + + +class TavilySearchAPIRetriever(BaseRetriever): + """Tavily Search API retriever.""" + + k: int = 10 + include_generated_answer: bool = False + include_raw_content: bool = False + include_images: bool = False + search_depth: SearchDepth = SearchDepth.BASIC + include_domains: Optional[List[str]] = None + exclude_domains: Optional[List[str]] = None + kwargs: Optional[Dict[str, Any]] = {} + api_key: Optional[str] = None + + def _get_relevant_documents( + self, query: str, *, run_manager: CallbackManagerForRetrieverRun + ) -> List[Document]: + try: + from tavily import Client + except ImportError: + raise ValueError( + "Tavily python package not found. " + "Please install it with `pip install tavily-python`." + ) + + tavily = Client(api_key=self.api_key or os.environ["TAVILY_API_KEY"]) + max_results = self.k if not self.include_generated_answer else self.k - 1 + response = tavily.search( + query=query, + max_results=max_results, + search_depth=self.search_depth.value, + include_answer=self.include_generated_answer, + include_domains=self.include_domains, + exclude_domains=self.exclude_domains, + include_raw_content=self.include_raw_content, + include_images=self.include_images, + **self.kwargs + ) + docs = [ + Document( + page_content=result.get("content", "") + if not self.include_raw_content + else result.get("raw_content", ""), + metadata={ + "title": result.get("title", ""), + "source": result.get("url", ""), + **{ + k: v + for k, v in result.items() + if k not in ("content", "title", "url", "raw_content") + }, + "images": response.get("images"), + }, + ) + for result in response.get("results") + ] + if self.include_generated_answer: + docs = [ + Document( + page_content=response.get("answer", ""), + metadata={ + "title": "Suggested Answer", + "source": "https://tavily.com/", + }, + ), + *docs, + ] + + return docs