diff --git a/.github/workflows/backend-tests.yml b/.github/workflows/backend-tests.yml index 9d4f8c5..8a3cfe5 100644 --- a/.github/workflows/backend-tests.yml +++ b/.github/workflows/backend-tests.yml @@ -10,7 +10,7 @@ permissions: contents: read jobs: - chunking-tests: + unit-tests: runs-on: ubuntu-latest defaults: run: @@ -28,5 +28,5 @@ jobs: - name: Install test dependency run: python -m pip install pytest - - name: Run chunking tests - run: python -m pytest tests/test_chunking.py -v + - name: Run backend unit tests + run: python -m pytest tests/test_chunking.py tests/test_retrieval.py -v diff --git a/backend/app/main.py b/backend/app/main.py index d7ff9e8..afd6c6d 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -8,6 +8,7 @@ from app.models.page_chunk_db import PageChunk from app.services.chunking import chunk_text from app.services.embeddings import generate_embedding +from app.services.retrieval import format_search_results app = FastAPI() @@ -82,27 +83,21 @@ def create_page(page: PageCreate, db: Session = Depends(get_db)): @app.get("/search") def search_pages(query: str, db: Session = Depends(get_db)): query_embedding = generate_embedding(query) - distance = PageChunk.embedding.cosine_distance(query_embedding) - results = ( + # Fetch more chunk candidates than we display because several of the best + # chunks can belong to the same source page. format_search_results keeps only + # the strongest chunk from each page. + candidate_limit = 50 + result_limit = 5 + + rows = ( db.query(PageChunk, Page, distance.label("distance")) .join(Page, PageChunk.page_id == Page.id) + .filter(PageChunk.embedding.is_not(None)) .order_by(distance) - .limit(5) + .limit(candidate_limit) .all() ) - return [ - { - "page_id": page.id, - "chunk_id": chunk.id, - "chunk_index": chunk.chunk_index, - "url": page.url, - "title": page.title, - "content": chunk.content, - "created_at": page.created_at, - "similarity": 1 - float(distance_value), - } - for chunk, page, distance_value in results - ] \ No newline at end of file + return format_search_results(rows, limit=result_limit) diff --git a/backend/app/services/retrieval.py b/backend/app/services/retrieval.py new file mode 100644 index 0000000..9b58b46 --- /dev/null +++ b/backend/app/services/retrieval.py @@ -0,0 +1,59 @@ +import re +from typing import Any, Iterable + + +_WHITESPACE = re.compile(r"\s+") + + +def clean_excerpt(text: str, max_length: int = 360) -> str: + """Return a compact, readable excerpt without cutting a word in half.""" + cleaned = _WHITESPACE.sub(" ", text).strip() + + if len(cleaned) <= max_length: + return cleaned + + cutoff = cleaned.rfind(" ", 0, max_length + 1) + if cutoff <= 0: + cutoff = max_length + + return cleaned[:cutoff].rstrip(" ,;:-") + "…" + + +def format_search_results( + rows: Iterable[tuple[Any, Any, float]], + limit: int = 5, + excerpt_length: int = 360, +) -> list[dict[str, Any]]: + """Keep the strongest chunk per page and expose a stable search-result shape. + + Rows are expected in ascending vector-distance order. The first occurrence of + a page is therefore its strongest chunk. This prevents one long source from + occupying every visible result while retaining chunk-level provenance. + """ + results: list[dict[str, Any]] = [] + seen_page_ids: set[int] = set() + + for chunk, page, distance_value in rows: + if page.id in seen_page_ids: + continue + + seen_page_ids.add(page.id) + similarity = max(-1.0, min(1.0, 1 - float(distance_value))) + + results.append( + { + "page_id": page.id, + "matched_chunk_id": chunk.id, + "matched_chunk_index": chunk.chunk_index, + "url": page.url, + "title": page.title, + "excerpt": clean_excerpt(chunk.content, excerpt_length), + "created_at": page.created_at, + "similarity": similarity, + } + ) + + if len(results) >= limit: + break + + return results diff --git a/backend/tests/test_retrieval.py b/backend/tests/test_retrieval.py new file mode 100644 index 0000000..6e62f9e --- /dev/null +++ b/backend/tests/test_retrieval.py @@ -0,0 +1,93 @@ +from datetime import datetime, timezone +from types import SimpleNamespace + +from app.services.retrieval import clean_excerpt, format_search_results + + +def make_page(page_id: int, title: str = "Source"): + return SimpleNamespace( + id=page_id, + url=f"https://example.com/{page_id}", + title=title, + created_at=datetime(2026, 10, 5, tzinfo=timezone.utc), + ) + + +def make_chunk(chunk_id: int, index: int, content: str): + return SimpleNamespace( + id=chunk_id, + chunk_index=index, + content=content, + ) + + +def test_clean_excerpt_collapses_whitespace(): + text = "A useful sentence.\n\n Another useful sentence." + + assert clean_excerpt(text) == "A useful sentence. Another useful sentence." + + +def test_clean_excerpt_truncates_on_word_boundary(): + excerpt = clean_excerpt("alpha beta gamma delta epsilon", max_length=18) + + assert excerpt == "alpha beta gamma…" + assert len(excerpt) <= 19 + + +def test_results_keep_only_best_chunk_per_page(): + page_one = make_page(1, "First source") + page_two = make_page(2, "Second source") + + rows = [ + (make_chunk(10, 2, "Best chunk from first source."), page_one, 0.10), + (make_chunk(11, 4, "Another chunk from first source."), page_one, 0.15), + (make_chunk(20, 0, "Best chunk from second source."), page_two, 0.20), + ] + + results = format_search_results(rows) + + assert len(results) == 2 + assert results[0]["page_id"] == 1 + assert results[0]["matched_chunk_id"] == 10 + assert results[0]["matched_chunk_index"] == 2 + assert results[1]["page_id"] == 2 + + +def test_results_expose_excerpt_instead_of_raw_content(): + page = make_page(1) + chunk = make_chunk(10, 0, "word " * 100) + + result = format_search_results([(chunk, page, 0.25)], excerpt_length=40)[0] + + assert "excerpt" in result + assert "content" not in result + assert result["excerpt"].endswith("…") + assert result["similarity"] == 0.75 + + +def test_result_limit_counts_unique_pages(): + rows = [] + for page_id in range(1, 8): + rows.append( + ( + make_chunk(page_id, 0, f"Chunk for page {page_id}"), + make_page(page_id), + page_id / 100, + ) + ) + + results = format_search_results(rows, limit=5) + + assert len(results) == 5 + assert [result["page_id"] for result in results] == [1, 2, 3, 4, 5] + + +def test_similarity_is_clamped_to_cosine_range(): + page = make_page(1) + chunk = make_chunk(1, 0, "content") + + high = format_search_results([(chunk, page, -1.0)])[0] + low = format_search_results([(chunk, page, 3.0)])[0] + + assert high["similarity"] == 1.0 + assert low["similarity"] == -1.0 diff --git a/frontend/app.js b/frontend/app.js index 147b0d4..68c0a4d 100644 --- a/frontend/app.js +++ b/frontend/app.js @@ -45,12 +45,11 @@ searchForm.addEventListener("submit", async (event) => { const meta = document.createElement("div"); meta.className = "result-meta"; - meta.textContent = - `Similarity: ${result.similarity.toFixed(3)} · Chunk ${result.chunk_index}`; + meta.textContent = `Similarity: ${result.similarity.toFixed(3)}`; const content = document.createElement("p"); content.className = "result-content"; - content.textContent = result.content; + content.textContent = result.excerpt; const link = document.createElement("a"); link.href = result.url; @@ -66,4 +65,4 @@ searchForm.addEventListener("submit", async (event) => { statusText.textContent = "Search failed. Make sure the Synapse backend is running."; } -}); \ No newline at end of file +});