From d7ceecfa1879860da98d6bf444b20fc3dccdc32e Mon Sep 17 00:00:00 2001 From: Ariel Memory Date: Sat, 4 Jul 2026 05:18:34 +0300 Subject: [PATCH 1/2] refactor(rag): fix N+1 query, extract _insert_page, simplify router - Fix N+1 query in _search_rrf: batch page lookups with IN clause - Extract _insert_page helper from ingest_file/ingest_text (dedup FTS insert) - Extract _match_route helper from route() (flatten nested matching) --- rag/engine.py | 71 +++++++++++++++++++++++++++------------------------ rag/router.py | 62 +++++++++++++++++++++++--------------------- 2 files changed, 70 insertions(+), 63 deletions(-) diff --git a/rag/engine.py b/rag/engine.py index b1f07baf..25bcf1d3 100644 --- a/rag/engine.py +++ b/rag/engine.py @@ -148,19 +148,16 @@ async def _ingest_single_file(self, conn, page_id: int, content: str) -> int: ) return len(chunks) - async def ingest_file(self, filepath: Path, user_id: str = "default", wiki_type: Optional[str] = None) -> str: - content = filepath.read_text(encoding="utf-8") - file_hash = hashlib.sha256(content.encode()).hexdigest() - conn = await self._cm.get("memory.db") - - cur = await conn.execute("SELECT id FROM rag_pages WHERE sha256_hash = ? AND user_id = ?", (file_hash, user_id)) + async def _insert_page(self, conn, title: str, content: str, user_id: str, page_hash: str, wiki_type: Optional[str] = None, path: str = "") -> int | None: + """Insert a page into rag_pages + rag_fts + chunks. Returns page_id or None if duplicate.""" + cur = await conn.execute("SELECT id FROM rag_pages WHERE sha256_hash = ? AND user_id = ?", (page_hash, user_id)) existing = await cur.fetchone() if existing: - return "[SKIP] %s (already ingested)" % filepath.name + return None cursor = await conn.execute( "INSERT INTO rag_pages (layer, user_id, title, path, content, sha256_hash, wiki_type) VALUES (?, ?, ?, ?, ?, ?, ?)", - (self.layer, user_id, filepath.stem, str(filepath), content, file_hash, wiki_type), + (self.layer, user_id, title, path, content, page_hash, wiki_type), ) page_id = cursor.lastrowid @@ -168,15 +165,25 @@ async def ingest_file(self, filepath: Path, user_id: str = "default", wiki_type: try: await conn.execute( "INSERT INTO rag_fts(rowid, title, content, wiki_type) VALUES (?, ?, ?, ?)", - (page_id, filepath.stem, content, wiki_type or ""), + (page_id, title, content, wiki_type or ""), ) except Exception: pass - num_chunks = await self._ingest_single_file(conn, page_id, content) + await self._ingest_single_file(conn, page_id, content) + return page_id + + async def ingest_file(self, filepath: Path, user_id: str = "default", wiki_type: Optional[str] = None) -> str: + content = filepath.read_text(encoding="utf-8") + file_hash = hashlib.sha256(content.encode()).hexdigest() + conn = await self._cm.get("memory.db") + + page_id = await self._insert_page(conn, filepath.stem, content, user_id, file_hash, wiki_type, str(filepath)) + if page_id is None: + return "[SKIP] %s (already ingested)" % filepath.name await conn.commit() - return "[OK] %s (%d chunks)" % (filepath.name, num_chunks) + return "[OK] %s" % filepath.name async def ingest_text( self, @@ -191,27 +198,11 @@ async def ingest_text( text_hash = hashlib.sha256(text.encode()).hexdigest() conn = await self._cm.get("memory.db") - cur = await conn.execute("SELECT id FROM rag_pages WHERE sha256_hash = ? AND user_id = ?", (text_hash, user_id)) - existing = await cur.fetchone() - if existing: - return existing[0] - - cursor = await conn.execute( - "INSERT INTO rag_pages (layer, user_id, title, path, content, sha256_hash, wiki_type) VALUES (?, ?, ?, ?, ?, ?, ?)", - (self.layer, user_id, title, path, text, text_hash, wiki_type), - ) - page_id = cursor.lastrowid - - if self._fts_available: - try: - await conn.execute( - "INSERT INTO rag_fts(rowid, title, content, wiki_type) VALUES (?, ?, ?, ?)", - (page_id, title, text, wiki_type or ""), - ) - except Exception: - pass - - await self._ingest_single_file(conn, page_id, text) + page_id = await self._insert_page(conn, title, text, user_id, text_hash, wiki_type, path) + if page_id is None: + cur = await conn.execute("SELECT id FROM rag_pages WHERE sha256_hash = ? AND user_id = ?", (text_hash, user_id)) + existing = await cur.fetchone() + return existing[0] if existing else 0 if relation_to is not None: await conn.execute( @@ -398,19 +389,31 @@ def rrf(rank: int) -> float: merged[doc_id] = score sorted_ids = sorted(merged.keys(), key=lambda x: -merged[x])[:limit] + if not sorted_ids: + return [] + conn = await self._cm.get("memory.db") + placeholders = ",".join(["?"] * len(sorted_ids)) + cur = await conn.execute( + f"SELECT id, title, content, wiki_type FROM rag_pages WHERE id IN ({placeholders})", + sorted_ids, + ) + rows = await cur.fetchall() + by_id = {r[0]: r for r in rows} + results = [] for doc_id in sorted_ids: - row = await (await conn.execute("SELECT id, title, content, wiki_type FROM rag_pages WHERE id=?", (doc_id,))).fetchone() + row = by_id.get(doc_id) if row: has_fts = doc_id in fts_ranks has_bin = doc_id in bin_ranks source = "rrf(fts+mib)" if (has_fts and has_bin) else ("fts5" if has_fts else "mib") + content = row[2] results.append( { "id": row[0], "title": row[1], - "content": row[2][:500] + "..." if len(row[2]) > 500 else row[2], + "content": content[:500] + "..." if len(content) > 500 else content, "wiki_type": row[3], "score": merged[doc_id], "source": source, diff --git a/rag/router.py b/rag/router.py index 71330512..98b30098 100644 --- a/rag/router.py +++ b/rag/router.py @@ -144,42 +144,46 @@ def _flat_keywords(self, kind: str) -> list[str]: async def route(self, query: str, recent_context: Optional[list[dict]] = None) -> RouterResult: q = query.lower() - # B2.5: Data-driven route matching for route in _ROUTE_TABLE: - strategy_name = route["strategy"] - strategy = Strategy[strategy_name] + strategy = Strategy[route["strategy"]] confidence = route["confidence"] keyword_kind = route["keywords"] - if keyword_kind == "recent": - if self._is_recent_query(q) and recent_context: - return RouterResult(strategy, recent_context, confidence) - - elif keyword_kind == "wiki": - if self._is_wiki_query(q): - return await self._route_wiki(query, strategy, confidence) - - elif keyword_kind == "entity": - entities = self._extract_entities(query) - if entities: - result = await self._route_entities(entities, strategy, confidence) - if result: - return result - - elif keyword_kind == "graph": - if self._is_graph_query(q): - result = await self._route_graph(q, strategy, confidence) - if result: - return result - - elif keyword_kind is None: - # Fallback to semantic - results = await self._rag.search(query, self.user_id, strategy="hybrid", limit=3) - if results: - return RouterResult(strategy, results, confidence) + result = await self._match_route(q, query, keyword_kind, strategy, confidence, recent_context) + if result is not None: + return result return RouterResult(Strategy.SEMANTIC, [], 0.0) + async def _match_route( + self, q_lower: str, query: str, keyword_kind: Optional[str], + strategy: Strategy, confidence: float, recent_context: Optional[list[dict]], + ) -> RouterResult | None: + """Try to match a single route. Returns RouterResult or None.""" + if keyword_kind == "recent": + if self._is_recent_query(q_lower) and recent_context: + return RouterResult(strategy, recent_context, confidence) + + elif keyword_kind == "wiki": + if self._is_wiki_query(q_lower): + return await self._route_wiki(query, strategy, confidence) + + elif keyword_kind == "entity": + entities = self._extract_entities(query) + if entities: + return await self._route_entities(entities, strategy, confidence) + + elif keyword_kind == "graph": + if self._is_graph_query(q_lower): + return await self._route_graph(q_lower, strategy, confidence) + + elif keyword_kind is None: + results = await self._rag.search(query, self.user_id, strategy="hybrid", limit=3) + if results: + return RouterResult(strategy, results, confidence) + + return None + async def _route_wiki(self, query: str, strategy: Strategy, confidence: float) -> RouterResult: results = await self._rag.search(query, self.user_id, strategy="hybrid", limit=3) if results: From 4d4a9879cfae56981e786fd6d8232fa05cd369c0 Mon Sep 17 00:00:00 2001 From: Ariel Memory Date: Sat, 4 Jul 2026 05:21:41 +0300 Subject: [PATCH 2/2] fix: apply ruff format to rag/engine.py and rag/router.py --- rag/engine.py | 4 +++- rag/router.py | 9 +++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/rag/engine.py b/rag/engine.py index 25bcf1d3..1b8696d5 100644 --- a/rag/engine.py +++ b/rag/engine.py @@ -148,7 +148,9 @@ async def _ingest_single_file(self, conn, page_id: int, content: str) -> int: ) return len(chunks) - async def _insert_page(self, conn, title: str, content: str, user_id: str, page_hash: str, wiki_type: Optional[str] = None, path: str = "") -> int | None: + async def _insert_page( + self, conn, title: str, content: str, user_id: str, page_hash: str, wiki_type: Optional[str] = None, path: str = "" + ) -> int | None: """Insert a page into rag_pages + rag_fts + chunks. Returns page_id or None if duplicate.""" cur = await conn.execute("SELECT id FROM rag_pages WHERE sha256_hash = ? AND user_id = ?", (page_hash, user_id)) existing = await cur.fetchone() diff --git a/rag/router.py b/rag/router.py index 98b30098..39fc7b46 100644 --- a/rag/router.py +++ b/rag/router.py @@ -156,8 +156,13 @@ async def route(self, query: str, recent_context: Optional[list[dict]] = None) - return RouterResult(Strategy.SEMANTIC, [], 0.0) async def _match_route( - self, q_lower: str, query: str, keyword_kind: Optional[str], - strategy: Strategy, confidence: float, recent_context: Optional[list[dict]], + self, + q_lower: str, + query: str, + keyword_kind: Optional[str], + strategy: Strategy, + confidence: float, + recent_context: Optional[list[dict]], ) -> RouterResult | None: """Try to match a single route. Returns RouterResult or None.""" if keyword_kind == "recent":