Spaces:
Running
Running
| import asyncio | |
| from typing import Any, Optional | |
| from fastapi import HTTPException | |
| from src.config import config | |
| from src.lib.embeddings.profile import EmbeddingProfileError | |
| from src.core.search.catalog_graph import node_id_from_passage_text | |
| from src.core.search.personalize import personalize_ranked_ids | |
| from src.core.search.graph_channels import ( | |
| COMMUNITIES_CHANNEL, | |
| ENTITIES_CHANNEL, | |
| EVENTS_CHANNEL, | |
| GraphHit, | |
| NEIGHBORS_CHANNEL, | |
| PASSAGES_CHANNEL, | |
| collect_community_hits, | |
| collect_entity_hits, | |
| collect_event_hits, | |
| expand_neighbor_hits, | |
| is_item_entity, | |
| selected_graph_channels, | |
| ) | |
| from src.core.search.hooks import ( | |
| SearchPluginError, | |
| rerank_max_k_for_mode, | |
| resolve_reranker, | |
| resolve_retrievers, | |
| retrieve_k_for_mode, | |
| ) | |
| from src.core.search.hybrid import ( | |
| collect_bm25_passages, | |
| collect_dense_passages, | |
| collect_literal_residual, | |
| dense_similarity, | |
| extras_from_metadata, | |
| facet_counts_from_extras, | |
| frozen_head_merge, | |
| fuse_passage_lists, | |
| hit_matches_extras, | |
| merge_hit_extras, | |
| passage_snippet, | |
| ) | |
| from src.lib.tracing.profiler import profile_request, profile_stage | |
| from src.services.api.constants.requests import ( | |
| SearchHit, | |
| SearchHitScores, | |
| SearchRequestBody, | |
| SearchResponse, | |
| ) | |
| from src.services.data.main import data_adapter | |
| from src.services.kg_agent.main import embeddings_adapter, graph_adapter, vector_store_adapter | |
| from src.utils.vector_search import VectorSearchFacade | |
| vector_search = VectorSearchFacade(vector_store_adapter) | |
| def _require_search_enabled() -> None: | |
| if not config.search_enabled: | |
| raise HTTPException( | |
| status_code=404, | |
| detail="Search is disabled. Set SEARCH_ENABLED=true to enable /retrieve/search.", | |
| ) | |
| def _plugin_http(exc: SearchPluginError) -> HTTPException: | |
| return HTTPException(status_code=400, detail=str(exc)) | |
| def _fetch_snippets(ids: list[str], known: dict[str, str], brain_id: str) -> dict[str, str]: | |
| texts, _ = _fetch_chunk_fields(ids, known, {}, brain_id) | |
| return texts | |
| def _fetch_chunk_fields( | |
| ids: list[str], | |
| known_texts: dict[str, str], | |
| known_extras: dict[str, dict[str, Any]], | |
| brain_id: str, | |
| ) -> tuple[dict[str, str], dict[str, dict[str, Any]]]: | |
| texts = dict(known_texts) | |
| extras: dict[str, dict[str, Any]] = dict(known_extras) | |
| missing = [ | |
| chunk_id | |
| for chunk_id in ids | |
| if chunk_id not in texts or chunk_id not in extras | |
| ] | |
| if not missing: | |
| return texts, extras | |
| try: | |
| chunks, _ = data_adapter.get_text_chunks_by_ids(missing, False, brain_id) | |
| except Exception: | |
| return texts, extras | |
| if not isinstance(chunks, (list, tuple)): | |
| return texts, extras | |
| for chunk in chunks: | |
| chunk_id = str(getattr(chunk, "id", "") or "") | |
| if not chunk_id: | |
| continue | |
| if chunk_id not in texts: | |
| texts[chunk_id] = getattr(chunk, "text", "") or "" | |
| parsed = extras_from_metadata(getattr(chunk, "metadata", None)) | |
| if parsed: | |
| extras[chunk_id] = merge_hit_extras(extras.get(chunk_id), parsed) or parsed | |
| elif chunk_id not in extras: | |
| extras[chunk_id] = {} | |
| return texts, extras | |
| def _run_dense( | |
| query_vector: list[float], | |
| brain_id: str, | |
| k: int, | |
| extras_out: dict[str, dict[str, str]] | None = None, | |
| ): | |
| with profile_stage("search.dense"): | |
| return collect_dense_passages( | |
| vector_search, query_vector, brain_id, k, extras_out=extras_out | |
| ) | |
| def _run_bm25( | |
| query: str, | |
| brain_id: str, | |
| k: int, | |
| extras_out: dict[str, dict[str, str]] | None = None, | |
| ): | |
| with profile_stage("search.bm25"): | |
| return collect_bm25_passages( | |
| data_adapter, query, brain_id, k, extras_out=extras_out | |
| ) | |
| def _want_passages(channels: list[str] | None) -> bool: | |
| if not channels: | |
| return True | |
| return any(str(item).strip().lower() == PASSAGES_CHANNEL for item in channels) | |
| def _collect_graph_hits( | |
| request: SearchRequestBody, | |
| *, | |
| query_vector: list[float] | None, | |
| graph_channels: list[str], | |
| ) -> list[GraphHit]: | |
| k = request.k | |
| brain_id = request.brain_id | |
| fanout = int(getattr(config, "search_neighbor_fanout", 50) or 50) | |
| community_labels = request.community_labels | |
| if community_labels is None: | |
| community_labels = list(config.search_community_labels) | |
| hits: list[GraphHit] = [] | |
| if ENTITIES_CHANNEL in graph_channels: | |
| hits.extend( | |
| collect_entity_hits( | |
| query=request.query, | |
| brain_id=brain_id, | |
| k=k, | |
| graph=graph_adapter, | |
| vector_search=vector_search, | |
| query_vector=query_vector, | |
| node_labels=request.node_labels, | |
| ) | |
| ) | |
| if EVENTS_CHANNEL in graph_channels: | |
| hits.extend( | |
| collect_event_hits( | |
| query=request.query, | |
| brain_id=brain_id, | |
| k=k, | |
| graph=graph_adapter, | |
| vector_search=vector_search, | |
| query_vector=query_vector, | |
| ) | |
| ) | |
| if COMMUNITIES_CHANNEL in graph_channels: | |
| hits.extend( | |
| collect_community_hits( | |
| query=request.query, | |
| brain_id=brain_id, | |
| k=k, | |
| graph=graph_adapter, | |
| vector_search=vector_search, | |
| query_vector=query_vector, | |
| community_labels=community_labels, | |
| fanout=fanout, | |
| ) | |
| ) | |
| if request.expand == "neighbors" and hits: | |
| hits.extend( | |
| expand_neighbor_hits( | |
| hits, | |
| brain_id=brain_id, | |
| k=k, | |
| graph=graph_adapter, | |
| community_labels=community_labels, | |
| fanout=fanout, | |
| ) | |
| ) | |
| return hits | |
| async def search(request: SearchRequestBody) -> SearchResponse: | |
| _require_search_enabled() | |
| with profile_request( | |
| "retrieve.search", enabled=request.profile_stages | |
| ) as profiler: | |
| response = await _search(request) | |
| if profiler is not None: | |
| response.stage_timings = profiler.last_report | |
| return response | |
| async def _search(request: SearchRequestBody) -> SearchResponse: | |
| try: | |
| reranker = resolve_reranker(request.rerank) | |
| plugin_retrievers = resolve_retrievers(request.channels) | |
| except SearchPluginError as exc: | |
| raise _plugin_http(exc) from exc | |
| k = request.k | |
| k_ret = retrieve_k_for_mode(request.mode, k) | |
| rerank_max = rerank_max_k_for_mode(request.mode) | |
| brain_id = request.brain_id | |
| fusion = request.fusion or config.search_fusion | |
| fusion_alpha = ( | |
| request.fusion_alpha | |
| if request.fusion_alpha is not None | |
| else config.search_fusion_alpha | |
| ) | |
| want_passages = _want_passages(request.channels) | |
| graph_channels = selected_graph_channels(request.channels) | |
| want_graph = bool(graph_channels) | |
| use_dense = want_passages and config.search_use_dense | |
| use_bm25 = want_passages and config.search_use_bm25 | |
| query_vector: Optional[list[float]] = None | |
| if use_dense or want_graph: | |
| with profile_stage("embed.query", blocking=False): | |
| embedded = await asyncio.to_thread( | |
| embeddings_adapter.embed_query, request.query | |
| ) | |
| query_vector = list(embedded.embeddings or []) | |
| dense_extras: dict[str, dict[str, str]] = {} | |
| bm25_extras: dict[str, dict[str, str]] = {} | |
| retrieve_jobs = [] | |
| if use_dense: | |
| retrieve_jobs.append( | |
| asyncio.to_thread( | |
| _run_dense, query_vector or [], brain_id, k_ret, dense_extras | |
| ) | |
| ) | |
| if use_bm25: | |
| retrieve_jobs.append( | |
| asyncio.to_thread( | |
| _run_bm25, request.query, brain_id, k_ret, bm25_extras | |
| ) | |
| ) | |
| dense_ids: list[str] = [] | |
| dense_distances: dict[str, float] = {} | |
| bm25_ids: list[str] = [] | |
| bm25_scores: dict[str, float] = {} | |
| texts: dict[str, str] = {} | |
| plugin_id_lists: list[list[str]] = [] | |
| plugin_scores: dict[str, dict[str, float]] = {} | |
| plugin_channels: dict[str, str] = {} | |
| graph_hits: list[GraphHit] = [] | |
| literal_ids: list[str] = [] | |
| with profile_stage("search.retrieve", blocking=False): | |
| gathered = await asyncio.gather(*retrieve_jobs) if retrieve_jobs else [] | |
| offset = 0 | |
| if use_dense: | |
| dense_ids, dense_distances = gathered[offset] | |
| offset += 1 | |
| if use_bm25: | |
| bm25_ids, bm25_scores, texts = gathered[offset] | |
| for name, retriever in plugin_retrievers: | |
| with profile_stage(f"search.plugin.{name}"): | |
| ids, scores, plugin_texts = await asyncio.to_thread( | |
| retriever, request.query, brain_id, k_ret | |
| ) | |
| plugin_id_lists.append(ids) | |
| plugin_scores[name] = scores | |
| texts.update(plugin_texts) | |
| for item_id in ids: | |
| plugin_channels.setdefault(str(item_id), f"plugin:{name}") | |
| if want_graph: | |
| try: | |
| graph_hits = await asyncio.to_thread( | |
| _collect_graph_hits, | |
| request, | |
| query_vector=query_vector, | |
| graph_channels=graph_channels, | |
| ) | |
| except EmbeddingProfileError: | |
| raise | |
| except Exception: | |
| graph_hits = [] | |
| if getattr(config, "search_literal_fill", False): | |
| lit_ids, literal_texts = collect_literal_residual( | |
| data_adapter, request.query, brain_id, k_ret | |
| ) | |
| literal_ids = list(lit_ids) | |
| texts.update(literal_texts) | |
| for item_id in literal_ids: | |
| plugin_channels.setdefault(str(item_id), PASSAGES_CHANNEL) | |
| passage_extras = {**dense_extras, **bm25_extras} | |
| dense_sims = { | |
| chunk_id: dense_similarity(distance) | |
| for chunk_id, distance in dense_distances.items() | |
| } | |
| extra_id_lists: list[list[str]] = [] | |
| graph_by_channel: dict[str, list[str]] = {} | |
| graph_by_id: dict[str, GraphHit] = {} | |
| for hit in graph_hits: | |
| graph_by_channel.setdefault(hit.channel, []).append(hit.id) | |
| current = graph_by_id.get(hit.id) | |
| if current is None or hit.score > current.score: | |
| graph_by_id[hit.id] = hit | |
| def _item_ids(channel_name: str) -> list[str]: | |
| ids = graph_by_channel.get(channel_name) or [] | |
| if channel_name != EVENTS_CHANNEL: | |
| ids = [ | |
| item_id | |
| for item_id in ids | |
| if is_item_entity( | |
| item_id, | |
| graph_by_id[item_id].labels if item_id in graph_by_id else [], | |
| ) | |
| ] | |
| return ids | |
| graph_channel_ids = { | |
| ENTITIES_CHANNEL: _item_ids(ENTITIES_CHANNEL), | |
| EVENTS_CHANNEL: _item_ids(EVENTS_CHANNEL), | |
| COMMUNITIES_CHANNEL: _item_ids(COMMUNITIES_CHANNEL), | |
| NEIGHBORS_CHANNEL: _item_ids(NEIGHBORS_CHANNEL), | |
| } | |
| for ids in graph_channel_ids.values(): | |
| if ids: | |
| extra_id_lists.append(ids) | |
| for item_id in ids: | |
| texts.setdefault(item_id, graph_by_id[item_id].snippet) | |
| fused = fuse_passage_lists( | |
| dense_ids if use_dense else [], | |
| bm25_ids if use_bm25 else [], | |
| fusion=fusion, | |
| alpha=fusion_alpha, | |
| dense_similarities=dense_sims, | |
| bm25_scores=bm25_scores, | |
| extra_id_lists=extra_id_lists, | |
| ) | |
| ranked_ids = [item for item, _ in fused] | |
| sidecar_lists = [ids for ids in plugin_id_lists if ids] | |
| if literal_ids: | |
| sidecar_lists.append(literal_ids) | |
| if sidecar_lists: | |
| ranked_ids = frozen_head_merge( | |
| ranked_ids, | |
| sidecar_lists, | |
| head_k=10, | |
| k=k_ret, | |
| ) | |
| else: | |
| ranked_ids = ranked_ids[:k_ret] | |
| rerank_scores: dict[str, float] = {} | |
| if reranker is not None and ranked_ids: | |
| head_n = min(rerank_max, len(ranked_ids)) | |
| head_ids = ranked_ids[:head_n] | |
| tail_ids = ranked_ids[head_n:] | |
| with profile_stage("search.rerank", blocking=False, chunks=head_n): | |
| texts = await asyncio.to_thread( | |
| _fetch_snippets, head_ids, texts, brain_id | |
| ) | |
| candidates = [ | |
| { | |
| "id": chunk_id, | |
| "text": texts.get(chunk_id) or "", | |
| "score": float(dict(fused).get(chunk_id, 0.0)), | |
| } | |
| for chunk_id in head_ids | |
| ] | |
| reranked = await asyncio.to_thread( | |
| reranker, request.query, candidates, head_n | |
| ) | |
| ordered: list[str] = [] | |
| seen: set[str] = set() | |
| for item in reranked or []: | |
| chunk_id = str((item or {}).get("id") or "") | |
| if not chunk_id or chunk_id not in head_ids or chunk_id in seen: | |
| continue | |
| ordered.append(chunk_id) | |
| seen.add(chunk_id) | |
| if item.get("score") is not None: | |
| rerank_scores[chunk_id] = float(item["score"]) | |
| for chunk_id in head_ids: | |
| if chunk_id not in seen: | |
| ordered.append(chunk_id) | |
| ranked_ids = ordered + tail_ids | |
| ranked_ids = ranked_ids[:k_ret] | |
| fused_map = dict(fused) | |
| extras_by_id: dict[str, dict[str, Any]] = dict(passage_extras) | |
| for item_id, graph_hit in graph_by_id.items(): | |
| extras_by_id[item_id] = ( | |
| merge_hit_extras(extras_by_id.get(item_id), graph_hit.extras) or {} | |
| ) | |
| with profile_stage("search.snippets", blocking=False, chunks=len(ranked_ids)): | |
| texts, extras_by_id = await asyncio.to_thread( | |
| _fetch_chunk_fields, ranked_ids, texts, extras_by_id, brain_id | |
| ) | |
| if request.extras: | |
| ranked_ids = [ | |
| item_id | |
| for item_id in ranked_ids | |
| if hit_matches_extras(extras_by_id.get(item_id), request.extras) | |
| ] | |
| node_id_by_hit: dict[str, str | None] = {} | |
| for item_id in ranked_ids: | |
| graph_hit = graph_by_id.get(item_id) | |
| if graph_hit is not None and is_item_entity(graph_hit.id, graph_hit.labels): | |
| node_id_by_hit[item_id] = graph_hit.id | |
| else: | |
| node_id_by_hit[item_id] = node_id_from_passage_text( | |
| texts.get(item_id) or "" | |
| ) | |
| personalize_scores: dict[str, float] = {} | |
| if request.target: | |
| retrieve_for_blend = { | |
| item_id: float( | |
| rerank_scores.get(item_id, fused_map.get(item_id, 0.0)) | |
| ) | |
| for item_id in ranked_ids | |
| } | |
| ranked_ids, personalize_scores = personalize_ranked_ids( | |
| query=request.query, | |
| ranked_ids=ranked_ids, | |
| retrieve_scores=retrieve_for_blend, | |
| node_id_by_hit=node_id_by_hit, | |
| target=request.target, | |
| brain_id=brain_id, | |
| ) | |
| ranked_ids = ranked_ids[:k] | |
| channel_of: dict[str, str] = {} | |
| for chunk_id in list(dense_ids) + list(bm25_ids): | |
| channel_of.setdefault(chunk_id, PASSAGES_CHANNEL) | |
| for hit in graph_hits: | |
| channel_of.setdefault(hit.id, hit.channel) | |
| for chunk_id, plugin_channel in plugin_channels.items(): | |
| channel_of.setdefault(chunk_id, plugin_channel) | |
| hits: list[SearchHit] = [] | |
| for chunk_id in ranked_ids: | |
| body = texts.get(chunk_id) or "" | |
| rrf = fused_map.get(chunk_id) | |
| plugin = { | |
| name: scores[chunk_id] | |
| for name, scores in plugin_scores.items() | |
| if chunk_id in scores | |
| } | |
| graph_hit = graph_by_id.get(chunk_id) | |
| extras = merge_hit_extras(extras_by_id.get(chunk_id), graph_hit.extras if graph_hit else None) | |
| scores = SearchHitScores( | |
| bm25=bm25_scores.get(chunk_id), | |
| dense=dense_sims.get(chunk_id), | |
| rrf=rrf if fusion == "rrf" else None, | |
| cc=rrf if fusion == "cc" else None, | |
| rerank=rerank_scores.get(chunk_id), | |
| plugin=plugin or None, | |
| graph=graph_hit.score if graph_hit is not None else None, | |
| personalize=personalize_scores.get(chunk_id), | |
| ) | |
| node_id = node_id_by_hit.get(chunk_id) | |
| hits.append( | |
| SearchHit( | |
| id=chunk_id, | |
| channel=channel_of.get(chunk_id, PASSAGES_CHANNEL), | |
| score=float( | |
| personalize_scores.get( | |
| chunk_id, | |
| rerank_scores.get(chunk_id, fused_map.get(chunk_id, 0.0)), | |
| ) | |
| ), | |
| scores=scores, | |
| snippet=passage_snippet(body), | |
| labels=list(graph_hit.labels) if graph_hit is not None else [], | |
| extras=extras, | |
| node_id=node_id, | |
| ) | |
| ) | |
| node_ids: list[str] = [] | |
| seen_nodes: set[str] = set() | |
| for hit in hits: | |
| if not hit.node_id or hit.node_id in seen_nodes: | |
| continue | |
| seen_nodes.add(hit.node_id) | |
| node_ids.append(hit.node_id) | |
| return SearchResponse( | |
| hits=hits, | |
| facets=facet_counts_from_extras([hit.extras for hit in hits]), | |
| channel_lists={ | |
| "dense": list(dense_ids), | |
| "bm25": list(bm25_ids), | |
| "entities": list(graph_channel_ids[ENTITIES_CHANNEL]), | |
| "events": list(graph_channel_ids[EVENTS_CHANNEL]), | |
| "communities": list(graph_channel_ids[COMMUNITIES_CHANNEL]), | |
| "neighbors": list(graph_channel_ids[NEIGHBORS_CHANNEL]), | |
| "literal": list(literal_ids), | |
| }, | |
| node_ids=node_ids, | |
| ) | |