File size: 3,824 Bytes
99b8067
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
93
94
95
96
97
98
99
100
101
102
103
104
#!/usr/bin/env python3
"""
Test script for improved Ollama function calling
"""

import sys
import os
from pathlib import Path

# Add the atles directory to the Python path
sys.path.insert(0, str(Path(__file__).parent / 'atles'))

def test_function_calling():
    """Test the improved function calling capabilities."""
    print("πŸ§ͺ Testing Improved Ollama Function Calling...")
    print("=" * 60)
    
    try:
        from atles.ollama_client_enhanced import OllamaFunctionCaller
        
        # Create the enhanced client
        print("πŸ“± Creating Enhanced Ollama Client...")
        client = OllamaFunctionCaller()
        
        # Check availability
        print("πŸ” Checking Ollama availability...")
        if not client.is_available():
            print("❌ Ollama is not running. Please start Ollama with: ollama serve")
            return
        
        print("βœ… Ollama is available!")
        
        # Show available functions
        print("\nπŸ”§ Available Functions:")
        schema = client.get_function_schema()
        for func in schema["functions"]:
            print(f"  β€’ {func['name']}: {func['description']}")
        
        # Test function execution directly
        print("\nπŸ§ͺ Testing Direct Function Execution...")
        
        # Test search_code function
        print("πŸ” Testing search_code function...")
        try:
            result = client.search_code_datasets("python flask", "python", "github_code")
            print(f"  βœ… Found {len(result)} code examples")
            if result:
                first_result = result[0]
                print(f"    - Example: {first_result.get('title', 'No title')}")
        except Exception as e:
            print(f"  ❌ search_code failed: {e}")
        
        # Test system info function
        print("\nπŸ’» Testing get_system_info function...")
        try:
            info = client.get_system_info()
            print(f"  βœ… Platform: {info.get('platform', 'Unknown')}")
            print(f"  βœ… Python: {info.get('python_version', 'Unknown')}")
        except Exception as e:
            print(f"  ❌ get_system_info failed: {e}")
        
        # Test file listing function
        print("\nπŸ“ Testing list_files function...")
        try:
            result = client.list_files(".", "*.py")
            print(f"  βœ… Found {len(result)} Python files")
            for file in result[:3]:  # Show first 3
                print(f"    - {file}")
        except Exception as e:
            print(f"  ❌ list_files failed: {e}")
        
        # Test function call handling with different formats
        print("\nπŸ”„ Testing Function Call Handling...")
        
        # Test 1: Standard FUNCTION_CALL format
        test_response1 = "FUNCTION_CALL:get_system_info:{}"
        print(f"  Testing format 1: {test_response1}")
        result1 = client.handle_function_call(test_response1)
        print(f"    Result: {result1[:100]}...")
        
        # Test 2: Alternative format
        test_response2 = "get_system_info:{}"
        print(f"  Testing format 2: {test_response2}")
        result2 = client.handle_function_call(test_response2)
        print(f"    Result: {result2[:100]}...")
        
        # Test 3: Search code format
        test_response3 = 'search_code:{"query": "python flask", "language": "python"}'
        print(f"  Testing format 3: {test_response3}")
        result3 = client.handle_function_call(test_response3)
        print(f"    Result: {result3[:100]}...")
        
        # Close client
        client.close()
        print("\nβœ… Function calling test completed!")
        
    except Exception as e:
        print(f"❌ Test failed: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    test_function_calling()