diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py index cdd1ad6c9..86bed0fcd 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py @@ -1,6 +1,7 @@ import ast import asyncio import os +import re import time import traceback from pathlib import Path @@ -161,10 +162,28 @@ async def generate_code(section, description, llm=None, verbose=False) -> None: async def generate_title(outline, llm=None, verbose: bool = False): """Generate a title of a notebook outline using an LLM.""" + MAX_TITLE_LENGTH = 50 title_chain = NotebookTitleChain.from_llm(llm=llm, verbose=verbose) title = await title_chain.apredict(content=outline) - title = title.strip() - title = title.strip("'\"") + if title is not None: + title = title.strip().strip("'\"") + if ( + len(title) > MAX_TITLE_LENGTH + ): # in case the title is too long because it returns chain of thought + pattern = r'"(.+?)"' # Match any text between quotes to get suggested title + title_matches = re.findall(pattern, title) # Get all matches, if available + if title_matches: # use the last match + title = ( + title_matches[-1][:MAX_TITLE_LENGTH] + .replace("'", "") + .replace('"', "") + ) # remove quotes in title + else: + title = outline["sections"][0]["content"][ + :MAX_TITLE_LENGTH + ] # use the first section content as title + if title is None or len(title) == 0: + title = "Generated_Notebook" # in case there is no title outline["title"] = title diff --git a/packages/jupyter-ai/jupyter_ai/tests/completions/test_handlers.py b/packages/jupyter-ai/jupyter_ai/tests/completions/test_handlers.py index b3d00159e..d747d9859 100644 --- a/packages/jupyter-ai/jupyter_ai/tests/completions/test_handlers.py +++ b/packages/jupyter-ai/jupyter_ai/tests/completions/test_handlers.py @@ -1,8 +1,10 @@ import json from types import SimpleNamespace from typing import Union +from unittest.mock import AsyncMock, patch import pytest +from jupyter_ai.chat_handlers.generate import generate_title from jupyter_ai.completions.handlers.default import DefaultInlineCompletionHandler from jupyter_ai.completions.models import ( InlineCompletionReply, @@ -212,3 +214,102 @@ async def test_handle_request_with_error(inline_handler): await inline_handler.tasks[0] error = inline_handler.messages[-1].model_dump().get("error", None) assert error is not None + + +# Test cases for generate_title function +@pytest.mark.asyncio +async def test_generate_title_valid_title(): + outline = { + "sections": [{"title": "Create a New File", "content": "Generated Notebook"}] + } + mock_llm = AsyncMock() + mock_llm.apredict.return_value = "Valid Title" + + with patch( + "jupyter_ai.chat_handlers.generate.NotebookTitleChain.from_llm", + return_value=mock_llm, + ): + await generate_title(outline, llm=mock_llm, verbose=False) + assert outline["title"] == "Valid Title" + + +@pytest.mark.asyncio +async def test_generate_title_long_title(): + outline = { + "sections": [{"title": "Create a New File", "content": "Generated Notebook"}] + } + mock_llm = AsyncMock() + max_title_length = 50 + + mock_llm.apredict.return_value = '"This is a very long title to "Generate Notebook" that exceeds fifty characters"' + with patch( + "jupyter_ai.chat_handlers.generate.NotebookTitleChain.from_llm", + return_value=mock_llm, + ): + await generate_title(outline, llm=mock_llm, verbose=False) + assert outline["title"] == "Generate Notebook"[:max_title_length] + + +@pytest.mark.asyncio +async def test_generate_title_long_title_nosubstring(): + outline = { + "sections": [{"title": "Create a New File", "content": "Generated Notebook"}] + } + mock_llm = AsyncMock() + max_title_length = 50 + + mock_llm.apredict.return_value = ( + '"This is a very long title to Generate Notebook that exceeds fifty characters"' + ) + with patch( + "jupyter_ai.chat_handlers.generate.NotebookTitleChain.from_llm", + return_value=mock_llm, + ): + await generate_title(outline, llm=mock_llm, verbose=False) + assert outline["title"] == "Generated Notebook" + + +@pytest.mark.asyncio +async def test_generate_title_with_quotes(): + outline = { + "sections": [{"title": "Create a New File", "content": "Generated Notebook"}] + } + mock_llm = AsyncMock() + mock_llm.apredict.return_value = "'\"Quoted Title\"'" + + with patch( + "jupyter_ai.chat_handlers.generate.NotebookTitleChain.from_llm", + return_value=mock_llm, + ): + await generate_title(outline, llm=mock_llm, verbose=False) + assert outline["title"] == "Quoted Title" + + +@pytest.mark.asyncio +async def test_generate_title_none_returned(): + outline = { + "sections": [{"title": "Create a New File", "content": "Generated Notebook"}] + } + mock_llm = AsyncMock() + mock_llm.apredict.return_value = None + + with patch( + "jupyter_ai.chat_handlers.generate.NotebookTitleChain.from_llm", + return_value=mock_llm, + ): + await generate_title(outline, llm=mock_llm, verbose=False) + assert outline["title"] == "Generated_Notebook" + + +@pytest.mark.asyncio +async def test_generate_title_none_returned_no_content(): + outline = {"sections": [{"title": "Create a New File", "content": ""}]} + mock_llm = AsyncMock() + mock_llm.apredict.return_value = "" + + with patch( + "jupyter_ai.chat_handlers.generate.NotebookTitleChain.from_llm", + return_value=mock_llm, + ): + await generate_title(outline, llm=mock_llm, verbose=False) + assert outline["title"] == "Generated_Notebook"