diff --git a/axon/routers/tavily.py b/axon/routers/tavily.py new file mode 100644 index 0000000..a4103f4 --- /dev/null +++ b/axon/routers/tavily.py @@ -0,0 +1,52 @@ +from fastapi import APIRouter + +from axon.models import ( + CrawlRequest, + ExtractRequest, + MapRequest, + SearchRequest, + TavilyExtractRequest, + TavilySearchRequest, +) +from axon.routers.crawl import crawl +from axon.routers.extract import extract +from axon.routers.map import map_site +from axon.routers.search import search + +router = APIRouter(prefix="/tavily") + + +@router.post("/search") +async def tavily_search(request: TavilySearchRequest): + include_answer = request.include_answer not in (False, "false", "none", "basic") + response = await search( + SearchRequest( + query=request.query, + max_results=request.max_results, + search_depth=request.search_depth, + include_answer=include_answer, + include_raw_content=request.include_raw_content, + domains_whitelist=request.include_domains, + domains_blacklist=request.exclude_domains, + topic=request.topic, + time_range=request.time_range, + ) + ) + return response.model_dump() + + +@router.post("/extract") +async def tavily_extract(request: TavilyExtractRequest): + urls = [request.urls] if isinstance(request.urls, str) else request.urls + response = await extract(ExtractRequest(urls=urls, include_images=request.include_images)) + return response.model_dump() + + +@router.post("/crawl") +async def tavily_crawl(request: CrawlRequest): + return (await crawl(request)).model_dump() + + +@router.post("/map") +async def tavily_map(request: MapRequest): + return (await map_site(request)).model_dump()