test_new_questions.py

"""Test script to generate example questions with new question types"""
import asyncio
import json
import random
from pathlib import Path

from gsma_dataset_creation.question_generator import generate_questions_for_chunk
from gsma_dataset_creation.qa_config import QAConfig


def load_random_chunks(chunked_dir: Path, num_chunks: int = 3):
    """Load random chunks from chunked directory"""
    chunk_files = list(chunked_dir.glob("*.json"))

    if not chunk_files:
        return None

    # Select random files
    selected_files = random.sample(chunk_files, min(num_chunks, len(chunk_files)))

    chunks = []
    for chunk_file in selected_files:
        try:
            with open(chunk_file, 'r') as f:
                data = json.load(f)

            # Get chunks from file
            file_chunks = data.get("chunks", [])
            if file_chunks:
                # Pick a random chunk from the file
                chunk = random.choice(file_chunks)
                chunks.append({
                    "content": chunk.get("text", ""),
                    "filename": data.get("source_document", chunk_file.stem),
                    "position": file_chunks.index(chunk),
                    "file": chunk_file.name
                })
        except Exception as e:
            print(f"Error loading {chunk_file}: {e}")
            continue

    return chunks


async def test_question_generation():
    """Generate test questions and display results"""

    # Try to load from actual chunked data first
    chunked_dir = Path("data/discover/chunked_late_1000")

    chunks_to_test = load_random_chunks(chunked_dir, num_chunks=3)

    if not chunks_to_test:
        print("⚠️  No chunked data found, using sample content")
        chunks_to_test = [{
            "content": """Network Slicing in 5G Networks

Network slicing is a fundamental capability of 5G networks that enables the creation of multiple virtual networks on a shared physical infrastructure. Each network slice is an end-to-end logical network that provides specific network capabilities and characteristics tailored to particular use cases or customer requirements.

Key characteristics of network slicing include:

1. Isolation: Each slice operates independently, ensuring that traffic and performance in one slice does not affect others.

2. Customization: Slices can be customized with specific Quality of Service (QoS) parameters, including latency, bandwidth, reliability, and security requirements.

3. Dynamic Resource Allocation: Resources can be dynamically allocated and adjusted based on real-time demand and service requirements.

4. End-to-End Management: Network slicing spans across the radio access network (RAN), transport network, and core network, providing comprehensive service delivery.

Network slicing enables diverse 5G use cases such as enhanced Mobile Broadband (eMBB), Ultra-Reliable Low-Latency Communications (URLLC), and massive Machine-Type Communications (mMTC). For example, an autonomous vehicle application requiring URLLC characteristics would utilize a network slice optimized for ultra-low latency and high reliability, while an IoT sensor network would use a slice optimized for massive connectivity with lower data rate requirements.""",
            "filename": "5g_network_slicing.md",
            "position": 0,
            "file": "sample"
        }]
    else:
        print(f"✅ Loaded {len(chunks_to_test)} random chunks from {chunked_dir}")

    config = QAConfig(
        num_questions=5,
        max_concurrent=1,
        limit_docs=None,
        limit_chunks=None
    )

    print("=" * 80)
    print("GENERATING QUESTIONS WITH NEW QUESTION TYPES")
    print("=" * 80)

    all_results = []

    for chunk_idx, chunk_data in enumerate(chunks_to_test, 1):
        print(f"\n{'#' * 80}")
        print(f"CHUNK {chunk_idx}/{len(chunks_to_test)}: {chunk_data['filename']} (position {chunk_data['position']})")
        print(f"{'#' * 80}")
        print("\nContent preview:")
        print("-" * 80)
        print(chunk_data['content'][:300] + "..." if len(chunk_data['content']) > 300 else chunk_data['content'])
        print("-" * 80)

        print(f"\n🔄 Generating 5 questions...\n")

        chunk_metadata = {
            "filename": chunk_data['filename'],
            "position": chunk_data['position']
        }

        result = await generate_questions_for_chunk(
            chunk_content=chunk_data['content'],
            chunk_metadata=chunk_metadata,
            config=config,
            model="openai/gpt-oss-120b",
            provider="Cerebras"
        )

        if not result.success:
            print(f"❌ Generation failed: {result.error_message}")
            continue

        print(f"✅ Successfully generated {len(result.questions)} questions!\n")

        # Display reasoning if available
        if result.reasoning:
            print("\n🧠 MODEL REASONING:")
            print("-" * 80)
            print(result.reasoning)
            print("-" * 80)

        # Display reasoning details if available
        if result.reasoning_details:
            print("\n📋 REASONING DETAILS:")
            print("-" * 80)
            print(json.dumps(result.reasoning_details, indent=2))
            print("-" * 80)

        # Display each question
        for i, q in enumerate(result.questions, 1):
            print(f"\n{'=' * 80}")
            print(f"QUESTION {i} - Type: {q.question_type.upper().replace('_', ' ')}")
            print(f"{'=' * 80}")
            print(f"\n❓ Question:\n{q.question}")
            print(f"\n💬 Answer:\n{q.answer}")
            print(f"\n📊 Tokens: Question={q.question_tokens}")

        print(f"\n⏱️  Processing time: {result.processing_time_seconds:.2f}s, Tokens: {result.tokens_used}")

        # Store result
        all_results.append({
            "chunk_info": chunk_data,
            "result": result
        })

    print("\n" + "=" * 80)
    print("GENERATION COMPLETE")
    print("=" * 80)
    print(f"\nProcessed {len(all_results)} chunks")
    print(f"Total tokens: {sum(r['result'].tokens_used for r in all_results)}")
    print(f"Total time: {sum(r['result'].processing_time_seconds for r in all_results):.2f}s")

    # Save to file for detailed review
    output_file = Path("test_questions_output.json")
    output_data = {
        "chunks": [
            {
                "chunk_info": {
                    "filename": r["chunk_info"]["filename"],
                    "position": r["chunk_info"]["position"],
                    "source_file": r["chunk_info"]["file"]
                },
                "questions": [
                    {
                        "question": q.question,
                        "answer": q.answer,
                        "question_type": q.question_type,
                        "question_tokens": q.question_tokens,
                    }
                    for q in r["result"].questions
                ],
                "metadata": {
                    "model": "openai/gpt-oss-120b",
                    "reasoning": r["result"].reasoning,
                    "reasoning_details": r["result"].reasoning_details,
                    "tokens_used": r["result"].tokens_used,
                    "processing_time_seconds": r["result"].processing_time_seconds
                }
            }
            for r in all_results
        ],
        "summary": {
            "total_chunks": len(all_results),
            "total_questions": sum(len(r["result"].questions) for r in all_results),
            "total_tokens": sum(r["result"].tokens_used for r in all_results),
            "total_time_seconds": sum(r["result"].processing_time_seconds for r in all_results)
        }
    }

    with open(output_file, "w") as f:
        json.dump(output_data, f, indent=2)

    print(f"\n📁 Full output saved to: {output_file}")

if __name__ == "__main__":
    asyncio.run(test_question_generation())