Fix title bug; simplify & update test; add section_title metadata

max-svistunov · max-svistunov · commit 996596688a3d · 2025-07-10T13:49:41.000+02:00
diff --git a/scripts/html_chunking/chunker.py b/scripts/html_chunking/chunker.py
@@ -9,7 +9,7 @@
 from bs4 import BeautifulSoup, Tag, NavigableString
 import warnings
 
-from tokenizer import count_html_tokens
+from .tokenizer import count_html_tokens
 
 # Constants
 DEFAULT_CHARS_PER_TOKEN_RATIO = 3.5
@@ -39,9 +39,9 @@ def find_first_anchor(chunk_soup: BeautifulSoup) -> Optional[str]:
 
 
 def get_document_title(soup: BeautifulSoup) -> str:
-    """Extracts the document title from the <h1> tag."""
-    h1_tag = soup.find('h1')
-    return h1_tag.get_text(strip=True) if h1_tag else "Untitled"
+    """Extracts the document title from the <title> tag."""
+    title_tag = soup.find('title')
+    return title_tag.get_text(strip=True) if title_tag else "Untitled"
 
 
 def chunk_html(
@@ -70,14 +70,18 @@ def chunk_html(
 
     try:
         soup = BeautifulSoup(html_content, 'html.parser')
-        doc_title = get_document_title(soup)
+        document_title = get_document_title(soup)
 
         if count_html_tokens(html_content, options.count_tag_tokens) <= options.max_token_limit:
-            metadata = {"docs_url": source_url, "title": doc_title}
+            metadata = {
+                "docs_url": source_url,
+                "title": document_title,
+                "section_title": document_title
+            }
             return [Chunk(text=html_content, metadata=metadata)]
     except Exception as e:
         warnings.warn("Could not pre-calculate total tokens: %s. Proceeding with chunking." % e)
-        doc_title = "Untitled"
+        document_title = "Untitled"
 
     try:
         body = soup.body or soup
@@ -86,9 +90,11 @@ def chunk_html(
         warnings.warn("A critical error occurred during semantic chunking: %s. Falling back to linear splitting." % e)
         string_chunks = _linear_split(html_content, options)
 
-    # Post-process string chunks to add stateful anchor metadata
+    # Post-process string chunks to add stateful anchor and title metadata
     final_chunks = []
     last_seen_anchor = None
+    last_heading_text = document_title
+
     for s_chunk in string_chunks:
         if not s_chunk.strip():
             continue
@@ -101,11 +107,29 @@ def chunk_html(
             
         final_anchor = last_seen_anchor
         
+        chunk_headings = chunk_soup.find_all(['h1', 'h2', 'h3', 'h4', 'h5', 'h6'])
+        if chunk_headings:
+            last_heading_text = chunk_headings[-1].get_text(strip=True)
+
+        section_title = last_heading_text
+
         full_source_url = f"{source_url}#{final_anchor}" if final_anchor else source_url
-        metadata = {"docs_url": full_source_url, "title": doc_title}
+        metadata = {
+            "docs_url": full_source_url,
+            "title": document_title,
+            "section_title": section_title
+        }
         final_chunks.append(Chunk(text=s_chunk, metadata=metadata))
 
-    return final_chunks if final_chunks else [Chunk(text=html_content, metadata={"docs_url": source_url, "title": doc_title})]
+    if not final_chunks:
+        metadata = {
+            "docs_url": source_url,
+            "title": document_title,
+            "section_title": document_title
+        }
+        return [Chunk(text=html_content, metadata=metadata)]
+
+    return final_chunks
 
 
 def _split_element_by_children(element: Tag, options: ChunkingOptions) -> List[str]:
diff --git a/scripts/html_chunking/html-stripper.py b/scripts/html_chunking/html-stripper.py
@@ -111,9 +111,18 @@ def strip_html_content(
 
         soup = BeautifulSoup(html_content, "html.parser")
 
+        # Extract title from the original document's <title> tag.
+        title_tag = soup.find('title')
+        title_text = title_tag.get_text(strip=True) if title_tag else "Untitled"
+
         if strip_mode in ['sections', 'all']:
             body_content = soup.body or soup
-            new_soup = BeautifulSoup("<html><body></body></html>", "html.parser")
+            new_soup = BeautifulSoup("<html><head></head><body></body></html>", "html.parser")
+
+            if new_soup.head:
+                new_title_tag = new_soup.new_tag("title")
+                new_title_tag.string = title_text
+                new_soup.head.append(new_title_tag)
 
             chapters = body_content.find_all("section", class_="chapter")
             if not chapters:
diff --git a/scripts/html_chunking/test_chunker.py b/scripts/html_chunking/test_chunker.py
@@ -31,18 +31,19 @@ def mock_count_html_tokens(html_string, count_tag_tokens=True):
 @patch('html_chunking.chunker.count_html_tokens', new=mock_count_html_tokens)
 class TestHtmlChunker(unittest.TestCase):
 
-    def test_chunk_html_small_input(self):
+    def test_small_input_no_chunking(self):
         """Tests that HTML smaller than the max_token_limit is not chunked."""
-        html = "<html><body><h1>My Title</h1><p>This is a small test.</p></body></html>"
+        html = "<html><head><title>Test Title</title></head><body><p>This is a small test.</p></body></html>"
         chunks = chunk_html(html, "http://example.com/small", max_token_limit=100)
         self.assertEqual(len(chunks), 1)
         self.assertEqual(chunks[0].text, html)
         self.assertEqual(chunks[0].metadata["docs_url"], "http://example.com/small")
-        self.assertEqual(chunks[0].metadata["title"], "My Title")
+        self.assertEqual(chunks[0].metadata["title"], "Test Title")
+        self.assertEqual(chunks[0].metadata["section_title"], "Test Title")
 
     def test_basic_splitting(self):
         """Tests basic splitting of multiple paragraphs."""
-        html = "<html><body>"
+        html = "<html><head><title>Basic Splitting</title></head><body>"
         for i in range(10):
             html += f"<p>This is paragraph {i}. It contains several words to simulate content.</p>"
         html += "</body></html>"
@@ -51,131 +52,96 @@ def test_basic_splitting(self):
         self.assertTrue(all(mock_count_html_tokens(c.text) <= 110 for c in chunks))
         self.assertIn("paragraph 0", chunks[0].text)
         self.assertIn("paragraph 9", chunks[-1].text)
+        self.assertEqual(chunks[0].metadata["title"], "Basic Splitting")
 
     def test_oversized_element_splitting(self):
         """Tests that a single element larger than the limit is recursively split."""
         long_text = "word " * 200
-        html = f"<html><body><div>{long_text}</div></body></html>"
+        html = f"<html><head><title>Oversized</title></head><body><div>{long_text}</div></body></html>"
         chunks = chunk_html(html, "http://example.com/oversized", max_token_limit=100)
         self.assertGreater(len(chunks), 1)
         full_text = "".join(BeautifulSoup(c.text, 'html.parser').get_text() for c in chunks)
         self.assertIn("word", full_text)
         self.assertGreater(len(full_text), 500)
+        self.assertEqual(chunks[0].metadata["title"], "Oversized")
 
     def test_table_splitting(self):
         """Tests that large tables are split, preserving the header in each chunk."""
         header = "<thead><tr><th>Header 1</th><th>Header 2</th></tr></thead>"
         rows = "".join([f"<tr><td>Row {i} Col 1</td><td>Row {i} Col 2</td></tr>" for i in range(20)])
-        html = f"<html><body><table>{header}<tbody>{rows}</tbody></table></body></html>"
+        html = f"<html><head><title>Table Test</title></head><body><table>{header}<tbody>{rows}</tbody></table></body></html>"
         chunks = chunk_html(html, "http://example.com/table", max_token_limit=100)
         self.assertGreater(len(chunks), 1)
         for chunk in chunks:
             self.assertIn("<thead>", chunk.text)
             self.assertIn("Header 1", chunk.text)
             self.assertIn("</table>", chunk.text)
+            self.assertEqual(chunk.metadata["title"], "Table Test")
         self.assertIn("Row 0", chunks[0].text)
         self.assertNotIn("Row 19", chunks[0].text)
         self.assertIn("Row 19", chunks[-1].text)
 
     def test_list_splitting(self):
         """Tests that large lists are split correctly."""
         items = "".join([f"<li>Item {i} is here.</li>" for i in range(30)])
-        html = f"<html><body><ul>{items}</ul></body></html>"
+        html = f"<html><head><title>List Test</title></head><body><ul>{items}</ul></body></html>"
         chunks = chunk_html(html, "http://example.com/list", max_token_limit=100)
         self.assertGreater(len(chunks), 1)
         for chunk in chunks:
             self.assertIn("<ul ", chunk.text)
             self.assertIn("</ul>", chunk.text)
+            self.assertEqual(chunk.metadata["title"], "List Test")
         self.assertIn("Item 0", chunks[0].text)
         self.assertIn("Item 29", chunks[-1].text)
 
-    def test_definition_list_splitting(self):
-        """Tests splitting of a definition list."""
-        items = "".join([f"<dt>Term {i}</dt><dd>Definition {i} is quite long and elaborate.</dd>" for i in range(15)])
-        html = f"<html><body><div class='variablelist'><dl>{items}</dl></div></body></html>"
-        chunks = chunk_html(html, "http://example.com/dl", max_token_limit=100)
-        self.assertGreater(len(chunks), 1)
-        for chunk in chunks:
-            self.assertIn("<dl>", chunk.text)
-            self.assertIn("</dl>", chunk.text)
-        self.assertIn("Term 0", chunks[0].text)
-        self.assertIn("Term 14", chunks[-1].text)
-
-    def test_code_splitting(self):
-        """Tests that preformatted code blocks are split by lines."""
-        code_lines = "\n".join([f"line_{i} = 'some code here';" for i in range(50)])
-        html = f"<html><body><pre>{code_lines}</pre></body></html>"
-        chunks = chunk_html(html, "http://example.com/code", max_token_limit=50)
-        self.assertGreater(len(chunks), 1)
-        for chunk in chunks:
-            self.assertIn("<pre ", chunk.text)
-            self.assertIn("</pre>", chunk.text)
-        self.assertIn("line_0", chunks[0].text)
-        self.assertIn("line_49", chunks[-1].text)
-        self.assertNotIn("line_49", chunks[0].text)
-
-    def test_heading_grouping(self):
-        """Tests that headings are grouped with the following element."""
-        html = "<html><body>"
-        for i in range(5):
-            html += f"<h2>Title {i}</h2><p>This is paragraph for title {i}. It has text.</p>"
-        html += "</body></html>"
-        chunks = chunk_html(html, "http://example.com/headings", max_token_limit=50)
-        self.assertEqual(len(chunks), 5)
-        for i, chunk in enumerate(chunks):
-            self.assertIn(f"Title {i}", chunk.text)
-            self.assertIn(f"paragraph for title {i}", chunk.text)
-
-    def test_paragraph_ending_with_colon_grouping(self):
-        """Tests grouping of a paragraph ending with a colon with the next list/table."""
-        html = ("<html><body><p>Here are the items:</p>"
-                "<ul><li>Item 1</li><li>Item 2</li></ul></body></html>")
-        chunks = chunk_html(html, "http://example.com/colon", max_token_limit=100)
-        self.assertEqual(len(chunks), 1)
-        self.assertIn("Here are the items:", chunks[0].text)
-        self.assertIn("<li>Item 1</li>", chunks[0].text)
-
-    def test_metadata_anchor_handling(self):
-        """Tests the generation of source metadata with correct anchors."""
+    def test_metadata_and_section_titles(self):
+        """Tests the generation of metadata with correct anchors and section titles."""
         html = """
-        <html><body>
-            <section id="intro"><h1>Intro</h1><p>Text</p></section>
+        <html><head><title>Main Document Title</title></head><body>
+            <section id="intro"><h1>Introduction</h1><p>Text about intro.</p></section>
             <div id="main-content">
                 <h2 id="topic1">Topic 1</h2><p>Content 1</p>
-                <p>More content 1</p>
+                <p>More content 1, still under Topic 1.</p>
             </div>
             <section id="conclusion">
-                <p>Conclusion text</p>
+                <p>Conclusion text, still under Topic 1 technically.</p>
                 <h3 id="final-thoughts">Final Thoughts</h3><p>Final words.</p>
             </section>
         </body></html>
         """
         chunks = chunk_html(html, "http://example.com/meta", max_token_limit=25)
         
-        self.assertGreaterEqual(len(chunks), 3)
+        self.assertGreaterEqual(len(chunks), 4)
 
-        # The first chunk might not have a specific anchor if it's just the title
-        self.assertIn(chunks[0].metadata["docs_url"], ["http://example.com/meta", "http://example.com/meta#intro"])
-        self.assertEqual(chunks[0].metadata["title"], "Intro")
+        # Check document title consistency
+        for chunk in chunks:
+            self.assertEqual(chunk.metadata["title"], "Main Document Title")
+
+        # Check section titles and anchors
+        intro_chunk = next(c for c in chunks if "Introduction" in c.text)
+        self.assertIn(intro_chunk.metadata["docs_url"], ["http://example.com/meta#intro", "http://example.com/meta"])
+        self.assertEqual(intro_chunk.metadata["section_title"], "Introduction")
 
-        topic1_chunks = [c for c in chunks if "Topic 1" in c.text or "Content 1" in c.text or "More content 1" in c.text]
+        topic1_chunks = [c for c in chunks if "Topic 1" in c.text or "Content 1" in c.text]
         self.assertTrue(all(c.metadata["docs_url"] == "http://example.com/meta#topic1" for c in topic1_chunks))
-        
-        final_thoughts_chunk = next((c for c in chunks if "Final words" in c.text), None)
-        
-        self.assertIsNotNone(final_thoughts_chunk, "Final thoughts chunk not found")
-        
+        self.assertTrue(all(c.metadata["section_title"] == "Topic 1" for c in topic1_chunks))
+
+        conclusion_chunk = next(c for c in chunks if "Conclusion text" in c.text)
+        self.assertEqual(conclusion_chunk.metadata["section_title"], "Topic 1") # Inherited from previous heading
+
+        final_thoughts_chunk = next(c for c in chunks if "Final words" in c.text)
         self.assertEqual(final_thoughts_chunk.metadata["docs_url"], "http://example.com/meta#final-thoughts")
-        self.assertEqual(final_thoughts_chunk.metadata["title"], "Intro")
+        self.assertEqual(final_thoughts_chunk.metadata["section_title"], "Final Thoughts")
 
     def test_no_anchor_found(self):
         """Tests that the source URL has no anchor if no IDs are present."""
-        html = "<html><body><h1>No Anchor Title</h1><p>Paragraph 1.</p><p>Paragraph 2.</p></body></html>"
+        html = "<html><head><title>No Anchor Title</title></head><body><p>Paragraph 1.</p><p>Paragraph 2.</p></body></html>"
         chunks = chunk_html(html, "http://example.com/no-anchor", max_token_limit=15)
         self.assertEqual(len(chunks), 2)
         self.assertEqual(chunks[0].metadata["docs_url"], "http://example.com/no-anchor")
         self.assertEqual(chunks[1].metadata["docs_url"], "http://example.com/no-anchor")
         self.assertEqual(chunks[0].metadata["title"], "No Anchor Title")
+        self.assertEqual(chunks[0].metadata["section_title"], "No Anchor Title")
 
     def test_empty_html(self):
         """Tests that empty or minimal HTML does not cause errors."""
diff --git a/scripts/html_embeddings/chunk_html.py b/scripts/html_embeddings/chunk_html.py
@@ -6,6 +6,7 @@
 import logging
 import re
 import sys
+from bs4 import BeautifulSoup
 from pathlib import Path
 from typing import Dict, List, Any, Optional
 from urllib.parse import urlparse
@@ -88,7 +89,7 @@ def chunk_html_documents(
             success, chunk_count = chunk_single_html_file(
                 input_file=html_file.resolve(),
                 output_dir=doc_specific_output_dir, # Pass the new doc-specific dir
-                input_base_dir=base_dir_for_relative_paths, # Pass the consistent version-level base path
+                input_base_dir=base_dir_for_relative_paths.resolve(), # Pass the consistent version-level base path
                 source_url=source_url,
                 max_token_limit=max_token_limit,
                 count_tag_tokens=count_tag_tokens,
@@ -182,15 +183,21 @@ def chunk_single_html_file(
 
         chunk_count = 0
         for i, chunk_obj in enumerate(chunks):
+            chunker_metadata = chunk_obj.metadata or {}
             chunk_metadata = {
                 **base_metadata,
-                **chunk_obj.metadata,
+                "docs_url": chunker_metadata.get("docs_url"),
+                "title": chunker_metadata.get("title"),
+                "section_title": chunker_metadata.get("section_title"),
                 "chunk_index": i,
                 "total_chunks": len(chunks),
                 "token_count": count_html_tokens(chunk_obj.text, count_tag_tokens),
                 "source_file": str(relative_path),
             }
 
+            # Filter out any keys that have None values to keep the JSON clean
+            chunk_metadata = {k: v for k, v in chunk_metadata.items() if v is not None}
+
             chunk_data = {
                 "id": f"{base_metadata['doc_id']}_chunk_{i:04d}",
                 "content": chunk_obj.text,
