Update /generate.py to handle "thinking" models

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
 
 

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"