Upload app.py

app.py

@@ -3,34 +3,33 @@
 from __future__ import annotations
 
 import gzip
+import hashlib
 import json
-import posixpath
 import random
 import re
+import sqlite3
 import time
 from contextlib import asynccontextmanager
-from datetime import datetime
 from pathlib import Path
 from typing import Optional
 from urllib.parse import quote
 
 from fastapi import FastAPI, Query
-from fastapi import HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
 from fastapi.responses import HTMLResponse, JSONResponse, PlainTextResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 
-
 BASE_DIR = Path(__file__).resolve().parent
 DATA_PATH = BASE_DIR / "data/search_data.json.gz"
 FOLDER_TREE_PATH = BASE_DIR / "data/folder_tree.json.gz"
 FOLDER_BROWSER_PATH = BASE_DIR / "data/folder_browser.json.gz"
-FULLTEXT_MANIFEST_PATH = BASE_DIR / "data/fulltext_manifest.json.gz"
+FULLTEXT_DIR = BASE_DIR / "data/fulltext"
 
 records: list[dict] = []
 record_map: dict[str, dict] = {}
+record_map_index: dict[str, int] = {}
 sources: list[dict] = []
 source_counts: dict[str, int] = {}
 folder_tree_data: dict[str, list[dict]] = {}
@@ -42,6 +41,72 @@ word_index: dict[str, set[int]] = {}
 did_you_mean_vocab: dict[str, int] = {}
 
 
+def query_tokens(text: str) -> list[str]:
+    return list(dict.fromkeys(re.findall(r"[a-z0-9]+|[\u4e00-\u9fff\u3400-\u4dbf]+", (text or "").lower())))
+
+
+def decode_doc_ids(payload: bytes, source: str) -> set[str]:
+    results = set()
+    number = shift = previous = 0
+    for byte in payload:
+        number |= (byte & 0x7F) << shift
+        if byte & 0x80:
+            shift += 7
+            continue
+        previous += number
+        results.add(f"{source}:{previous}")
+        number = shift = 0
+    return results
+
+
+class FulltextDatabases:
+    def __init__(self, directory: Path):
+        self.connections = {}
+        for path in directory.glob("*.sqlite3"):
+            uri = f"file:{path.resolve()}?mode=ro&immutable=1"
+            self.connections[path.stem] = sqlite3.connect(uri, uri=True, check_same_thread=False)
+
+    def close(self) -> None:
+        for connection in self.connections.values():
+            connection.close()
+
+    def search_source(self, source: str, query: str) -> set[str]:
+        connection = self.connections.get(source)
+        tokens = query_tokens(query)
+        if connection is None or not tokens:
+            return set()
+        matched = None
+        for token in tokens:
+            token_hash = hashlib.sha256(token.encode("utf-8")).digest()[:16]
+            row = connection.execute("SELECT docs FROM postings WHERE token_hash = ?", (token_hash,)).fetchone()
+            token_matches = decode_doc_ids(row[0], source) if row else set()
+            matched = token_matches if matched is None else matched & token_matches
+            if not matched:
+                return set()
+        return matched or set()
+
+    def summaries(self, doc_ids: list[str]) -> dict[str, str]:
+        grouped = {}
+        for doc_id in doc_ids:
+            grouped.setdefault(doc_id.split(":", 1)[0], []).append(doc_id)
+        summaries = {}
+        for source, source_doc_ids in grouped.items():
+            connection = self.connections.get(source)
+            if connection is None:
+                continue
+            for offset in range(0, len(source_doc_ids), 500):
+                batch = source_doc_ids[offset:offset + 500]
+                rows = connection.execute(
+                    f"SELECT doc_id, summary FROM documents WHERE doc_id IN ({','.join('?' for _ in batch)})",
+                    batch,
+                )
+                summaries.update(rows)
+        return summaries
+
+
+fulltext_databases: Optional[FulltextDatabases] = None
+
+
 def tokenize(text: str) -> list[str]:
     text_lower = (text or "").lower()
     return list(set(re.findall(r"[a-z0-9]+|[\u4e00-\u9fff\u3400-\u4dbf]+", text_lower)))
@@ -125,13 +190,14 @@ def build_indexes() -> None:
 
 
 def load_data() -> None:
-    global records, record_map, sources, source_counts, folder_tree_data, folder_browser_data
+    global records, record_map, record_map_index, sources, source_counts, folder_tree_data, folder_browser_data
     start = time.time()
     payload = load_json_gz(DATA_PATH)
     records = payload.get("records", [])
     sources = payload.get("sources", [])
     source_counts = {item["slug"]: item.get("count", 0) for item in sources}
     record_map = {rec["doc_id"]: rec for rec in records}
+    record_map_index = {rec["doc_id"]: idx for idx, rec in enumerate(records)}
     folder_tree_data = load_json_gz(FOLDER_TREE_PATH)
     folder_browser_data = load_json_gz(FOLDER_BROWSER_PATH)
     build_indexes()
@@ -184,6 +250,15 @@ def trim_record(rec: dict) -> dict:
     }
 
 
+def add_summaries(items: list[dict]) -> list[dict]:
+    if fulltext_databases is None or not items:
+        return items
+    summaries = fulltext_databases.summaries([item["doc_id"] for item in items])
+    for item in items:
+        item["snippet"] = summaries.get(item["doc_id"], "")
+    return items
+
+
 def search(q="", sources_filter=None, folders=None, min_size=None, max_size=None, page=1, page_size=100, sort="relevance", exact=False, search_paths=True):
     q = q.strip()
     if not q:
@@ -227,7 +302,33 @@ def search(q="", sources_filter=None, folders=None, min_size=None, max_size=None
 
     total = len(filtered)
     start = (page - 1) * page_size
-    result_items = [trim_record(records[idx]) for idx in filtered[start:start + page_size]]
+    result_items = add_summaries([trim_record(records[idx]) for idx in filtered[start:start + page_size]])
+    return {"results": result_items, "total": total, "page": page, "page_size": page_size, "did_you_mean": None}
+
+
+def fulltext_search(q="", sources_filter=None, folders=None, min_size=None, max_size=None, page=1, page_size=100, sort="relevance", exact=False):
+    q = q.strip()
+    if not q or fulltext_databases is None:
+        return search(q, sources_filter, folders, min_size, max_size, page, page_size, sort, exact, True)
+
+    source_slugs = sources_filter or list(source_counts)
+    matched_doc_ids = set()
+    for source_slug in source_slugs:
+        matched_doc_ids.update(fulltext_databases.search_source(source_slug, q))
+    indices = [idx for doc_id in matched_doc_ids if (idx := record_map_index.get(doc_id)) is not None]
+    filtered = apply_filters(indices, sources_filter, folders, min_size, max_size)
+
+    tokens = query_tokens(q)
+    if sort == "relevance":
+        filtered.sort(key=lambda idx: (-score_record(idx, tokens, True), records[idx]["display_rel_path"].lower()))
+    elif sort == "name":
+        filtered.sort(key=lambda idx: records[idx]["display_rel_path"].lower())
+    elif sort == "size":
+        filtered.sort(key=lambda idx: (-records[idx].get("size", 0), records[idx]["display_rel_path"].lower()))
+
+    total = len(filtered)
+    start = (page - 1) * page_size
+    result_items = add_summaries([trim_record(records[idx]) for idx in filtered[start:start + page_size]])
     return {"results": result_items, "total": total, "page": page, "page_size": page_size, "did_you_mean": None}
 
 
@@ -246,8 +347,13 @@ def get_folder_contents(source_slug: str, path: str) -> dict:
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
+    global fulltext_databases
     load_data()
-    yield
+    fulltext_databases = FulltextDatabases(FULLTEXT_DIR)
+    try:
+        yield
+    finally:
+        fulltext_databases.close()
 
 
 app = FastAPI(title="VOMEBOOK Search", version="1.0", lifespan=lifespan)
@@ -266,22 +372,30 @@ class SearchRequest(BaseModel):
     sort: str = "relevance"
     exact: bool = False
     search_paths: bool = True
+    fulltext: bool = False
 
 
 class ZipRequest(BaseModel):
     doc_ids: list[str] = []
 
 
+def run_search(body: SearchRequest, sources_filter=None):
+    selected_sources = sources_filter if sources_filter is not None else body.sources
+    if body.fulltext:
+        return fulltext_search(body.q, selected_sources, body.folders, body.min_size, body.max_size, body.page, body.page_size, body.sort, body.exact)
+    return search(body.q, selected_sources, body.folders, body.min_size, body.max_size, body.page, body.page_size, body.sort, body.exact, body.search_paths)
+
+
 @app.post("/api/search")
 def api_search(body: SearchRequest):
-    return JSONResponse(search(body.q, body.sources, body.folders, body.min_size, body.max_size, body.page, body.page_size, body.sort, body.exact, body.search_paths))
+    return JSONResponse(run_search(body))
 
 
 @app.post("/api/search/{source_slug}")
 def api_search_source(source_slug: str, body: SearchRequest):
     if source_slug not in source_counts:
         return JSONResponse({"error": "source not found", "results": [], "total": 0}, status_code=404)
-    return JSONResponse(search(body.q, [source_slug], body.folders, body.min_size, body.max_size, body.page, body.page_size, body.sort, body.exact, body.search_paths))
+    return JSONResponse(run_search(body, [source_slug]))
 
 
 @app.get("/api/sources")
@@ -397,12 +511,6 @@ def api_zip(req: ZipRequest):
         return JSONResponse({"error": f"zip failed: {exc}"}, status_code=500)
 
 
-@app.get("/api/fulltext-manifest")
-def api_fulltext_manifest():
-    return JSONResponse(load_json_gz(FULLTEXT_MANIFEST_PATH))
-
-
-app.mount("/data", StaticFiles(directory=str(BASE_DIR / "data")), name="data")
 app.mount("/static", StaticFiles(directory=str(BASE_DIR / "static"), html=True), name="static")
 app.mount("/icons", StaticFiles(directory=str(BASE_DIR / "static/icons")), name="icons")