Method _embedding_based_retrieve

src/memu/app/retrieve.py:867–941

Embedding-based retrieval with query rewriting and judging at each tier

(
        self,
        query: str,
        top_k: int,
        context_queries: list[dict[str, Any]] | None,
        ctx: Context,
        store: Database,
        llm_client: Any | None = None,
        where: Mapping[str, Any] | None = None,
    )

        return None

    async def _embedding_based_retrieve(
        self,
        query: str,
        top_k: int,
        context_queries: list[dict[str, Any]] | None,
        ctx: Context,
        store: Database,
        llm_client: Any | None = None,
        where: Mapping[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Embedding-based retrieval with query rewriting and judging at each tier"""
        where_filters = self._normalize_where(where)
        category_pool = store.memory_category_repo.list_categories(where_filters)
        items_pool = store.memory_item_repo.list_items(where_filters)
        resource_pool = store.resource_repo.list_resources(where_filters)
        client = llm_client or self._get_llm_client()
        current_query = query
        qvec = (await client.embed([current_query]))[0]
        response: dict[str, Any] = {"resources": [], "items": [], "categories": [], "next_step_query": None}
        content_sections: list[str] = []

        # Tier 1: Categories
        cat_hits, summary_lookup = await self._rank_categories_by_summary(
            qvec,
            top_k,
            ctx,
            store,
            embed_client=client,
            categories=category_pool,
        )
        if cat_hits:
            response["categories"] = self._materialize_hits(cat_hits, category_pool)
            content_sections.append(
                self._format_category_content(cat_hits, summary_lookup, store, categories=category_pool)
            )

            needs_more, current_query = await self._decide_if_retrieval_needed(
                current_query,
                context_queries,
                retrieved_content="\n\n".join(content_sections),
                llm_client=client,
            )
            response["next_step_query"] = current_query
            if not needs_more:
                return response
            # Re-embed with rewritten query
            qvec = (await client.embed([current_query]))[0]

        # Tier 2: Items
        item_hits = store.memory_item_repo.vector_search_items(qvec, top_k, where=where_filters)
        if item_hits:
            response["items"] = self._materialize_hits(item_hits, items_pool)
            content_sections.append(self._format_item_content(item_hits, store, items=items_pool))

            needs_more, current_query = await self._decide_if_retrieval_needed(
                current_query,
                context_queries,
                retrieved_content="\n\n".join(content_sections),