teng-lin--notebooklm-py/tests/e2e/test_chat.py

"""E2E tests for chat functionality.

These tests require valid NotebookLM authentication.
Run with: pytest tests/e2e/test_chat.py -m e2e
"""

import pytest

from notebooklm import AskResult, ChatReference

from .conftest import requires_auth


@pytest.mark.e2e
@pytest.mark.live_chat_ask
@requires_auth
class TestChatE2E:
    """E2E tests for chat API."""

    @pytest.mark.asyncio
    async def test_ask_question_returns_answer(self, client, multi_source_notebook_id):
        """Test asking a question returns a valid answer."""
        result = await client.chat.ask(
            multi_source_notebook_id,
            "What is the main topic of these sources?",
        )

        assert isinstance(result, AskResult)
        assert result.answer
        assert len(result.answer) > 10
        assert result.conversation_id
        assert result.turn_number >= 1

    @pytest.mark.asyncio
    async def test_ask_returns_references_with_source_ids(self, client, multi_source_notebook_id):
        """Test that ask returns references with valid source IDs."""
        # Ask a question likely to generate citations
        result = await client.chat.ask(
            multi_source_notebook_id,
            "Summarize the key points with specific citations.",
        )

        assert isinstance(result, AskResult)
        assert result.answer

        # If the answer contains citations [1], [2], etc., there should be references
        if "[1]" in result.answer:
            assert len(result.references) >= 1, "Answer has citations but no references"

            # Verify references have valid source IDs
            for ref in result.references:
                assert isinstance(ref, ChatReference)
                assert ref.source_id
                # Source ID should be a UUID
                assert len(ref.source_id) == 36
                assert ref.source_id.count("-") == 4

    @pytest.mark.asyncio
    async def test_ask_returns_references_with_cited_text(self, client, multi_source_notebook_id):
        """Test that references include cited text when available."""
        result = await client.chat.ask(
            multi_source_notebook_id,
            "Quote specific passages that explain the main concept.",
        )

        assert isinstance(result, AskResult)

        # Check if any references have cited_text
        refs_with_text = [ref for ref in result.references if ref.cited_text]

        # Note: Not all citations may have cited_text depending on the response
        # So we just verify the structure is correct when present
        for ref in refs_with_text:
            assert isinstance(ref.cited_text, str)
            assert len(ref.cited_text) > 0

    @pytest.mark.asyncio
    async def test_ask_follow_up_conversation(self, client, multi_source_notebook_id):
        """Test follow-up questions use the same conversation."""
        # First question
        result1 = await client.chat.ask(
            multi_source_notebook_id,
            "What is the main topic?",
        )
        assert result1.conversation_id
        assert result1.is_follow_up is False

        # Follow-up question
        result2 = await client.chat.ask(
            multi_source_notebook_id,
            "Can you elaborate on that?",
            conversation_id=result1.conversation_id,
        )
        assert result2.conversation_id == result1.conversation_id
        assert result2.is_follow_up is True
        assert result2.turn_number > result1.turn_number

    @pytest.mark.asyncio
    async def test_delete_conversation_forces_fresh_next_ask(
        self, client, multi_source_notebook_id
    ):
        """``delete_conversation`` is the supported way to force a fresh conversation.

        Without an explicit delete, a null-``conversation_id`` ``ask()`` extends
        the most-recent server conversation (see ``ChatAPI.ask`` Note). After
        deleting, the next null-conv ask starts a brand-new turn-1 conversation.

        Note on conversation_id: ``hPTbtc`` (GET_LAST_CONVERSATION_ID) returns a
        notebook-scoped "current conversation" slot that the server reuses for
        the next turn-1 after a delete. So ``result2.conversation_id`` may equal
        ``result1.conversation_id`` even though result2 is a genuinely fresh
        conversation. The freshness signal is ``turn_number == 1`` /
        ``is_follow_up is False``, plus the server-side turn count.
        """
        result1 = await client.chat.ask(
            multi_source_notebook_id,
            "What is covered in these sources?",
        )
        assert result1.conversation_id

        await client.chat.delete_conversation(multi_source_notebook_id, result1.conversation_id)

        result2 = await client.chat.ask(
            multi_source_notebook_id,
            "Start fresh - what are the main themes?",
        )

        assert result2.is_follow_up is False
        assert result2.turn_number == 1

        # Server-side confirmation that result2 really is a fresh conversation
        # and not a follow-up to result1: the conversation should hold only the
        # one Q&A pair we just posted. Pin ``result2.conversation_id`` first so
        # we don't silently fall through to the implicit "current conversation"
        # that ``get_history`` would resolve on a ``None`` argument.
        assert result2.conversation_id
        history = await client.chat.get_history(
            multi_source_notebook_id, conversation_id=result2.conversation_id
        )
        assert len(history) == 1

    @pytest.mark.asyncio
    async def test_ask_specific_sources(self, client, multi_source_notebook_id):
        """Test asking questions about specific sources."""
        # Get sources
        sources = await client.sources.list(multi_source_notebook_id)
        if not sources:
            pytest.skip("No sources in notebook")

        # Ask about first source only
        result = await client.chat.ask(
            multi_source_notebook_id,
            "What is this source about?",
            source_ids=[sources[0].id],
        )

        assert isinstance(result, AskResult)
        assert result.answer

    @pytest.mark.asyncio
    async def test_references_have_citation_numbers(self, client, multi_source_notebook_id):
        """Test that references have sequential citation numbers."""
        result = await client.chat.ask(
            multi_source_notebook_id,
            "List the key points with citations.",
        )

        if result.references:
            # Citation numbers should be assigned sequentially
            citation_numbers = [ref.citation_number for ref in result.references]
            assert all(n is not None for n in citation_numbers)
            assert citation_numbers == list(range(1, len(citation_numbers) + 1))


@pytest.mark.e2e
@requires_auth
class TestChatHistoryE2E:
    """E2E tests for chat history and conversation turns API (khqZz RPC).

    These tests use an existing read-only notebook with pre-existing conversation
    history. They do not ask new questions, since conversation persistence takes
    time and makes tests flaky.
    """

    @pytest.mark.asyncio
    @pytest.mark.readonly
    async def test_get_conversation_turns_returns_qa(self, client, read_only_notebook_id):
        """get_conversation_turns returns Q&A turns for an existing conversation."""
        conv_id = await client.chat.get_conversation_id(read_only_notebook_id)
        if not conv_id:
            pytest.skip("No conversation history available in read-only notebook")

        turns_data = await client.chat.get_conversation_turns(
            read_only_notebook_id,
            conv_id,
            limit=2,
        )

        assert turns_data is not None
        if not turns_data:
            pytest.skip(
                "Read-only notebook has a conversation but no chat turns — "
                "cannot verify turn structure. Seed the notebook with chat messages to enable this test."
            )
        assert isinstance(turns_data[0], list)
        turns = turns_data[0]
        assert len(turns) >= 1

        turn_types = [turn[2] for turn in turns if isinstance(turn, list) and len(turn) > 2]
        assert any(t in (1, 2) for t in turn_types), "Expected question or answer turns"

    @pytest.mark.asyncio
    @pytest.mark.readonly
    async def test_get_conversation_turns_question_text(self, client, read_only_notebook_id):
        """get_conversation_turns includes question text in an existing conversation."""
        conv_id = await client.chat.get_conversation_id(read_only_notebook_id)
        if not conv_id:
            pytest.skip("No conversation history available in read-only notebook")

        turns_data = await client.chat.get_conversation_turns(
            read_only_notebook_id,
            conv_id,
            limit=2,
        )

        assert turns_data is not None
        if not turns_data:
            pytest.skip(
                "Read-only notebook has a conversation but no chat turns — "
                "cannot verify question text. Seed the notebook with chat messages to enable this test."
            )
        turns = turns_data[0]
        question_turns = [t for t in turns if isinstance(t, list) and len(t) > 3 and t[2] == 1]
        assert question_turns, "No question turn found in response"
        assert isinstance(question_turns[0][3], str)
        assert len(question_turns[0][3]) > 0

    @pytest.mark.asyncio
    @pytest.mark.readonly
    async def test_get_conversation_turns_answer_text(self, client, read_only_notebook_id):
        """get_conversation_turns includes AI answer text in an existing conversation."""
        conv_id = await client.chat.get_conversation_id(read_only_notebook_id)
        if not conv_id:
            pytest.skip("No conversation history available in read-only notebook")

        turns_data = await client.chat.get_conversation_turns(
            read_only_notebook_id,
            conv_id,
            limit=2,
        )

        assert turns_data is not None
        if not turns_data:
            pytest.skip(
                "Read-only notebook has a conversation but no chat turns — "
                "cannot verify answer text. Seed the notebook with chat messages to enable this test."
            )
        turns = turns_data[0]
        answer_turns = [t for t in turns if isinstance(t, list) and len(t) > 4 and t[2] == 2]
        assert answer_turns, "No answer turn found in response"
        answer_text = answer_turns[0][4][0][0]
        assert isinstance(answer_text, str)
        assert len(answer_text) > 0

    @pytest.mark.asyncio
    @pytest.mark.readonly
    async def test_get_conversation_id(self, client, read_only_notebook_id):
        """get_conversation_id returns an existing conversation ID."""
        conv_id = await client.chat.get_conversation_id(read_only_notebook_id)
        if not conv_id:
            pytest.skip("No conversation history available in read-only notebook")

        assert isinstance(conv_id, str)
        assert len(conv_id) > 0

    @pytest.mark.asyncio
    @pytest.mark.readonly
    async def test_get_history_returns_qa_pairs(self, client, read_only_notebook_id):
        """get_history returns Q&A pairs from existing conversation history."""
        qa_pairs = await client.chat.get_history(read_only_notebook_id)
        if not qa_pairs:
            pytest.skip("No conversation history available in read-only notebook")

        # Each entry is a (question, answer) tuple
        q, a = qa_pairs[-1]  # most recent Q&A
        assert isinstance(q, str) and q, "Question should be non-empty string"
        assert isinstance(a, str) and a, "Answer should be non-empty string"


@pytest.mark.e2e
@pytest.mark.live_chat_ask
@requires_auth
class TestChatReferencesE2E:
    """E2E tests specifically for chat references and citations."""

    @pytest.mark.asyncio
    async def test_reference_source_ids_exist_in_notebook(self, client, multi_source_notebook_id):
        """Test that reference source IDs correspond to actual sources."""
        # Get all sources in the notebook
        sources = await client.sources.list(multi_source_notebook_id)
        source_ids = {s.id for s in sources}

        # Ask a question that generates citations
        result = await client.chat.ask(
            multi_source_notebook_id,
            "Explain the main concepts with references to sources.",
        )

        # All reference source IDs should exist in the notebook
        for ref in result.references:
            assert ref.source_id in source_ids, (
                f"Reference source_id {ref.source_id} not found in notebook sources"
            )

    @pytest.mark.asyncio
    async def test_cited_text_matches_source_content(self, client, multi_source_notebook_id):
        """Test that cited text comes from the actual source content."""
        result = await client.chat.ask(
            multi_source_notebook_id,
            "Quote a specific passage from the sources.",
        )

        # For references with cited_text, verify it's non-empty
        for ref in result.references:
            if ref.cited_text:
                assert len(ref.cited_text) > 0

                # Could optionally verify against source fulltext:
                # fulltext = await client.sources.get_fulltext(
                #     multi_source_notebook_id, ref.source_id
                # )
                # assert ref.cited_text in fulltext.content