scratch/test_student_simulation.py

"""Phase 2 End-to-End Student Simulation Test.

Simulates a real Kerala HSE Plus Two Physics student using DocDoe AI:
1. Creates a test source (Physics chapter text)
2. Tests that embeddings are generated during chunking
3. Tests hybrid retrieval (TF-IDF + Vector)
4. Tests weak-topic memory tracking
5. Tests chat with source-grounded answers
6. Tests streaming chat response

Run: .venv\Scripts\python scratch\test_student_simulation.py
"""
import asyncio
import json
import os
import sys
import time

# Ensure app is importable
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from dotenv import load_dotenv
load_dotenv(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".env"))


def separator(title: str) -> None:
    print(f"\n{'='*70}")
    print(f"  {title}")
    print(f"{'='*70}")


def test_step(name: str, passed: bool, detail: str = "") -> None:
    status = "✅ PASS" if passed else "❌ FAIL"
    print(f"  {status} | {name}")
    if detail:
        print(f"         → {detail}")


# ─── Test Data: Plus Two Physics Chapter ──────────────────────────────────────

PHYSICS_TEXT = """
Chapter 1: Electric Charges and Fields

1.1 Introduction
Electrostatics deals with the study of forces, fields and potentials arising from static charges.

1.2 Electric Charge
Electric charge is a fundamental property of matter. There are two types of charges:
positive charge and negative charge. Like charges repel each other and unlike charges
attract each other. The SI unit of charge is Coulomb (C).

Conservation of charge: The total charge of an isolated system is always conserved.
Quantization of charge: Charge exists in discrete packets. q = ne, where n is an integer
and e = 1.6 × 10^-19 C is the elementary charge.

1.3 Coulomb's Law
The force between two point charges is directly proportional to the product of their
magnitudes and inversely proportional to the square of the distance between them.
F = kq1q2/r², where k = 9 × 10^9 Nm²/C² (Coulomb's constant)

The force is along the line joining the two charges. It is repulsive for like charges
and attractive for unlike charges.

1.4 Electric Field
Electric field is the space around a charge where its influence can be felt.
E = F/q₀ where q₀ is a small positive test charge.

Electric field due to a point charge: E = kQ/r²
The direction of E is radially outward for positive charge and radially inward
for negative charge.

1.5 Electric Field Lines
Electric field lines are imaginary lines drawn in such a way that the tangent at
any point gives the direction of the electric field at that point.
Properties:
- Field lines start from positive charges and end at negative charges
- Two field lines never intersect
- Field lines are perpendicular to the surface of a conductor
- The density of field lines represents the strength of the field

1.6 Electric Dipole
An electric dipole consists of two equal and opposite charges separated by a small
distance. Dipole moment p = q × 2a (direction from -q to +q)

Electric field on the axial line of a dipole: E = 2kp/r³
Electric field on the equatorial line: E = kp/r³

1.7 Gauss's Law
The total electric flux through any closed surface is 1/ε₀ times the total charge
enclosed by the surface. Φ = q/ε₀

Applications:
- Field due to an infinite long straight charged wire: E = λ/(2πε₀r)
- Field due to a uniformly charged infinite plane sheet: E = σ/(2ε₀)
- Field due to a uniformly charged thin spherical shell:
  Outside: E = kQ/r²
  Inside: E = 0

Previous Year Questions (Kerala HSE):
2023: State and explain Coulomb's law. Derive the expression for electric field
due to a point charge. (5 marks)
2022: What is an electric dipole? Derive the expression for the electric field
at a point on the axial line of an electric dipole. (5 marks)
2024: State Gauss's law. Using Gauss's law, derive the expression for the electric
field due to a uniformly charged spherical shell. (5 marks)
"""


def main() -> None:
    print("\n🎓 DocDoe AI — Phase 2 Student Simulation Test")
    print("   Simulating: Kerala HSE Plus Two Physics Student\n")

    # ── Step 1: Database & Model Check ──
    separator("Step 1: Database & Model Initialization")

    from app.core.database import init_db, SessionLocal
    init_db()
    test_step("Database initialized with Phase 2 tables", True)

    db = SessionLocal()

    # Check if embedding column exists
    from sqlalchemy import inspect as sa_inspect
    from app.core.database import engine
    inspector = sa_inspect(engine)
    chunk_columns = {c["name"] for c in inspector.get_columns("document_chunks")}
    has_embedding = "embedding" in chunk_columns
    test_step("document_chunks.embedding column exists", has_embedding, f"Columns: {sorted(chunk_columns)}")

    # Check weak_topics table
    tables = set(inspector.get_table_names())
    has_weak = "student_weak_topics" in tables
    test_step("student_weak_topics table exists", has_weak)

    # ── Step 2: Create a Physics source with embeddings ──
    separator("Step 2: Upload Physics Chapter (with Embedding Generation)")

    from app.models.document import Document
    from app.services.chunking import replace_document_chunks

    # Create document
    doc = Document(
        user_id="usr_demo_student",
        title="Physics Ch1: Electric Charges and Fields",
        file_name="physics_ch1.txt",
        file_type="text/plain",
        file_path="text://chapter_text",
        subject="Physics",
        chapter="Electric Charges and Fields",
        syllabus="Kerala HSE",
        status="ready",
        extracted_text=PHYSICS_TEXT,
        source_type="chapter_text",
    )
    db.add(doc)
    db.commit()
    db.refresh(doc)
    test_step("Document created", True, f"ID: {doc.id}")

    # Generate chunks (this now also generates embeddings!)
    t0 = time.time()
    chunks = replace_document_chunks(db, doc)
    db.commit()
    elapsed = time.time() - t0
    test_step(f"Chunks generated: {len(chunks)}", len(chunks) > 0, f"Time: {elapsed:.2f}s")

    # Check if embeddings were stored
    chunks_with_emb = [c for c in chunks if c.embedding]
    test_step(
        f"Embeddings stored: {len(chunks_with_emb)}/{len(chunks)}",
        len(chunks_with_emb) > 0,
        f"Vector dim: {len(json.loads(chunks_with_emb[0].embedding)) if chunks_with_emb else 'N/A'}",
    )

    # ── Step 3: Test Hybrid Retrieval ──
    separator("Step 3: Hybrid Search (TF-IDF + Vector + RRF)")

    from app.services.retrieval import retrieve_relevant_chunks

    test_queries = [
        "What is Coulomb's law?",
        "Explain electric dipole moment",
        "Gauss's law applications",
        "Previous year questions on electric field",
        "What is the formula for electric field due to point charge?",
    ]

    for query in test_queries:
        t0 = time.time()
        results = retrieve_relevant_chunks(db, doc.id, query, limit=3, user_id="usr_demo_student")
        elapsed = time.time() - t0
        top_score = results[0].score if results else 0
        top_heading = results[0].chunk.heading[:50] if results and results[0].chunk.heading else "N/A"
        test_step(
            f"Query: \"{query[:40]}...\"",
            len(results) > 0 and top_score > 0,
            f"Top: {top_heading} | Score: {top_score:.4f} | {elapsed:.2f}s",
        )

    # ── Step 4: Weak Topic Memory ──
    separator("Step 4: Weak Topic Memory Tracking")

    from app.services.weak_topic_service import record_weak_topic, get_user_weak_topics

    # Record some weak topics
    record_weak_topic(db, "usr_demo_student", "Coulomb's Law", "Physics")
    record_weak_topic(db, "usr_demo_student", "Coulomb's Law", "Physics")  # frequency += 1
    record_weak_topic(db, "usr_demo_student", "Gauss's Law", "Physics")
    record_weak_topic(db, "usr_demo_student", "Electric Dipole", "Physics")
    record_weak_topic(db, "usr_demo_student", "Coulomb's Law", "Physics")  # frequency += 1 again
    db.commit()

    topics = get_user_weak_topics(db, "usr_demo_student", subject="Physics", limit=5)
    test_step(
        f"Weak topics recorded & ranked",
        len(topics) >= 3 and topics[0].lower() == "coulomb's law",
        f"Top: {topics}",
    )

    # Test weak-topic boosted retrieval
    results_boosted = retrieve_relevant_chunks(db, doc.id, "Explain forces between charges", limit=3, user_id="usr_demo_student")
    results_no_boost = retrieve_relevant_chunks(db, doc.id, "Explain forces between charges", limit=3, user_id=None)
    boosted_score = results_boosted[0].score if results_boosted else 0
    normal_score = results_no_boost[0].score if results_no_boost else 0
    test_step(
        "Weak-topic boost active",
        boosted_score >= normal_score,
        f"Boosted: {boosted_score:.4f} vs Normal: {normal_score:.4f}",
    )

    # ── Step 5: Test Chat AI (Non-streaming) ──
    separator("Step 5: Chat with Sarvam AI (Source-Grounded)")

    from openai import OpenAI
    from app.core.config import get_settings

    settings = get_settings()
    if settings.sarvam_api_key:
        client = OpenAI(base_url=settings.sarvam_base_url, api_key=settings.sarvam_api_key)

        # Build context from hybrid retrieval
        from app.services.retrieval import chunks_to_context
        ctx_chunks = retrieve_relevant_chunks(db, doc.id, "What is Coulomb's law?", limit=3, user_id="usr_demo_student")
        context = chunks_to_context(ctx_chunks)

        system_prompt = (
            "You are DocDoe, an expert AI tutor for Indian students preparing for Kerala HSE board exams. "
            "Answer the student's question using the provided source context. Be clear, exam-focused, "
            "and include relevant formulas."
        )
        user_msg = f"Source context:\n{context}\n\n---\nStudent's question: What is Coulomb's law? Give the formula and SI units."

        try:
            t0 = time.time()
            response = client.chat.completions.create(
                model=settings.sarvam_model_main,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_msg},
                ],
                max_tokens=500,
                temperature=0.3,
                timeout=30.0,
            )
            elapsed = time.time() - t0
            answer = response.choices[0].message.content or ""

            # Strip <think> blocks
            import re
            answer = re.sub(r"<think>.*?</think>", "", answer, flags=re.DOTALL).strip()

            has_formula = "F" in answer and ("q" in answer.lower() or "r" in answer.lower())
            is_relevant = "coulomb" in answer.lower() or "force" in answer.lower()

            test_step(
                "Sarvam AI response received",
                bool(answer),
                f"{elapsed:.1f}s | {len(answer)} chars",
            )
            test_step(
                "Response is relevant (mentions Coulomb/force)",
                is_relevant,
            )
            test_step(
                "Response contains formula",
                has_formula,
            )

            # Print first 300 chars of answer
            print(f"\n  📝 AI Answer Preview:")
            for line in answer[:400].split("\n"):
                print(f"     {line}")
            if len(answer) > 400:
                print(f"     ... ({len(answer) - 400} more chars)")

        except Exception as exc:
            test_step("Sarvam AI chat", False, str(exc))
    else:
        test_step("Sarvam AI chat", False, "No SARVAM_API_KEY configured")

    # ── Step 6: Streaming Chat Test ──
    separator("Step 6: Streaming Chat (SSE)")

    if settings.sarvam_api_key:
        try:
            t0 = time.time()
            stream = client.chat.completions.create(
                model=settings.sarvam_model_main,
                messages=[
                    {"role": "system", "content": "You are DocDoe, a friendly AI study buddy."},
                    {"role": "user", "content": "Hi DocDoe! I'm studying for my Kerala HSE Physics exam. Can you help?"},
                ],
                max_tokens=200,
                temperature=0.5,
                stream=True,
                timeout=25.0,
            )

            tokens = []
            past_think = False
            buf = ""
            for chunk in stream:
                if not chunk.choices:
                    continue
                token = chunk.choices[0].delta.content or ""
                if not token:
                    continue

                if past_think:
                    tokens.append(token)
                else:
                    buf += token
                    if "</think>" in buf:
                        past_think = True
                        after = buf.split("</think>", 1)[1].lstrip("\n")
                        if after:
                            tokens.append(after)
                        buf = ""

            if not past_think and buf.strip():
                tokens.append(buf.strip())

            full_response = "".join(tokens)
            elapsed = time.time() - t0

            test_step(
                f"Streaming response received",
                bool(full_response),
                f"{elapsed:.1f}s | {len(tokens)} tokens | {len(full_response)} chars",
            )

            print(f"\n  💬 Streaming Response:")
            for line in full_response[:300].split("\n"):
                print(f"     {line}")

        except Exception as exc:
            test_step("Streaming chat", False, str(exc))

    # ── Cleanup ──
    separator("Cleanup")
    db.delete(doc)
    from app.models.weak_topic import StudentWeakTopic
    from sqlalchemy import delete
    db.execute(delete(StudentWeakTopic).where(StudentWeakTopic.user_id == "usr_demo_student"))
    db.commit()
    db.close()
    test_step("Test data cleaned up", True)

    print(f"\n{'='*70}")
    print(f"  🎓 Student Simulation Complete!")
    print(f"{'='*70}\n")


if __name__ == "__main__":
    main()