ds_zagres_ai/test_team_chats.py

#!/usr/bin/env python3
"""
Test script for team chats functionality.
This script tests the team chat features of the AI service, including:
- Creating team chats
- Adding members to team chats
- Sending messages to team chats
- Verifying integration with OpenWebUI channels
"""

import requests
import json
import argparse
import time
import sys
from typing import Dict, Any, List, Tuple, Optional

def test_health(api_url: str) -> bool:
    """Test the health endpoint."""
    print(f"Testing health endpoint at: {api_url}/health")

    try:
        response = requests.get(f"{api_url}/health", timeout=30)
        response.raise_for_status()

        print("✅ Health check successful!")
        print(f"Response: {response.json()}")

        return True
    except Exception as e:
        print(f"❌ ERROR: Health check failed: {str(e)}")
        return False

def test_create_team_chat(api_url: str, user_id: str, title: str, model_id: str) -> Tuple[bool, Optional[str]]:
    """
    Test creating a team chat.
    
    Args:
        api_url: API URL
        user_id: User ID
        title: Chat title
        model_id: Model ID
        
    Returns:
        Tuple of (success, chat_id)
    """
    print(f"Creating team chat with title: {title}")
    
    try:
        response = requests.post(
            f"{api_url}/chats",
            headers={"Content-Type": "application/json"},
            json={
                "user_id": user_id,
                "title": title,
                "model_id": model_id,
                "is_team_chat": True
            },
            timeout=60
        )
        
        response.raise_for_status()
        chat = response.json()
        chat_id = chat.get("id")
        
        print(f"✅ Team chat created successfully with ID: {chat_id}")
        print(f"OpenWebUI Channel ID: {chat.get('openwebui_channel_id', 'Not available')}")
        print(f"Team Members: {chat.get('team_members', [])}")
        
        return True, chat_id
    except Exception as e:
        print(f"❌ ERROR: Failed to create team chat: {str(e)}")
        return False, None

def test_add_team_member(api_url: str, chat_id: str, user_id: str) -> bool:
    """
    Test adding a member to a team chat.
    
    Args:
        api_url: API URL
        chat_id: Chat ID
        user_id: User ID to add
        
    Returns:
        Success status
    """
    print(f"Adding user {user_id} to team chat {chat_id}")
    
    try:
        response = requests.post(
            f"{api_url}/chats/{chat_id}/members/{user_id}",
            timeout=30
        )
        
        response.raise_for_status()
        
        print(f"✅ User {user_id} added to team chat successfully")
        return True
    except Exception as e:
        print(f"❌ ERROR: Failed to add user to team chat: {str(e)}")
        return False

def test_send_team_message(api_url: str, chat_id: str, user_id: str, message: str) -> bool:
    """
    Test sending a message to a team chat.
    
    Args:
        api_url: API URL
        chat_id: Chat ID
        user_id: User ID
        message: Message content
        
    Returns:
        Success status
    """
    print(f"Sending message to team chat {chat_id}: {message}")
    
    try:
        response = requests.post(
            f"{api_url}/chats/{chat_id}/messages",
            headers={"Content-Type": "application/json"},
            json={
                "message": message,
                "user_id": user_id,
                "temperature": 0.7,
                "max_tokens": 500
            },
            timeout=120
        )
        
        response.raise_for_status()
        
        result = response.json()
        print(f"✅ Message sent successfully")
        print(f"AI Response: {result.get('content', 'No content in response')[:100]}...")
        
        return True
    except Exception as e:
        print(f"❌ ERROR: Failed to send message to team chat: {str(e)}")
        return False

def test_get_user_chats(api_url: str, user_id: str) -> Tuple[bool, List[Dict[str, Any]]]:
    """
    Test getting all chats for a user.
    
    Args:
        api_url: API URL
        user_id: User ID
        
    Returns:
        Tuple of (success, chats)
    """
    print(f"Getting chats for user {user_id}")
    
    try:
        response = requests.get(
            f"{api_url}/chats/user/{user_id}",
            timeout=30
        )
        
        response.raise_for_status()
        
        chats = response.json()
        team_chats = [chat for chat in chats if chat.get("is_team_chat", False)]
        
        print(f"✅ Found {len(chats)} chats for user {user_id}, {len(team_chats)} are team chats")
        
        return True, chats
    except Exception as e:
        print(f"❌ ERROR: Failed to get user chats: {str(e)}")
        return False, []

def main():
    """Main function."""
    parser = argparse.ArgumentParser(description='Test team chats functionality')
    parser.add_argument('--api-url', type=str, default='http://localhost:5252', help='AI service API URL')
    parser.add_argument('--model', type=str, default='llama3.1', help='Model to use for testing')
    parser.add_argument('--user-id', type=str, default='test_user', help='User ID for testing')
    parser.add_argument('--second-user-id', type=str, default='test_user2', help='Second user ID for testing team membership')
    parser.add_argument('--message', type=str, default='Hello team! Can you help me with a project?', help='Message to send to the team chat')
    
    args = parser.parse_args()
    
    print("=== Team Chats Functionality Test ===")
    print(f"API URL: {args.api_url}")
    print(f"Model: {args.model}")
    print(f"User ID: {args.user_id}")
    print()
    
    # Test health endpoint
    print("=== Testing Health Endpoint ===")
    health_success = test_health(args.api_url)
    print()
    
    if not health_success:
        print("❌ Health check failed. Aborting tests.")
        sys.exit(1)
    
    # Test creating a team chat
    print("=== Creating Team Chat ===")
    chat_success, chat_id = test_create_team_chat(
        args.api_url, 
        args.user_id, 
        f"Team Chat Test {time.strftime('%Y-%m-%d %H:%M:%S')}", 
        args.model
    )
    print()
    
    if not chat_success or not chat_id:
        print("❌ Failed to create team chat. Aborting tests.")
        sys.exit(1)
    
    # Test adding a team member
    print("=== Adding Team Member ===")
    member_success = test_add_team_member(args.api_url, chat_id, args.second_user_id)
    print()
    
    # Test sending a message to the team chat
    print("=== Sending Message to Team Chat ===")
    message_success = test_send_team_message(args.api_url, chat_id, args.user_id, args.message)
    print()
    
    # Test getting user chats
    print("=== Getting User Chats ===")
    chats_success, chats = test_get_user_chats(args.api_url, args.user_id)
    print()
    
    # Print summary
    print("=== Test Summary ===")
    print(f"Health Endpoint: {'✅ SUCCESS' if health_success else '❌ FAILED'}")
    print(f"Create Team Chat: {'✅ SUCCESS' if chat_success else '❌ FAILED'}")
    print(f"Add Team Member: {'✅ SUCCESS' if member_success else '❌ FAILED'}")
    print(f"Send Team Message: {'✅ SUCCESS' if message_success else '❌ FAILED'}")
    print(f"Get User Chats: {'✅ SUCCESS' if chats_success else '❌ FAILED'}")
    
    if chat_success and chat_id:
        print(f"\nCreated team chat ID: {chat_id}")
        print("You can now check the OpenWebUI interface to see if the channel was created")
        print("and if the message was sent to the channel.")
        print(f"OpenWebUI URL: http://104.225.217.215:8080/")

if __name__ == "__main__":
    main()