dynamic-rag-bot/simple_integration_test.py at master · vblearnstowritecode/dynamic-rag-bot · GitHub

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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
#!/usr/bin/env python3
"""
Simple integration test to verify both basic and detailed queries work.
"""

from bot import ask_bot
from document_reader import load_document

def parse_fetch_action(action: str) -> tuple[str, str]:
    """Parse fetch action to extract country and question."""
    if not action.startswith("fetch:"):
        return None, None

    # Remove "fetch:" prefix and split on first two colons
    parts = action[6:].split(":", 2)

    if len(parts) >= 2:
        return parts[0], parts[1]
    elif len(parts) == 1:
        return parts[0], "general information"

    return None, None

def test_integration():
    """Test the full integration flow."""

    print("=" * 60)
    print("INTEGRATION TEST: Dynamic RAG Bot")
    print("=" * 60)

    # Test 1: Basic query - should answer directly
    print("\n1. Testing basic query: 'What's the capital of France?'")
    print("-" * 40)

    response = ask_bot("What's the capital of France?")
    print(f"Thoughts: {response.thoughts}")
    print(f"Action: {response.action}")
    print(f"Bot: {response.content}")

    # Test 2: Detailed query - should fetch document
    print("\n2. Testing detailed query: 'Tell me about French wine regions'")
    print("-" * 40)

    response = ask_bot("Tell me about French wine regions")
    print(f"Thoughts: {response.thoughts}")
    print(f"Action: {response.action}")

    if response.action.startswith("fetch:"):
        country, question = parse_fetch_action(response.action)
        if country and question:
            print(f"Fetching document for '{country}' about: '{question}'")
            doc_response = load_document(country, question)
            print(f"Bot: {doc_response.format_with_sources()}")
        else:
            print("ERROR: Could not parse fetch action")
    else:
        print(f"Bot: {response.content}")

    # Test 3: Another basic query
    print("\n3. Testing population query: 'What's the population of Japan?'")
    print("-" * 40)

    response = ask_bot("What's the population of Japan?")
    print(f"Thoughts: {response.thoughts}")
    print(f"Action: {response.action}")
    print(f"Bot: {response.content}")

    # Test 4: Another document query
    print("\n4. Testing another detailed query: 'What are the main tourist attractions in Japan?'")
    print("-" * 40)

    response = ask_bot("What are the main tourist attractions in Japan?")
    print(f"Thoughts: {response.thoughts}")
    print(f"Action: {response.action}")

    if response.action.startswith("fetch:"):
        country, question = parse_fetch_action(response.action)
        if country and question:
            print(f"Fetching document for '{country}' about: '{question}'")
            doc_response = load_document(country, question)
            print(f"Bot: {doc_response.format_with_sources()}")
        else:
            print("ERROR: Could not parse fetch action")
    else:
        print(f"Bot: {response.content}")

    print("\n" + "=" * 60)
    print("INTEGRATION TEST COMPLETE")
    print("=" * 60)

if __name__ == "__main__":
    test_integration()