test.py

import asyncio
import os
import logging
import tempfile
from pathlib import Path
from app import JARVISAgent, llm_client, llm_type, llm_model, embedder
from tools.search import search_tool, multi_hop_search_tool
from tools.file_parser import file_parser_tool
from tools.image_parser import image_parser_tool
from tools.calculator import calculator_tool
from tools.document_retriever import document_retriever_tool
from tools.duckduckgo_search import duckduckgo_search_tool
from tools.weather_info import weather_info_tool
from tools.hub_stats import hub_stats_tool
from tools.guest_info import guest_info_retriever_tool
from tools.file_fetcher import fetch_task_file
from tools.answer_generator import preprocess_question, filter_results
from state import validate_state, reset_state, JARVISState

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

async def test_tools():
    """Test all tools."""
    logger.info("Testing Search Tool (SerpAPI)...")
    try:
        if not os.getenv("SERPAPI_API_KEY"):
            logger.warning("Search Warning: SERPAPI_API_KEY not set")
        else:
            result = await search_tool.ainvoke({"query": "What is the capital of France?"})
            logger.info(f"Search Result: {result}")
    except Exception as e:
        logger.error(f"Search Error: {e}")

    logger.info("Testing Multi-Hop Search Tool...")
    try:
        result = await multi_hop_search_tool.ainvoke({
            "query": "What is the population of France's capital?",
            "steps": 2,
            "llm_client": llm_client,
            "llm_type": llm_type,
            "llm_model": llm_model
        })
        logger.info(f"Multi-Hop Search Result: {result}")
    except Exception as e:
        logger.error(f"Multi-Hop Search Error: {e}")

    logger.info("Testing DuckDuckGo Search Tool...")
    try:
        result = await duckduckgo_search_tool.ainvoke({
            "query": "What is the capital of France?",
            "original_query": "What is the capital of France?",
            "embedder": embedder
        })
        logger.info(f"DuckDuckGo Result: {result}")
    except Exception as e:
        logger.error(f"DuckDuckGo Error: {e}")

    logger.info("Testing Weather Info Tool...")
    try:
        if not os.getenv("OPENWEATHERMAP_API_KEY"):
            logger.warning("Weather Warning: OPENWEATHERMAP_API_KEY not set")
        else:
            result = await weather_info_tool.ainvoke({"location": "London"})
            logger.info(f"Weather Result: {result}")
    except Exception as e:
        logger.error(f"Weather Error: {e}")

    logger.info("Testing Document Retriever Tool...")
    try:
        from PyPDF2 import PdfWriter
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
            writer = PdfWriter()
            from PyPDF2.generic import NameObject, create_string_object
            page = writer.add_blank_page(width=72, height=72)
            page[NameObject("/Contents")] = create_string_object("Sample document content for testing.")
            writer.write(tmp)
            tmp_path = tmp.name
        result = await document_retriever_tool.ainvoke({
            "task_id": "test_task",
            "query": "Sample question",
            "file_path": tmp_path
        })
        logger.info(f"Document Retriever Result: {result}")
        os.unlink(tmp_path)
    except Exception as e:
        logger.error(f"Document Retriever Error: {e}")

    logger.info("Testing Image Parser Tool...")
    try:
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            # Create a minimal PNG (1x1 pixel)
            from PIL import Image
            img = Image.new('RGB', (1, 1), color='white')
            img.save(tmp.name, 'PNG')
            tmp_path = tmp.name
        result = await image_parser_tool.ainvoke({"task_id": "test_task", "file_path": tmp_path})
        logger.info(f"Image Parser Result: {result}")
        os.unlink(tmp_path)
    except Exception as e:
        logger.error(f"Image Parser Error: {e}")

    logger.info("Testing File Parser Tool...")
    try:
        with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as tmp:
            tmp.write(b"Sample text file content")
            tmp_path = tmp.name
        result = await file_parser_tool.ainvoke({
            "task_id": "test_task",
            "file_type": "txt",
            "file_path": tmp_path,
            "query": "What is in the file?"
        })
        logger.info(f"File Parser Result: {result}")
        os.unlink(tmp_path)
    except Exception as e:
        logger.error(f"File Parser Error: {e}")

    logger.info("Testing Calculator Tool...")
    try:
        result = await calculator_tool.ainvoke({"expression": "2 + 2"})
        logger.info(f"Calculator Result: {result}")
    except Exception as e:
        logger.error(f"Calculator Error: {e}")

    logger.info("Testing Hub Stats Tool...")
    try:
        if not os.getenv("HUGGINGFACEHUB_API_TOKEN"):
            logger.warning("Hub Stats Warning: HUGGINGFACEHUB_API_TOKEN not set")
        else:
            result = await hub_stats_tool.ainvoke({"author": "meta-llama"})
            logger.info(f"Hub Stats Result: {result}")
    except Exception as e:
        logger.error(f"Hub Stats Error: {e}")

    logger.info("Testing Guest Info Retriever Tool...")
    try:
        result = await guest_info_retriever_tool.ainvoke({"query": "Who is the guest named John?"})
        logger.info(f"Guest Info Result: {result}")
    except Exception as e:
        logger.error(f"Guest Info Error: {e}")

async def test_file_fetcher():
    """Test file fetcher."""
    logger.info("Testing File Fetcher...")
    try:
        result = await fetch_task_file("8e867cd7-cff9-4e6c-867a-ff5ddc2550be", "Sample question with data")
        logger.info(f"File Fetcher Result: {result}")
    except Exception as e:
        logger.error(f"File Fetcher Error: {e}")

async def test_answer_generator():
    """Test answer generator functions."""
    logger.info("Testing Preprocess Question...")
    try:
        result = await preprocess_question("What's the weather in Paris?")
        logger.info(f"Preprocess Question Result: {result}")
    except Exception as e:
        logger.error(f"Preprocess Question Error: {e}")

    logger.info("Testing Filter Results...")
    try:
        results = ["Paris is the capital of France.", "Florida is a state.", "Paris is in Texas."]
        filtered = filter_results(results, "What is the capital of France?")
        logger.info(f"Filter Results: {filtered}")
    except Exception as e:
        logger.error(f"Filter Results Error: {e}")

async def test_state_management():
    """Test state management functions."""
    logger.info("Testing Reset State...")
    try:
        state = reset_state("test_task", "What is the capital of France?")
        logger.info(f"Reset State Result: {state}")
    except Exception as e:
        logger.error(f"Reset State Error: {e}")

    logger.info("Testing Validate State...")
    try:
        invalid_state = {"task_id": "", "question": ""}
        validate_state(invalid_state)
        logger.error("Validate State should have failed")
    except ValueError as e:
        logger.info(f"Validate State Error (expected): {e}")

    try:
        valid_state = reset_state("test_task", "Sample question")
        validated = validate_state(valid_state)
        logger.info(f"Validate State Result: {validated}")
    except Exception as e:
        logger.error(f"Validate State Error: {e}")

async def test_agent():
    """Test JARVISAgent with various cases."""
    logger.info("Testing JARVISAgent (Simple Question)...")
    try:
        agent = JARVISAgent()
        answer = await agent.process_question("test_task", "What is the capital of France?")
        logger.info(f"JARVISAgent Answer: {answer}")
    except Exception as e:
        logger.error(f"JARVISAgent Error: {e}")

    logger.info("Testing JARVISAgent (Edge Case: Empty Question)...")
    try:
        agent = JARVISAgent()
        answer = await agent.process_question("test_task", "")
        logger.info(f"JARVISAgent Empty Question Answer: {answer}")
    except Exception as e:
        logger.info(f"JARVISAgent Empty Question Error (expected): {e}")

async def main():
    required_envs = [
        "HUGGINGFACEHUB_API_TOKEN",
        "TOGETHER_API_KEY",
        "OPENWEATHERMAP_API_KEY",
        "SERPAPI_API_KEY"
    ]
    for env in required_envs:
        if not os.getenv(env):
            logger.warning(f"{env} not set, some tools may fail")

    await test_tools()
    await test_file_fetcher()
    await test_answer_generator()
    await test_state_management()
    await test_agent()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except Exception as e:
        logger.error(f"Test script failed: {e}")