""" File: /retrieve.py Created Date: Sunday October 26th 2025 Author: Christian Nonis ----- Last Modified: Thursday January 29th 2026 8:43:59 pm Modified By: Christian Nonis ----- """ import asyncio import json import re from collections import defaultdict from datetime import datetime from typing import Any, Optional from fastapi import HTTPException from starlette.responses import JSONResponse from src.constants.embeddings import Vector from src.constants.kg import IdentificationParams, Node, Predicate from src.core.search.entities import search_entities from src.core.search.fact_filter import ( filter_relevant_facts, personalized_pagerank, reciprocal_rank_fusion, ) from src.core.search.hybrid import ( collect_bm25_passages, collect_dense_passages, collect_ilike_passages, fuse_passage_lists, ) from src.core.search.relationships import search_relationships from src.lib.tracing.profiler import profile_request, profile_stage from src.utils.vector_search import VectorSearchFacade from src.services.api.constants.requests import ( GetContextRequestBody, GetContextResponse, GetContextTriple, RetrieveRequestResponse, RetrieveNeighborsRequestResponse, RetrievedNeighborNode, ) from src.services.kg_agent.main import graph_adapter, kg_agent from src.services.data.main import data_adapter from src.services.kg_agent.main import embeddings_adapter, vector_store_adapter from src.utils.dates import parse_date_string, to_naive_utc from src.utils.similarity.vectors import cosine_similarity from src.utils.nlp.ner import _entity_extractor from src.config import config vector_search = VectorSearchFacade(vector_store_adapter) _MAX_NOUN_CHUNKS = 5 _SEED_K = 25 _PASSAGE_K = 24 _PPR_DAMPING = 0.85 _PPR_ITERS = 20 _SESSION_RE = re.compile(r"session_(\d+)", re.IGNORECASE) _BATCH_TURN_RE = re.compile(r"\bb(\d+)_t(\d+)\b", re.IGNORECASE) _HISTORY_MODE_RE = re.compile( r"\b(" r"order(?:ing|ed)?|chronolog|sequence|walk me through|" r"contradict|conflict(?:ing)?|inconsisten|" r"previously|originally|before (?:the |that )?(?:update|change|extension)|" r"first (?:said|stated|mentioned|sprint|deadline)|" r"how many (?:weeks|days)|" r"between .+ and |" r"have i|did i" r")\b", re.I, ) _CURRENT_TRUTH_RE = re.compile( r"\b(" r"now|currently|latest|updated (?:to|value)|after the (?:update|change)|" r"how many commits|average response time|response time of" r")\b", re.I, ) _SPINE_ACTOR = ("MADE", "INITIATED", "PERFORMED", "EXPERIENCED", "COVERED") _SPINE_TARGET = ("TARGETED", "AFFECTED", "RESULTED") _CONTEXT_REL = ("OCCURRED", "WITHIN") _DIVERSIFY_LAMBDA = 0.65 _BRIDGE_SCORE_PENALTY = 0.15 _BRIDGE_SEED_HUB_CAP = 12 _BRIDGE_RESERVE_HUB_CAP = 6 _BRIDGE_RESERVE_MIN = 3 _BRIDGE_RESERVE_FRAC_DENOM = 4 _DISTANCE_QUANT_DECIMALS = 3 def _quantize_distance(distance: float) -> float: try: value = float(distance) except (TypeError, ValueError): return float("inf") if value != value or value == float("inf") or value == float("-inf"): return value return round(value, _DISTANCE_QUANT_DECIMALS) def _candidate_event_uuid(candidate: dict[str, Any]) -> str: return _event_hub_id(candidate) def _best_hub_fact_texts(candidates: list[dict[str, Any]]) -> dict[str, str]: best: dict[str, tuple[float, str]] = {} for candidate in candidates: hub = _candidate_event_uuid(candidate) if not hub: continue text = str(candidate.get("text") or "").strip() if not text: continue score = _candidate_distance(candidate) prev = best.get(hub) if prev is None or score < prev[0] or (score == prev[0] and text < prev[1]): best[hub] = (score, text) return {hub: text for hub, (_score, text) in best.items()} def _paths_for_curated( paths: list[dict[str, Any]], curated: list[dict[str, Any]], ) -> list[dict[str, Any]]: """Keep bridge paths whose hubs both survived curation; attach fact legs.""" hub_texts = _best_hub_fact_texts(curated) if not hub_texts or not paths: return [] curated_hubs = set(hub_texts) out: list[dict[str, Any]] = [] for path in paths: hubs = [str(h) for h in (path.get("hubs") or []) if h] if len(hubs) < 2 or not all(hub in curated_hubs for hub in hubs): continue out.append( { "hubs": hubs, "shared_entity": path.get("shared_entity"), "shared_entity_name": path.get("shared_entity_name"), "weight": path.get("weight"), "legs": [hub_texts[hub] for hub in hubs], } ) out.sort( key=lambda p: ( tuple(p["hubs"]), str(p.get("shared_entity") or ""), ) ) return out def _candidate_distance(candidate: dict[str, Any]) -> float: raw = candidate.get("score") if raw is None: return float("inf") try: return _quantize_distance(float(raw)) except (TypeError, ValueError): return float("inf") def _rank_bridge_seed_hubs( candidates: list[dict[str, Any]], *, cap: int, ) -> list[str]: hub_best_score: dict[str, float] = {} for candidate in candidates: hub = _candidate_event_uuid(candidate) if not hub: continue score = _candidate_distance(candidate) prev = hub_best_score.get(hub) if prev is None or score < prev: hub_best_score[hub] = score return sorted( hub_best_score.keys(), key=lambda hub: (hub_best_score[hub], hub), )[: max(0, cap)] def _resolve_chunk_sessions( chunk_ids: list[str], brain_id: str ) -> dict[str, list[str]]: unique: list[str] = [] seen: set[str] = set() for chunk_id in chunk_ids: cid = str(chunk_id or "").strip() if cid and cid not in seen: seen.add(cid) unique.append(cid) if not unique: return {} out: dict[str, list[str]] = {} try: chunks, _ = data_adapter.get_text_chunks_by_ids( unique[:120], False, brain_id ) except Exception: return {} for chunk in chunks: chunk_id = str(getattr(chunk, "id", "") or "") body = (getattr(chunk, "text", None) or "").strip() if not chunk_id or not body: continue sessions = _session_ids_from_text(body) if sessions: out[chunk_id] = sessions return out def _sessions_from_chunk_map( chunk_ids: list[str], chunk_sessions: dict[str, list[str]] ) -> list[str]: found: list[str] = [] seen: set[str] = set() for chunk_id in chunk_ids: for sid in chunk_sessions.get(str(chunk_id), []) or []: if sid not in seen: seen.add(sid) found.append(sid) return found def _expand_cross_event_bridges( candidates: list[dict[str, Any]], brain_id: str, *, max_per_hub: int, include_history: bool = False, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: from src.core.saving.hub_bridges import ( HubBridge, select_bridge_neighbors, ) if max_per_hub <= 0 or not candidates: return candidates, [] hub_best_score: dict[str, float] = {} seed_sessions: set[str] = set() seed_chunk_ids: list[str] = [] for candidate in candidates: hub = _candidate_event_uuid(candidate) if hub: score = _candidate_distance(candidate) prev = hub_best_score.get(hub) if prev is None or score < prev: hub_best_score[hub] = score for sid in _candidate_session_ids(candidate): seed_sessions.add(sid) seed_chunk_ids.extend(str(c) for c in (candidate.get("chunk_ids") or []) if c) ranked_seeds = _rank_bridge_seed_hubs( candidates, cap=_BRIDGE_SEED_HUB_CAP ) if not ranked_seeds: return candidates, [] reserve_hubs = set( ranked_seeds[: min(_BRIDGE_RESERVE_HUB_CAP, len(ranked_seeds))] ) try: bridges = graph_adapter.get_hub_bridges(ranked_seeds, brain_id=brain_id) except Exception: return candidates, [] if not bridges: return candidates, [] typed_bridges: list[HubBridge] = [] for item in bridges: if isinstance(item, HubBridge): typed_bridges.append(item) else: typed_bridges.append( HubBridge( event_a=str(item.get("event_a") or ""), event_b=str(item.get("event_b") or ""), shared_entity=str(item.get("shared_entity") or ""), shared_entity_name=str(item.get("shared_entity_name") or ""), weight=float(item.get("weight") or 1.0), ) ) expansions = select_bridge_neighbors( ranked_seeds, typed_bridges, max_per_hub=max_per_hub, ) if not expansions: return candidates, [] selected_neighbors = {neighbor for neighbor, _ in expansions} neighbor_uuids = sorted(selected_neighbors) with profile_stage( "graph.bridge_neighbors", neighbors=len(neighbor_uuids) ) as detail: neighbors = graph_adapter.get_event_hub_facts( neighbor_uuids, brain_id=brain_id ) detail["rows"] = len(neighbors) neighbor_chunk_ids: list[str] = [] facts_by_hub: dict[str, list[tuple[Any, Any, Any, Any, Any]]] = {} for n, r, m, r2, b in neighbors: event_uuid = str(getattr(m, "uuid", "") or "") if not event_uuid: continue facts_by_hub.setdefault(event_uuid, []).append((n, r, m, r2, b)) neighbor_chunk_ids.extend(_extract_source_chunk_ids(n, r, m, r2, b)) chunk_sessions = _resolve_chunk_sessions( seed_chunk_ids + neighbor_chunk_ids, brain_id ) if not seed_sessions: seed_sessions.update(_sessions_from_chunk_map(seed_chunk_ids, chunk_sessions)) hub_sessions: dict[str, set[str]] = {} for event_uuid, rows in facts_by_hub.items(): sessions: set[str] = set() for n, r, m, r2, b in rows: chunk_ids = _extract_source_chunk_ids(n, r, m, r2, b) sessions.update(_sessions_from_chunk_map(chunk_ids, chunk_sessions)) hub_sessions[event_uuid] = sessions bridge_by_neighbor = {neighbor: bridge for neighbor, bridge in expansions} seed_by_neighbor: dict[str, str] = {} for neighbor, bridge in expansions: if bridge.event_a == neighbor: seed_by_neighbor[neighbor] = bridge.event_b else: seed_by_neighbor[neighbor] = bridge.event_a path_meta: list[dict[str, Any]] = [] for neighbor, bridge in expansions: seed = seed_by_neighbor.get(neighbor, "") path_meta.append( { "hubs": sorted([seed, neighbor]), "shared_entity": bridge.shared_entity, "shared_entity_name": bridge.shared_entity_name, "weight": bridge.weight, } ) path_meta.sort( key=lambda p: ( tuple(p["hubs"]), p["shared_entity"], ) ) existing_keys = { candidate.get("key") for candidate in candidates if candidate.get("key") is not None } expanded = list(candidates) for event_uuid in sorted(selected_neighbors): bridge = bridge_by_neighbor.get(event_uuid) if bridge is None: continue seed = seed_by_neighbor.get(event_uuid, "") base = hub_best_score.get(seed, float("inf")) novel = hub_sessions.get(event_uuid, set()) - seed_sessions score = ( base + _BRIDGE_SCORE_PENALTY if base != float("inf") else float("inf") ) if novel and score != float("inf"): score = max(0.0, score - min(0.05, 0.02 * len(novel))) reserve_ok = seed in reserve_hubs for n, r, m, r2, b in facts_by_hub.get(event_uuid, []): if not _fact_predicates_allowed(r, r2, include_history=include_history): continue key = _fact_key(r, r2) if key in existing_keys: continue existing_keys.add(key) chunk_ids = _extract_source_chunk_ids(n, r, m, r2, b) session_ids = _sessions_from_chunk_map(chunk_ids, chunk_sessions) fact_text = _format_event_fact(n, r, m, r2, b) if include_history and not ( _is_currently_valid(r) and _is_currently_valid(r2) ): fact_text = f"[superseded] {fact_text}" expanded.append( { "identified_entity": f"bridge:{bridge.shared_entity}", "triple": (n, r, m, r2, b), "score": score, "key": key, "text": fact_text, "chunk_ids": chunk_ids, "session_ids": session_ids, "bridge": { "from_hub": seed, "to_hub": event_uuid, "shared_entity": bridge.shared_entity, "shared_entity_name": bridge.shared_entity_name, "reserve_ok": reserve_ok, }, } ) return expanded, path_meta async def retrieve_data( text: str, limit: int, preferred_entities: str, brain_id: str = "default" ) -> RetrieveRequestResponse: """ Retrieve data from the knowledge graph and data store. """ if preferred_entities: preferred_entities = [ e.strip() for e in preferred_entities.split(",") if e.strip() ] else: preferred_entities = [] def _get_data(): text_embeddings = embeddings_adapter.embed_query(text) data_vectors = vector_search.search_data( text_embeddings.embeddings, brain_id=brain_id, k=limit, ) triple_vectors = vector_search.search_triplets( text_embeddings.embeddings, brain_id=brain_id, k=limit, ) search_result = data_adapter.search(text, brain_id) ts_text_chunks = search_result.text_chunks ts_observations = search_result.observations v_text_chunks, v_observations = data_adapter.get_text_chunks_by_ids( [dv.metadata.get("resource_id") for dv in data_vectors], True, brain_id ) node_ids = [ node_id for tv in triple_vectors for node_id in tv.metadata.get("node_ids", []) ] nodes = graph_adapter.get_nodes_by_uuid( uuids=node_ids, brain_id=brain_id, with_relationships=True, relationships_depth=1, relationships_type=[ tv.metadata.get("predicate") for tv in triple_vectors if tv.metadata.get("predicate", None) ], preferred_labels=preferred_entities or [], ) return ts_text_chunks, ts_observations, v_text_chunks, v_observations, nodes ts_text_chunks, ts_observations, v_text_chunks, v_observations, nodes = ( await asyncio.to_thread(_get_data) ) return RetrieveRequestResponse( data=[*ts_text_chunks, *v_text_chunks], observations=[*ts_observations, *v_observations], relationships=nodes, ) async def retrieve_neighbors( uuid: Optional[str] = None, look_for: Optional[str] = None, identification_params: Optional[IdentificationParams] = None, limit: int = 10, brain_id: str = "default", ) -> RetrieveNeighborsRequestResponse: """ Retrieve neighboring nodes related to a specified main node. If `uuid` is provided it is used to locate the main node; otherwise `identification_params` is used. Optionally filters first-degree neighbors by semantic similarity to `look_for`, expands those matches to similar nodes, and returns a deduplicated list of neighbor nodes (up to `limit`) with their relationship and the most common matching similar node. Parameters: uuid (Optional[str]): UUID of the main node to retrieve neighbors for. If omitted, `identification_params` must be provided. look_for (Optional[str]): Text used to filter first-degree neighbors by embedding similarity before expanding to similar nodes. identification_params (Optional[IdentificationParams]): Identification parameters used to find the main node when `uuid` is not provided. limit (int): Maximum number of neighbor results to include in the response. brain_id (str): Identifier of the brain / dataset to query. Returns: RetrieveNeighborsRequestResponse: Object containing: - count: total number of unique neighbors found, - main_node: the resolved main Node, - neighbors: list of RetrievedNeighborNode objects (neighbor, relationship, most_common) limited to `limit`. Raises: HTTPException: 404 if the main node cannot be found. """ async def _get_neighbors(): # --------------------------------------------------------- # ================= GETTING THE MAIN NODE ================= # --------------------------------------------------------- def _get_node() -> Node: node = None if uuid: node = graph_adapter.get_by_uuid(uuid, brain_id) elif identification_params: node = graph_adapter.get_by_identification_params( identification_params, brain_id=brain_id, entity_types=identification_params.entity_types, ) if not node: raise HTTPException(status_code=404, detail="Entity not found") return node node = await asyncio.to_thread(_get_node) target_node_types = node.labels looking_for_v = embeddings_adapter.embed_query(look_for) if look_for else None # --------------------------------------------------------- # ===== Getting 1st degree neighbors of the main node ===== # --------------------------------------------------------- def _get_fd_neighbors() -> ( tuple[dict[str, list[tuple[Predicate, Node]]], list[str]] ): fd_neighbors = graph_adapter.get_neighbors( [node.uuid], limit=limit, brain_id=brain_id ) fd_v_neighbors_ids = [ fd[1].properties.get("v_id") for fd in fd_neighbors[node.uuid] if fd[1].properties.get("v_id") is not None ] if look_for: fd_v_neighbors_embeddings = vector_store_adapter.get_by_ids( fd_v_neighbors_ids, brain_id=brain_id, store="nodes" ) fd_v_neighbors_embeddings_map = { v.id: v.embeddings for v in fd_v_neighbors_embeddings if ( cosine_similarity(looking_for_v.embeddings, v.embeddings) > 0.5 and v.id and not v.id.replace( "-", "" ).isalpha() # likely not a UUID if all numeric (may have hyphens for uuid standard) ) } fd_v_neighbors_ids = list(fd_v_neighbors_embeddings_map.keys()) return fd_neighbors, fd_v_neighbors_ids fd_neighbors, fd_v_neighbors_ids = await asyncio.to_thread(_get_fd_neighbors) # --------------------------------------------------------- # === Getting nodes similar to the 1st degree neighbors === # --------------------------------------------------------- fd_v_similar_node_futures = [] for fd_v_neighbor_id in fd_v_neighbors_ids: fd_v_similar_node_futures.append( asyncio.to_thread( vector_store_adapter.search_similar_by_ids, [fd_v_neighbor_id], brain_id, "nodes", 0.5, limit, ) ) fd_v_similar_nodes_results: list[dict[str, list[Vector]]] = ( await asyncio.gather(*fd_v_similar_node_futures) ) fd_similar_node_ids = [ v.metadata.get("uuid") for result_dict in fd_v_similar_nodes_results for vectors in result_dict.values() for v in vectors if v.metadata.get("uuid") is not None ] fd_similar_nodes = await asyncio.to_thread( graph_adapter.get_by_uuids, fd_similar_node_ids, brain_id ) fd_similar_nodes_by_uuid = {n.uuid: n for n in fd_similar_nodes} # --------------------------------------------------------- # === Getting neighbors of the 1st degree similar nodes === # --------------------------------------------------------- def _get_fd_similar_node_neighbors() -> dict[str, list[tuple[Predicate, Node]]]: fd_similar_node_neighbors = graph_adapter.get_neighbors( fd_similar_node_ids, limit=limit, brain_id=brain_id, of_types=list(set(target_node_types)), ) return fd_similar_node_neighbors fd_similar_node_neighbors = await asyncio.to_thread( _get_fd_similar_node_neighbors ) seen_neighbor_uuids = set() unique_neighbors = [] for source_uuid, neighbors_list in fd_similar_node_neighbors.items(): for neighbor_tuple in neighbors_list: neighbor_uuid = neighbor_tuple[1].uuid if neighbor_uuid not in seen_neighbor_uuids: seen_neighbor_uuids.add(neighbor_uuid) unique_neighbors.append( RetrievedNeighborNode( neighbor=neighbor_tuple[1], relationship=neighbor_tuple[0], most_common=fd_similar_nodes_by_uuid.get(source_uuid), ) ) return RetrieveNeighborsRequestResponse( count=len(unique_neighbors), main_node=node, neighbors=unique_neighbors[:limit], ) return await _get_neighbors() async def retrieve_neighbors_ai_mode( identification_params: IdentificationParams, looking_for: Optional[list[str]], limit: int, brain_id: str = "default", ) -> RetrieveNeighborsRequestResponse: """ Retrieve neighbors of an entity from the knowledge graph. """ def _get_neighbors(): node = graph_adapter.get_by_identification_params( identification_params, brain_id=brain_id, entity_types=identification_params.entity_types, ) if not node: raise HTTPException(status_code=404, detail="Entity not found") result = kg_agent.retrieve_neighbors(node, looking_for, limit, brain_id) ids = [neighbor.uuid for neighbor in result.neighbors] descriptions = [neighbor.description for neighbor in result.neighbors] nodes = graph_adapter.get_nodes_by_uuid(uuids=ids, brain_id=brain_id) paired = list(zip(nodes, descriptions)) return RetrieveNeighborsRequestResponse(neighbors=paired) result = await asyncio.to_thread(_get_neighbors) return result async def get_relationships( limit: int = 10, skip: int = 0, relationship_types: Optional[list[str]] = None, from_node_labels: Optional[list[str]] = None, to_node_labels: Optional[list[str]] = None, query_text: Optional[str] = None, query_search_target: Optional[str] = "all", brain_id: str = "default", ): """ Retrieve relationships from the knowledge graph with optional filtering and pagination. Parameters: relationship_types (list[str], optional): Filter results to specific relationship types. from_node_labels (list[str], optional): Filter relationships originating from nodes with these labels. to_node_labels (list[str], optional): Filter relationships targeting nodes with these labels. query_text (str, optional): Text to search within relationship or node content. query_search_target (str, optional): Field to target for text search; commonly "all", "from", or "to". limit (int, optional): Maximum number of relationships to return. skip (int, optional): Number of relationships to skip (offset). brain_id (str, optional): Identifier of the brain/graph to query. Returns: JSONResponse: A response whose JSON content contains: - message: Confirmation string. - relationships: List of serialized relationship objects. - total: Total number of matching relationships. """ relationships = await asyncio.to_thread( search_relationships, limit, skip, relationship_types, from_node_labels, to_node_labels, query_text, query_search_target, brain_id, ) return JSONResponse( content={ "message": "Relationships retrieved successfully", "relationships": [r.model_dump(mode="json") for r in relationships.results], "total": relationships.total, } ) async def get_entities( limit: int = 10, skip: int = 0, node_labels: Optional[list[str]] = None, query_text: Optional[str] = None, brain_id: str = "default", ): """ Retrieve entities from the knowledge graph with optional label and text filters. Parameters: limit (int): Maximum number of entities to return (pagination). skip (int): Number of entities to skip (pagination offset). node_labels (Optional[list[str]]): If provided, only return entities whose labels match any value in this list. query_text (Optional[str]): If provided, filter entities by matching text content. brain_id (str): Identifier of the knowledge graph/brain to query. Returns: JSONResponse: Object containing: - message (str): Informational message. - entities (list): Serialized entity objects. - total (int): Total number of matching entities. """ entities = await asyncio.to_thread( search_entities, limit, skip, node_labels, query_text, brain_id ) return JSONResponse( content={ "message": "Entities retrieved successfully", "entities": [e.model_dump(mode="json") for e in entities.results], "total": entities.total, } ) def _format_kg_item(item: Any) -> str: name = getattr(item, "name", None) or "" description = getattr(item, "description", None) or "" happened_at = getattr(item, "happened_at", None) or "" base = f"{name}: {description}" if description else name if happened_at: return f"{base} @{happened_at}" return base def _format_event_fact(n: Node, r: Predicate, m: Node, r2: Predicate, b: Node) -> str: r2_name = (getattr(r2, "name", None) or "").upper() end_label = ( "Context" if any(token in r2_name for token in _CONTEXT_REL) else "Target" ) return ( f"Actor: {_format_kg_item(n)} | {_format_kg_item(r)} | " f"Event: {_format_kg_item(m)} | {_format_kg_item(r2)} | " f"{end_label}: {_format_kg_item(b)}" ) def _fact_key(r: Predicate, r2: Predicate) -> tuple[str, str]: return (getattr(r, "uuid", None) or "", getattr(r2, "uuid", None) or "") def _rel_name(predicate: Predicate | None) -> str: return (getattr(predicate, "name", None) or "").upper() def _is_spine_actor_rel(name: str) -> bool: return any(token in name for token in _SPINE_ACTOR) def _is_spine_target_rel(name: str) -> bool: return any(token in name for token in _SPINE_TARGET) def _is_context_rel(name: str) -> bool: return any(token in name for token in _CONTEXT_REL) def _event_hub_id(candidate: dict[str, Any]) -> str: triple = candidate.get("triple") or () r = triple[1] if len(triple) > 1 else None m = triple[2] if len(triple) > 2 else None r2 = triple[3] if len(triple) > 3 else None event_uuid = getattr(m, "uuid", None) if m is not None else None if event_uuid: return str(event_uuid) flow_key = ( getattr(r, "flow_key", None) or getattr(r2, "flow_key", None) or "" ) if flow_key: return str(flow_key) key = candidate.get("key") if isinstance(key, tuple): return "|".join(str(part or "") for part in key) return str(candidate.get("text") or "") def _candidate_session_ids(candidate: dict[str, Any]) -> list[str]: sessions = candidate.get("session_ids") if isinstance(sessions, list) and sessions: return [str(s) for s in sessions if s] return [] def _hub_leg_kind(candidate: dict[str, Any]) -> str: triple = candidate.get("triple") or () r2 = triple[3] if len(triple) > 3 else None name = _rel_name(r2) if _is_context_rel(name): return "context" if _is_spine_target_rel(name): return "spine" return "other" def _hub_completeness_score(candidate: dict[str, Any]) -> float: triple = candidate.get("triple") or () if len(triple) < 5: return 0.0 n, r, m, r2, b = triple score = 0.0 m_labels = {str(label).upper() for label in (getattr(m, "labels", None) or [])} if "EVENT" in m_labels: score += 1.0 r_name = _rel_name(r) r2_name = _rel_name(r2) if _is_spine_actor_rel(r_name): score += 1.0 elif r_name: score += 0.2 if _is_spine_target_rel(r2_name): score += 1.0 elif _is_context_rel(r2_name): score += 0.35 elif r2_name: score += 0.15 if getattr(n, "uuid", None) and getattr(b, "uuid", None): score += 0.5 if getattr(m, "happened_at", None): score += 0.15 return score def _fact_recency_score( candidate: dict[str, Any], *, now: datetime | None = None ) -> float: triple = candidate.get("triple") or () best = 0.0 now = now or datetime.now(tz=None) for node in triple[0::2] if triple else (): raw = getattr(node, "happened_at", None) if not raw: continue parsed = parse_date_string(raw) if isinstance(raw, str) else raw if not isinstance(parsed, datetime): continue days_ago = max(0, (now - to_naive_utc(parsed)).days) score = 1.0 if days_ago <= 0 else 1.0 / (1.0 + days_ago) if score > best: best = score return best def _relevance_score( candidate: dict[str, Any], *, now: datetime | None = None ) -> float: distance = _candidate_distance(candidate) distance_term = 1.0 / (1.0 + max(0.0, distance)) ppr_term = float(candidate.get("ppr_mass") or 0.0) completeness = _hub_completeness_score(candidate) / 3.8 recency = _fact_recency_score(candidate, now=now) topic_boost = 0.08 if candidate.get("topic_preferred") else 0.0 return ( 0.45 * distance_term + 0.25 * ppr_term + 0.20 * completeness + 0.10 * recency + topic_boost ) def _completeness_rank_key(candidate: dict[str, Any]) -> tuple: return ( -_relevance_score(candidate), _candidate_distance(candidate), ) + _candidate_tiebreak_key(candidate) def _rank_facts_with_completeness( candidates: list[dict[str, Any]], ) -> list[dict[str, Any]]: return sorted(candidates, key=_completeness_rank_key) def _is_complementary_hub_leg( selected_meta: list[dict[str, Any]], item: dict[str, Any] ) -> bool: if item["kind"] not in {"spine", "context"}: return False for prev in selected_meta: if prev["hub"] != item["hub"]: continue if {item["kind"], prev["kind"]} == {"spine", "context"}: return True return False def _bridge_reserved_slots(max_facts: int) -> int: if max_facts <= 0: return 0 return min( max_facts, max(_BRIDGE_RESERVE_MIN, max_facts // _BRIDGE_RESERVE_FRAC_DENOM), ) def _is_bridge_candidate(candidate: dict[str, Any]) -> bool: return bool(candidate.get("bridge")) def _bridge_reserve_ok(candidate: dict[str, Any]) -> bool: bridge = candidate.get("bridge") if not bridge: return False if isinstance(bridge, dict) and "reserve_ok" in bridge: return bool(bridge.get("reserve_ok")) return True def _prepare_diversify_item( candidate: dict[str, Any], *, now: datetime | None = None ) -> dict[str, Any]: key = candidate.get("key") or ("", "") is_bridge = _is_bridge_candidate(candidate) return { "candidate": candidate, "relevance": _relevance_score(candidate, now=now), "hub": _event_hub_id(candidate), "sessions": frozenset(_candidate_session_ids(candidate)), "kind": _hub_leg_kind(candidate), "tie": (str(key[0]) if key else "", str(key[1]) if len(key) > 1 else ""), "is_bridge": is_bridge, "reserve_ok": _bridge_reserve_ok(candidate) if is_bridge else False, } def _mmr_pick_index( remaining: list[dict[str, Any]], selected_meta: list[dict[str, Any]], selected_hubs: set[str], selected_sessions: set[str], ) -> int: best_idx = 0 best_mmr = None best_tie = None selected_kinds_by_hub: dict[str, set[str]] = defaultdict(set) for selected in selected_meta: selected_kinds_by_hub[selected["hub"]].add(selected["kind"]) for idx, item in enumerate(remaining): relevance = item["relevance"] hub = item["hub"] sessions = item["sessions"] complementary = ( item["kind"] in {"spine", "context"} and bool( selected_kinds_by_hub.get(item["hub"], set()) & ({"spine", "context"} - {item["kind"]}) ) ) if complementary: hub_pen = 0.0 sess_pen = 0.0 relevance = relevance + 0.5 else: hub_pen = 1.0 if hub in selected_hubs else 0.0 if sessions and selected_sessions: sess_pen = len(sessions & selected_sessions) / max( 1, len(sessions) ) else: sess_pen = 0.0 redundancy = 0.55 * hub_pen + 0.45 * sess_pen mmr = _DIVERSIFY_LAMBDA * relevance - (1.0 - _DIVERSIFY_LAMBDA) * redundancy tie = item["tie"] if ( best_mmr is None or mmr > best_mmr or (mmr == best_mmr and (best_tie is None or tie < best_tie)) ): best_mmr = mmr best_tie = tie best_idx = idx return best_idx def _mmr_take( remaining: list[dict[str, Any]], *, budget: int, selected: list[dict[str, Any]], selected_meta: list[dict[str, Any]], selected_hubs: set[str], selected_sessions: set[str], ) -> None: while budget > 0 and remaining: best_idx = _mmr_pick_index( remaining, selected_meta, selected_hubs, selected_sessions ) chosen = remaining.pop(best_idx) selected.append(chosen["candidate"]) selected_meta.append(chosen) selected_hubs.add(chosen["hub"]) selected_sessions.update(chosen["sessions"]) budget -= 1 def _take_novel_bridge_slots( remaining: list[dict[str, Any]], *, budget: int, covered_sessions: set[str], selected: list[dict[str, Any]], selected_meta: list[dict[str, Any]], selected_hubs: set[str], selected_sessions: set[str], ) -> None: while budget > 0 and remaining: novel_sessions: set[str] = set() for item in remaining: if not item["is_bridge"] or not item["reserve_ok"]: continue novel_sessions.update(item["sessions"] - covered_sessions) if not novel_sessions: break target = min(novel_sessions) best_idx = None best_key = None for idx, item in enumerate(remaining): if not item["is_bridge"] or not item["reserve_ok"]: continue if target not in item["sessions"]: continue key = (item["hub"], item["tie"]) if best_key is None or key < best_key: best_key = key best_idx = idx if best_idx is None: break chosen = remaining.pop(best_idx) selected.append(chosen["candidate"]) selected_meta.append(chosen) selected_hubs.add(chosen["hub"]) selected_sessions.update(chosen["sessions"]) covered_sessions.update(chosen["sessions"]) budget -= 1 def _diversify_facts( ranked: list[dict[str, Any]], *, max_facts: int, ) -> list[dict[str, Any]]: if max_facts <= 0 or not ranked: return [] if len(ranked) <= max_facts: return list(ranked) ranking_now = datetime.now(tz=None) prepared = [ _prepare_diversify_item(candidate, now=ranking_now) for candidate in ranked ] prepared.sort(key=lambda item: (-item["relevance"], item["tie"])) selected: list[dict[str, Any]] = [] selected_meta: list[dict[str, Any]] = [] selected_hubs: set[str] = set() selected_sessions: set[str] = set() reserve = _bridge_reserved_slots(max_facts) has_bridge = any(item["is_bridge"] and item["reserve_ok"] for item in prepared) if not has_bridge: reserve = 0 main_budget = max_facts - reserve non_bridge = [item for item in prepared if not item["is_bridge"]] bridge_items = [item for item in prepared if item["is_bridge"]] stable_bridges = [item for item in bridge_items if item["reserve_ok"]] _mmr_take( non_bridge, budget=main_budget, selected=selected, selected_meta=selected_meta, selected_hubs=selected_hubs, selected_sessions=selected_sessions, ) covered_by_non_bridge = set(selected_sessions) _take_novel_bridge_slots( stable_bridges, budget=reserve, covered_sessions=covered_by_non_bridge, selected=selected, selected_meta=selected_meta, selected_hubs=selected_hubs, selected_sessions=selected_sessions, ) remaining = non_bridge + stable_bridges remaining.sort(key=lambda item: (-item["relevance"], item["tie"])) leftover = max_facts - len(selected) _mmr_take( remaining, budget=leftover, selected=selected, selected_meta=selected_meta, selected_hubs=selected_hubs, selected_sessions=selected_sessions, ) return selected def _temporal_conflict_meta( candidates: list[dict[str, Any]], ) -> list[dict[str, Any]]: by_entity: dict[str, list[dict[str, Any]]] = defaultdict(list) for candidate in candidates: triple = candidate.get("triple") or () for node in (triple[0], triple[4]) if len(triple) >= 5 else (): uuid = getattr(node, "uuid", None) if not uuid: continue happened = None event = triple[2] if len(triple) > 2 else None raw = getattr(event, "happened_at", None) if event is not None else None if isinstance(raw, str) and raw.strip(): happened = raw.strip() by_entity[str(uuid)].append( { "fact_key": list(candidate.get("key") or ()), "hub_id": _event_hub_id(candidate), "happened_at": happened, } ) conflicts: list[dict[str, Any]] = [] for entity_uuid, entries in sorted(by_entity.items()): times = { entry["happened_at"] for entry in entries if entry.get("happened_at") } hubs = {entry["hub_id"] for entry in entries} if len(times) >= 2 and len(hubs) >= 2: conflicts.append( { "entity_uuid": entity_uuid, "hubs": sorted(hubs), "happened_ats": sorted(times), "facts": [entry["fact_key"] for entry in entries], } ) return conflicts def _edge_type_weight(predicate: Predicate | None, query: str | None = None) -> float: name = _rel_name(predicate) if _is_spine_actor_rel(name) or _is_spine_target_rel(name): weight = 1.0 elif _is_context_rel(name): weight = 0.35 else: weight = 0.55 if query: q_tokens = set(re.findall(r"[a-z0-9]+", query.lower())) p_tokens = set(re.findall(r"[a-z0-9]+", name.lower())) if q_tokens and p_tokens: overlap = len(q_tokens & p_tokens) / max(1, len(p_tokens)) weight *= 1.0 + 0.5 * overlap return weight def _add_weighted_edge( adjacency: dict[str, dict[str, float]], src: str | None, dst: str | None, weight: float, ) -> None: if not src or not dst or src == dst or weight <= 0: return bucket = adjacency.setdefault(str(src), {}) prev = bucket.get(str(dst), 0.0) if weight > prev: bucket[str(dst)] = weight def _candidate_tiebreak_key(candidate: dict[str, Any]) -> tuple[str, ...]: triple = candidate.get("triple") or () n = triple[0] if len(triple) > 0 else None r = triple[1] if len(triple) > 1 else None m = triple[2] if len(triple) > 2 else None r2 = triple[3] if len(triple) > 3 else None b = triple[4] if len(triple) > 4 else None flow_key = ( getattr(r, "flow_key", None) or getattr(r2, "flow_key", None) or "" ) key = candidate.get("key") if not (isinstance(key, tuple) and len(key) == 2): key = _fact_key(r, r2) if r is not None and r2 is not None else ("", "") r_uuid, r2_uuid = key return ( str(flow_key or ""), str(r_uuid or ""), str(r2_uuid or ""), str(getattr(n, "uuid", None) or ""), str(getattr(m, "uuid", None) or ""), str(getattr(b, "uuid", None) or ""), str(candidate.get("identified_entity") or ""), str(candidate.get("text") or ""), ) def _distance_rank_key(candidate: dict[str, Any]) -> tuple: return (_candidate_distance(candidate),) + _candidate_tiebreak_key( candidate ) def _ppr_rank_key( candidate: dict[str, Any], ppr_scores: dict[str, float] ) -> tuple: return (_score_candidate_with_ppr(candidate, ppr_scores),) + _candidate_tiebreak_key( candidate ) def _dedupe_candidates(candidates: list[dict[str, Any]]) -> list[dict[str, Any]]: deduped: dict[tuple[str, str], dict[str, Any]] = {} for candidate in candidates: key = candidate["key"] existing = deduped.get(key) if existing is None: deduped[key] = candidate continue if candidate["score"] < existing["score"]: deduped[key] = candidate continue if candidate["score"] == existing["score"] and _candidate_tiebreak_key( candidate ) < _candidate_tiebreak_key(existing): deduped[key] = candidate return sorted(deduped.values(), key=_distance_rank_key) def _merge_variant_candidate_lists( variant_lists: list[list[dict[str, Any]]], ) -> list[dict[str, Any]]: merged: list[dict[str, Any]] = [] for local in variant_lists: ordered = sorted(local, key=_distance_rank_key) merged.extend(ordered) return merged def _is_currently_valid(predicate: Predicate) -> bool: props = getattr(predicate, "properties", None) or {} if props.get("invalid_at"): return False if getattr(predicate, "deprecated", False): return False return True def _wants_historical_facts(query: str) -> bool: text = (query or "").strip() if not text: return False if _CURRENT_TRUTH_RE.search(text) and not _HISTORY_MODE_RE.search(text): return False return bool(_HISTORY_MODE_RE.search(text)) def _fact_predicates_allowed( r: Predicate, r2: Predicate, *, include_history: bool, ) -> bool: if include_history: return True return _is_currently_valid(r) and _is_currently_valid(r2) def _collect_query_variants(text: str, elements) -> list[str]: seen: set[str] = set() variants: list[str] = [] def _add(value: str | None) -> None: if not value: return cleaned = value.strip() if not cleaned: return key = cleaned.lower() if key in seen: return seen.add(key) variants.append(cleaned) _add(text) for token in getattr(elements, "tokens", None) or []: if isinstance(token, dict): _add(token.get("text")) noun_chunks = getattr(elements, "noun_chunks", None) or [] added_chunks = 0 for chunk in noun_chunks: if added_chunks >= _MAX_NOUN_CHUNKS: break before = len(variants) _add(chunk if isinstance(chunk, str) else None) if len(variants) > before: added_chunks += 1 for extra in _decompose_enumeration_queries(text, elements): _add(extra) for extra in _ordering_aspect_queries(text): _add(extra) return variants _ENUMERATION_RE = re.compile( r"\b(what|which|list|ways?|types?|kinds?|events?|books?|items?|" r"changes?|traits?|hobbies|activities)\b", re.IGNORECASE, ) _ORDERING_ASPECT_QUERIES_BUDGET = ( "core functionality authentication expense tracking visualization", "transaction CRUD error handling response management", "security deployment integration tests gunicorn workers", ) _ORDERING_ASPECT_QUERIES_TRANSLATION = ( "translation API integration error handling DeepL", "rate limiting request queue caching Redis", "language detection libraries franc evaluation", "database schema optimization contextual memory store", ) def _is_enumeration_question(text: str) -> bool: return bool(_ENUMERATION_RE.search(text or "")) def _is_ordering_question(text: str) -> bool: return bool( re.search( r"\b(order|ordered|ordering|chronolog|sequence|walk me through)\b", text or "", re.I, ) ) def _ordering_aspect_queries(text: str) -> list[str]: if not _is_ordering_question(text): return [] q = text or "" if re.search( r"\b(translation|language detection|deepl|franc|microservice|" r"multi-language|rate limit)\b", q, re.I, ): return list(_ORDERING_ASPECT_QUERIES_TRANSLATION)[:2] return list(_ORDERING_ASPECT_QUERIES_BUDGET)[:2] def _decompose_enumeration_queries(text: str, elements) -> list[str]: if not _is_enumeration_question(text): return [] queries: list[str] = [] entities: list[str] = [] for token in getattr(elements, "tokens", None) or []: if not isinstance(token, dict): continue label = str(token.get("label") or token.get("ent_type") or "").upper() value = str(token.get("text") or "").strip() if not value: continue if label in {"PERSON", "PER", "ORG", "GPE"} or ( value[:1].isupper() and " " not in value ): entities.append(value) relation_terms = [ m.group(0).lower() for m in re.finditer( r"\b(books?|read|ways?|types?|kinds?|events?|hobbies|" r"activities|changes?|traits?|friends?|places?)\b", text or "", re.IGNORECASE, ) ] seen: set[str] = set() for entity in entities[:4]: for relation in relation_terms[:4]: q = f"{entity} {relation}" key = q.lower() if key not in seen and key != (text or "").strip().lower(): seen.add(key) queries.append(q) key = entity.lower() if key not in seen: seen.add(key) queries.append(entity) for relation in relation_terms[:3]: if relation not in seen: seen.add(relation) queries.append(relation) return queries[:8] def _seed_personalization(seeds: list[tuple[str, float, str]]) -> dict[str, float]: weights: dict[str, float] = {} for uuid, distance, _ in seeds: weight = 1.0 / (1.0 + max(0.0, _quantize_distance(distance))) prev = weights.get(uuid) if prev is None or weight > prev: weights[uuid] = weight return weights def _build_adjacency_from_seeds( seed_uuids: list[str], brain_id: str, *, query: str | None = None, query_embedding: list[float] | None = None, include_history: bool = False, ) -> dict[str, list[tuple[str, float]]]: weighted: dict[str, dict[str, float]] = {} if not seed_uuids: return {} try: neighbors = graph_adapter.get_event_centric_neighbors( seed_uuids, brain_id=brain_id ) except Exception: return {} for n, r, m, r2, b in neighbors: if not _fact_predicates_allowed(r, r2, include_history=include_history): continue n_uuid = getattr(n, "uuid", None) m_uuid = getattr(m, "uuid", None) b_uuid = getattr(b, "uuid", None) w1 = _edge_type_weight(r, query=query) w2 = _edge_type_weight(r2, query=query) if query_embedding is not None: for pred, base in ((r, w1), (r2, w2)): emb = getattr(pred, "embeddings", None) or getattr( pred, "embedding", None ) if emb is None: continue try: sim = float(cosine_similarity(query_embedding, emb)) boost = 0.5 + 0.5 * max(0.0, sim) if pred is r: w1 *= boost else: w2 *= boost except Exception: pass _add_weighted_edge(weighted, n_uuid, m_uuid, w1) _add_weighted_edge(weighted, m_uuid, n_uuid, w1) _add_weighted_edge(weighted, m_uuid, b_uuid, w2) _add_weighted_edge(weighted, b_uuid, m_uuid, w2) if n_uuid and b_uuid: bridge = min(w1, w2) * 0.5 _add_weighted_edge(weighted, n_uuid, b_uuid, bridge) _add_weighted_edge(weighted, b_uuid, n_uuid, bridge) return { src: sorted(((dst, weight) for dst, weight in edges.items()), key=lambda x: x[0]) for src, edges in sorted(weighted.items(), key=lambda item: item[0]) } def _score_candidate_with_ppr( candidate: dict[str, Any], ppr_scores: dict[str, float] ) -> float: n, _r, m, _r2, b = candidate["triple"] scores = [ ppr_scores.get(str(getattr(node, "uuid", "")), 0.0) for node in (n, m, b) if getattr(node, "uuid", None) ] if not scores: return _candidate_distance(candidate) return -max(scores) def _context_looks_insufficient( question: str, text_lines: list[str], passages: list[str] ) -> bool: blob = "\n".join(text_lines + passages).lower() if len(blob.strip()) < 80: return True tokens = [ t for t in re.findall(r"[a-z0-9']+", (question or "").lower()) if len(t) > 3 and t not in { "what", "which", "when", "where", "who", "does", "did", "have", "has", "been", "with", "from", "about", "would", "could", "their", "there", "this", "that", "mentioned", } ] if not tokens: return False hits = sum(1 for token in tokens if token in blob) return hits / len(tokens) < 0.35 def _stabilize_seed_hits( seeds: list[tuple[str, float, str]], ) -> list[tuple[str, float, str]]: best: dict[str, tuple[float, str]] = {} for uuid, distance, name in seeds: uid = str(uuid or "").strip() if not uid: continue dist = _quantize_distance(distance) prev = best.get(uid) if prev is None or dist < prev[0]: best[uid] = (dist, str(name)) return sorted( ((uid, dist, name) for uid, (dist, name) in best.items()), key=lambda item: (item[1], item[0]), ) def _seed_nodes_for_text(text: str, brain_id: str) -> list[tuple[str, float, str]]: """Return (node_uuid, distance, entity_name) seeds from nodes + relationships.""" with profile_stage("embed.query"): text_embeddings = embeddings_adapter.embed_query(text) embeddings = text_embeddings.embeddings seeds: list[tuple[str, float, str]] = [] with profile_stage("vector.search_nodes", k=_SEED_K): node_vectors = vector_search.search_nodes( embeddings, brain_id=brain_id, k=_SEED_K, ) for vector in node_vectors: meta = vector.metadata or {} uuid = meta.get("uuid") or getattr(vector, "id", None) if not uuid: continue distance = ( vector.distance if vector.distance is not None else float("inf") ) name = meta.get("name") or text seeds.append((str(uuid), float(distance), str(name))) with profile_stage("vector.search_relationships", k=_SEED_K): rel_vectors = vector_search.search_relationships( embeddings, brain_id=brain_id, k=_SEED_K, ) for vector in rel_vectors: meta = vector.metadata or {} distance = ( vector.distance if vector.distance is not None else float("inf") ) node_ids = sorted(str(n) for n in (meta.get("node_ids") or []) if n) for node_id in node_ids: seeds.append((node_id, float(distance), text)) return _stabilize_seed_hits(seeds) def _extract_source_chunk_ids(*items: Any) -> list[str]: ids: list[str] = [] seen: set[str] = set() for item in items: props = getattr(item, "properties", None) or {} for value in props.get("source_chunk_ids") or []: chunk_id = str(value).strip() if chunk_id and chunk_id not in seen: seen.add(chunk_id) ids.append(chunk_id) single = props.get("source_chunk_id") if single: chunk_id = str(single).strip() if chunk_id and chunk_id not in seen: seen.add(chunk_id) ids.append(chunk_id) return ids def _session_ids_from_text(text: str) -> list[str]: found: list[str] = [] seen: set[str] = set() for match in _SESSION_RE.finditer(text or ""): sid = f"session_{match.group(1)}" if sid not in seen: seen.add(sid) found.append(sid) for match in _BATCH_TURN_RE.finditer(text or ""): sid = f"session_b{match.group(1)}_t{match.group(2)}" if sid not in seen: seen.add(sid) found.append(sid) return found def _session_ids_for_chunks( chunk_ids: list[str], chunk_sessions: dict[str, list[str]] ) -> list[str]: found: list[str] = [] seen: set[str] = set() for chunk_id in chunk_ids: for sid in chunk_sessions.get(chunk_id) or []: if sid not in seen: seen.add(sid) found.append(sid) return found def _build_fact_channel( curated: list[dict[str, Any]], chunk_sessions: dict[str, list[str]], ) -> tuple[list[str], list[GetContextTriple], list[str]]: text_lines: list[str] = [] triples: list[GetContextTriple] = [] graph_session_ids: list[str] = [] seen_sessions: set[str] = set() for candidate in curated: n, r, m, r2, b = candidate["triple"] chunk_ids = list(candidate.get("chunk_ids") or []) session_ids = _candidate_session_ids(candidate) if not session_ids: session_ids = _session_ids_for_chunks(chunk_ids, chunk_sessions) line = candidate["text"] if session_ids: line = f"{line} ({', '.join(session_ids)})" text_lines.append(line) triples.append( GetContextTriple( identified_entity=candidate["identified_entity"], triple=(n, r, m, r2, b), source_chunk_ids=chunk_ids or None, source_session_ids=session_ids or None, ) ) for sid in session_ids: if sid not in seen_sessions: seen_sessions.add(sid) graph_session_ids.append(sid) graph_session_ids.sort() return text_lines, triples, graph_session_ids def _retrieve_passages( text: str, brain_id: str, *, limit: int ) -> list[tuple[str, float, str]]: """Return (chunk_id, score, text) ranked passages via vector + keyword fusion.""" mode = ( config.context_passage_mode if config.search_enabled else "ilike" ) use_dense = mode in ("hybrid", "dense", "ilike") use_bm25 = mode in ("hybrid", "bm25") use_ilike = mode == "ilike" with profile_stage("passages.retrieve", blocking=True, queries=1): vector_ids: list[str] = [] id_to_distance: dict[str, float] = {} if use_dense: with profile_stage("embed.query"): text_embeddings = embeddings_adapter.embed_query(text) with profile_stage("vector.search_data", k=max(limit, _PASSAGE_K)): vector_ids, id_to_distance = collect_dense_passages( vector_search, text_embeddings.embeddings, brain_id, max(limit, _PASSAGE_K), ) keyword_ids: list[str] = [] id_to_text: dict[str, str] = {} if use_bm25: with profile_stage("data.bm25_search"): try: keyword_ids, _, id_to_text = collect_bm25_passages( data_adapter, text, brain_id, max(limit, _PASSAGE_K), ) except Exception: pass elif use_ilike: with profile_stage("data.keyword_search"): try: keyword_ids, id_to_text = collect_ilike_passages( data_adapter, text, brain_id ) except Exception: pass fused = fuse_passage_lists( vector_ids if use_dense else [], keyword_ids, ) ranked_ids = [item for item, _ in fused[:limit]] missing = [cid for cid in ranked_ids if cid not in id_to_text] if missing: with profile_stage("data.fetch_chunks", chunks=len(missing)): try: chunks, _ = data_adapter.get_text_chunks_by_ids( missing, False, brain_id ) for chunk in chunks: id_to_text[str(chunk.id)] = chunk.text or "" except Exception: pass passages: list[tuple[str, float, str]] = [] for chunk_id in ranked_ids: body = (id_to_text.get(chunk_id) or "").strip() if not body: continue score = id_to_distance.get(chunk_id, 1.0 / (1.0 + len(passages))) passages.append((chunk_id, float(score), body)) return passages async def get_context( request: GetContextRequestBody, *, fact_filter_adapter: Any = None, ) -> GetContextResponse: """ Retrieve contextual information for a text. `fact_filter_adapter` is only supplied by deep-navigation callers: the one-shot context path leaves it unset because an LLM call does not fit the sub-second budget, and without it no fact filtering happens. `profile_stages` attaches a per-stage wall-clock breakdown to the response and publishes it as a latency trace event; it is off unless the caller asks for it or `TRACE_STAGE_PROFILER_ENABLED` is set. """ with profile_request( "retrieve.context", enabled=request.profile_stages ) as profiler: response = await _build_context(request, fact_filter_adapter) if profiler is not None: response.stage_timings = profiler.last_report return response async def _build_context( request: GetContextRequestBody, fact_filter_adapter: Any, ) -> GetContextResponse: with profile_stage("nlp.extract_elements"): elements = _entity_extractor.extract_elements(request.text) with profile_stage("nlp.query_variants") as detail: variants = _collect_query_variants(request.text, elements) detail["variants"] = len(variants) max_passages = max(1, request.max_passages) historical_context: list[str] = [] source_passages: list[str] = [] candidates: list[dict[str, Any]] = [] seed_hits: list[tuple[str, float, str]] = [] passage_hits: list[tuple[str, float, str]] = [] include_history = _wants_historical_facts(request.text) def _collect_facts_for_variant( text: str, ) -> tuple[list[tuple[str, float, str]], list[dict[str, Any]]]: with profile_stage("facts.variant") as detail: with profile_stage("facts.seed_search"): seeds = _seed_nodes_for_text(text, request.brain_id) detail["seeds"] = len(seeds) if not seeds: return [], [] uuid_to_distance: dict[str, float] = {} for uuid, distance, _ in seeds: dist = _quantize_distance(distance) prev = uuid_to_distance.get(uuid) if prev is None or dist < prev: uuid_to_distance[uuid] = dist seed_uuids = sorted( uuid_to_distance.keys(), key=lambda uid: (uuid_to_distance[uid], uid), ) with profile_stage( "graph.event_neighbors", seed_uuids=len(seed_uuids) ) as neighbors_detail: neighbors = graph_adapter.get_event_centric_neighbors( seed_uuids, brain_id=request.brain_id ) neighbors_detail["rows"] = len(neighbors) local: list[dict[str, Any]] = [] with profile_stage("facts.assemble"): for n, r, m, r2, b in neighbors: if not _fact_predicates_allowed( r, r2, include_history=include_history ): continue distances = [ uuid_to_distance[u] for u in ( getattr(n, "uuid", None), getattr(m, "uuid", None), getattr(b, "uuid", None), ) if u in uuid_to_distance ] score = min(distances) if distances else float("inf") fact_text = _format_event_fact(n, r, m, r2, b) if include_history and not ( _is_currently_valid(r) and _is_currently_valid(r2) ): fact_text = f"[superseded] {fact_text}" local.append( { "identified_entity": text, "triple": (n, r, m, r2, b), "score": score, "key": _fact_key(r, r2), "text": fact_text, "chunk_ids": _extract_source_chunk_ids(n, r, m, r2, b), } ) return list(seeds), local def _collect_passages() -> None: nonlocal passage_hits with profile_stage("passages.collect") as detail: queries = [request.text] + [ v for v in variants if v.strip().lower() != request.text.strip().lower() ] if _is_enumeration_question(request.text): queries = queries[:6] else: queries = queries[:3] detail["queries"] = len(queries) fused_lists: list[list[str]] = [] id_to_text: dict[str, str] = {} id_to_score: dict[str, float] = {} for query in queries: hits = _retrieve_passages( query, request.brain_id, limit=max(max_passages, _PASSAGE_K), ) ordered_ids: list[str] = [] for chunk_id, score, body in hits: ordered_ids.append(chunk_id) id_to_text[chunk_id] = body prev = id_to_score.get(chunk_id) if prev is None or score < prev: id_to_score[chunk_id] = score if ordered_ids: fused_lists.append(ordered_ids) if not fused_lists: passage_hits = [] return fused = reciprocal_rank_fusion(fused_lists) ranked_ids = [item for item, _ in fused[: max(max_passages, _PASSAGE_K)]] passage_hits = [ (cid, id_to_score.get(cid, 1.0), id_to_text.get(cid, "")) for cid in ranked_ids if id_to_text.get(cid) ] async def _get_historical_context(): nonlocal historical_context with profile_stage("historical.context", blocking=False): passages = await asyncio.to_thread( _retrieve_passages, request.text, request.brain_id, limit=max(request.historical_limit, _PASSAGE_K), ) if passages: historical_context = [ text for _, _, text in passages[: request.historical_limit] ] return historical_context with profile_stage("historical.fallback", blocking=False): text_chunks, structured_data = await asyncio.gather( asyncio.to_thread( data_adapter.get_last_text_chunks, brain_id=request.brain_id, limit=request.historical_limit, ), asyncio.to_thread( data_adapter.get_last_structured_data, brain_id=request.brain_id, limit=request.historical_limit, ), ) historical_context = [text_chunk.text for text_chunk in text_chunks] + [ json.dumps(structured_data.data) for structured_data in structured_data if len(structured_data.data.items()) > 0 ] return historical_context with profile_stage( "retrieval.fanout", blocking=False, variants=len(variants) ): fact_futures = [ asyncio.to_thread(_collect_facts_for_variant, variant) for variant in variants ] other_futures = [ asyncio.to_thread(_collect_passages), _get_historical_context(), ] all_results = await asyncio.gather(*fact_futures, *other_futures) variant_outputs = all_results[: len(variants)] seed_hits = [] variant_lists: list[list[dict[str, Any]]] = [] for seeds, local in variant_outputs: seed_hits.extend(seeds) variant_lists.append(local) candidates = _merge_variant_candidate_lists(variant_lists) bridge_paths: list[dict[str, Any]] = [] with profile_stage( "facts.bridge_expand", bridges=request.cross_event_bridges, ) as detail: candidates, bridge_paths = _expand_cross_event_bridges( candidates, request.brain_id, max_per_hub=request.cross_event_bridges, include_history=include_history, ) detail["candidates"] = len(candidates) detail["paths"] = len(bridge_paths) with profile_stage("facts.dedup_rank") as detail: ranked = _dedupe_candidates(candidates) detail["candidates"] = len(candidates) detail["deduped"] = len(ranked) if request.use_ppr and seed_hits and ranked: with profile_stage("ppr"): personalization = _seed_personalization(seed_hits) with profile_stage( "ppr.adjacency", seeds=len(personalization) ) as adjacency_detail: adjacency = _build_adjacency_from_seeds( sorted(personalization.keys()), request.brain_id, query=request.text, include_history=include_history, ) adjacency_detail["nodes"] = len(adjacency) with profile_stage("ppr.iterations", iterations=_PPR_ITERS): ppr_scores = personalized_pagerank( adjacency, personalization, damping=_PPR_DAMPING, iterations=_PPR_ITERS, ) if ppr_scores: with profile_stage("ppr.reorder", scored=len(ppr_scores)): for candidate in ranked: n, _r, m, _r2, b = candidate["triple"] masses = [ ppr_scores.get(str(getattr(node, "uuid", "")), 0.0) for node in (n, m, b) if getattr(node, "uuid", None) ] candidate["ppr_mass"] = max(masses) if masses else 0.0 ranked = sorted( ranked, key=lambda c: _ppr_rank_key(c, ppr_scores) ) max_facts = max(0, request.max_facts) provenance_ids: list[str] = [] for candidate in ranked: provenance_ids.extend(candidate.get("chunk_ids") or []) for chunk_id, _, body in passage_hits: if body and body not in source_passages: source_passages.append(body) if chunk_id not in provenance_ids: provenance_ids.append(chunk_id) chunk_sessions: dict[str, list[str]] = {} if provenance_ids: with profile_stage( "provenance.chunks", blocking=False, chunk_ids=len(provenance_ids[:80]), ): try: chunks, _ = await asyncio.to_thread( data_adapter.get_text_chunks_by_ids, provenance_ids[:80], False, request.brain_id, ) for chunk in chunks: body = (chunk.text or "").strip() chunk_id = str(getattr(chunk, "id", "") or "") if chunk_id and body: sessions = _session_ids_from_text(body) if sessions: chunk_sessions[chunk_id] = sessions if body and body not in source_passages: source_passages.append(body) except Exception: pass for candidate in ranked: if not candidate.get("session_ids"): candidate["session_ids"] = _session_ids_for_chunks( list(candidate.get("chunk_ids") or []), chunk_sessions ) response_topics: list[dict[str, Any]] = [] with profile_stage("topics.coarse_to_fine") as topic_detail: try: from src.core.saving.topic_hyperedges import select_topics_and_sessions memberships = graph_adapter.list_topic_memberships(request.brain_id) seed_sessions: list[str] = [] seen_seed: set[str] = set() for candidate in ranked: for sid in _candidate_session_ids(candidate): if sid not in seen_seed: seen_seed.add(sid) seed_sessions.append(sid) response_topics, preferred_sessions = select_topics_and_sessions( request.text, memberships, seed_sessions, k_topics=10, k_sessions=20 if include_history else 10, ) preferred = set(preferred_sessions) topic_detail["topics"] = len(response_topics) topic_detail["preferred_sessions"] = len(preferred) if preferred: for candidate in ranked: sessions = set(_candidate_session_ids(candidate)) if sessions & preferred: candidate["topic_preferred"] = True except Exception as exc: topic_detail["error"] = str(exc)[:120] with profile_stage("fact_filter") as detail: if request.apply_fact_filter and fact_filter_adapter is not None and ranked: detail["filter_applied"] = True keep = filter_relevant_facts( request.text, [c["text"] for c in ranked], llm_adapter=fact_filter_adapter, max_keep=max(max_facts, len(ranked)), ) filtered = [ranked[i] for i in keep if 0 <= i < len(ranked)] if not filtered: filtered = ranked else: detail["filter_applied"] = False filtered = ranked with profile_stage("facts.diversify", max_facts=max_facts) as detail: filtered = _rank_facts_with_completeness(filtered) curated = _diversify_facts(filtered, max_facts=max_facts) detail["candidates"] = len(filtered) detail["curated"] = len(curated) with profile_stage("context.assemble") as detail: temporal_conflicts = _temporal_conflict_meta(curated) text_lines, triples, graph_session_ids = _build_fact_channel( curated, chunk_sessions ) source_passages = source_passages[:max_passages] detail["facts"] = len(text_lines) detail["passages"] = len(source_passages) insufficient = False if request.sufficiency_retry: with profile_stage("sufficiency.check") as detail: insufficient = _context_looks_insufficient( request.text, text_lines, source_passages + historical_context ) detail["fired"] = insufficient if insufficient: with profile_stage("sufficiency.retry", blocking=False) as detail: followups = _decompose_enumeration_queries(request.text, elements) if not followups: tokens = [ t for t in re.findall(r"[A-Za-z][A-Za-z']+", request.text) if len(t) > 3 ] followups = [" ".join(tokens[:4])] if tokens else [] detail["followups"] = len(followups[:3]) extra_passages: list[tuple[str, float, str]] = [] extra_variant_lists: list[list[dict[str, Any]]] = [] for query in followups[:3]: extra_passages.extend( await asyncio.to_thread( _retrieve_passages, query, request.brain_id, limit=max_passages, ) ) seeds, local = await asyncio.to_thread( _collect_facts_for_variant, query ) seed_hits.extend(seeds) extra_variant_lists.append(local) for chunk_id, _, body in extra_passages: if body and body not in source_passages: source_passages.append(body) if chunk_id not in provenance_ids: provenance_ids.append(chunk_id) source_passages = source_passages[:max_passages] candidates.extend(_merge_variant_candidate_lists(extra_variant_lists)) ranked = _dedupe_candidates(candidates) for candidate in ranked: if not candidate.get("session_ids"): candidate["session_ids"] = _session_ids_for_chunks( list(candidate.get("chunk_ids") or []), chunk_sessions ) curated = _diversify_facts( _rank_facts_with_completeness(ranked), max_facts=max_facts ) temporal_conflicts = _temporal_conflict_meta(curated) for candidate in curated: for chunk_id in candidate.get("chunk_ids") or []: if chunk_id not in provenance_ids: provenance_ids.append(chunk_id) missing_ids = [ cid for cid in provenance_ids[:80] if cid not in chunk_sessions ] if missing_ids: try: chunks, _ = await asyncio.to_thread( data_adapter.get_text_chunks_by_ids, missing_ids, False, request.brain_id, ) for chunk in chunks: body = (chunk.text or "").strip() chunk_id = str(getattr(chunk, "id", "") or "") if chunk_id and body: sessions = _session_ids_from_text(body) if sessions: chunk_sessions[chunk_id] = sessions except Exception: pass text_lines, triples, graph_session_ids = _build_fact_channel( curated, chunk_sessions ) response_paths = _paths_for_curated(bridge_paths, curated) with profile_stage("context.render", passages=len(source_passages)): topic_block = [ f"[topic] {t.get('label')}" for t in response_topics if isinstance(t, dict) and t.get("label") ] passage_block = [f"[passage] {p}" for p in source_passages] text_context = "\n".join(topic_block + passage_block + text_lines) return GetContextResponse( text_context=text_context, triples=triples, historical_context=historical_context, source_passages=source_passages, graph_session_ids=graph_session_ids or None, temporal_conflicts=temporal_conflicts or None, paths=response_paths or None, topics=response_topics or None, )