108 lines
3.4 KiB
Python
108 lines
3.4 KiB
Python
"""Tests for Tavliy web search."""
|
|
|
|
import pytest
|
|
import logging
|
|
|
|
from modules.interfaces.interfaceChatModel import ActionResult
|
|
from gateway.modules.interfaces.interfaceWebModel import (
|
|
WebSearchRequest,
|
|
WebCrawlRequest,
|
|
WebScrapeRequest,
|
|
)
|
|
from gateway.modules.connectors.connectorWebTavily import ConnectorTavily
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.expensive
|
|
async def test_tavily_connector_search_test_live_api():
|
|
logger.info("Testing Tavliy connector search with live API calls")
|
|
|
|
# Test request
|
|
request = WebSearchRequest(query="How old is the Earth?", max_results=5)
|
|
|
|
# Tavily instance
|
|
connectorWebTavily = await ConnectorTavily.create()
|
|
|
|
# Search test
|
|
action_result = await connectorWebTavily.search_urls(request=request)
|
|
|
|
# Check results
|
|
assert isinstance(action_result, ActionResult)
|
|
|
|
logger.info("=" * 20)
|
|
logger.info(f"Action result success status: {action_result.success}")
|
|
logger.info(f"Action result error: {action_result.error}")
|
|
logger.info(f"Action result label: {action_result.resultLabel}")
|
|
|
|
logger.info("Documents:")
|
|
for doc in action_result.documents:
|
|
logger.info("-" * 10)
|
|
logger.info(f" - Document Name: {doc.documentName}")
|
|
logger.info(f" - Document Mime Type: {doc.mimeType}")
|
|
logger.info(f" - Document Data: {doc.documentData}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.expensive
|
|
async def test_tavily_connector_crawl_test_live_api():
|
|
logger.info("Testing Tavily connector crawl with live API calls")
|
|
|
|
# Test request
|
|
urls = [
|
|
"https://en.wikipedia.org/wiki/Earth",
|
|
"https://valueon.ch",
|
|
]
|
|
request = WebCrawlRequest(urls=urls)
|
|
|
|
# Tavily instance
|
|
connectorWebTavily = await ConnectorTavily.create()
|
|
|
|
# Crawl test
|
|
action_result = await connectorWebTavily.crawl_urls(request=request)
|
|
|
|
# Check results
|
|
assert isinstance(action_result, ActionResult)
|
|
|
|
logger.info("=" * 20)
|
|
logger.info(f"Action result success status: {action_result.success}")
|
|
logger.info(f"Action result error: {action_result.error}")
|
|
logger.info(f"Action result label: {action_result.resultLabel}")
|
|
|
|
logger.info("Documents:")
|
|
for doc in action_result.documents:
|
|
logger.info("-" * 10)
|
|
logger.info(f" - Document Name: {doc.documentName}")
|
|
logger.info(f" - Document Mime Type: {doc.mimeType}")
|
|
logger.info(f" - Document Data: {doc.documentData}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.expensive
|
|
async def test_tavily_connector_scrape_test_live_api():
|
|
logger.info("Testing Tavily connector scrape with live API calls")
|
|
|
|
# Test request with query
|
|
request = WebScrapeRequest(query="How old is the Earth?", max_results=3)
|
|
|
|
# Tavily instance
|
|
connectorWebTavily = await ConnectorTavily.create()
|
|
|
|
# Scrape test
|
|
action_result = await connectorWebTavily.scrape(request=request)
|
|
|
|
# Check results
|
|
assert isinstance(action_result, ActionResult)
|
|
|
|
logger.info("=" * 20)
|
|
logger.info(f"Action result success status: {action_result.success}")
|
|
logger.info(f"Action result error: {action_result.error}")
|
|
logger.info(f"Action result label: {action_result.resultLabel}")
|
|
|
|
logger.info("Documents:")
|
|
for doc in action_result.documents:
|
|
logger.info("-" * 10)
|
|
logger.info(f" - Document Name: {doc.documentName}")
|
|
logger.info(f" - Document Mime Type: {doc.mimeType}")
|
|
logger.info(f" - Document Data: {doc.documentData}")
|