1267 lines
48 KiB
Python
1267 lines
48 KiB
Python
"""
|
|
Retrieval module for 4-way parallel search.
|
|
|
|
Implements:
|
|
1. Semantic retrieval (vector similarity)
|
|
2. BM25 retrieval (keyword/full-text search)
|
|
3. Graph retrieval (via pluggable GraphRetriever interface)
|
|
4. Temporal retrieval (time-aware search with spreading)
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from datetime import UTC, datetime
|
|
from typing import Optional
|
|
|
|
from ...config import get_config
|
|
from ..db_utils import acquire_with_retry
|
|
from ..memory_engine import fq_table
|
|
from .graph_retrieval import BFSGraphRetriever, GraphRetriever
|
|
from .link_expansion_retrieval import LinkExpansionRetriever
|
|
from .mpfp_retrieval import MPFPGraphRetriever
|
|
from .types import MPFPTimings, RetrievalResult
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ParallelRetrievalResult:
|
|
"""Result from parallel retrieval across all methods."""
|
|
|
|
semantic: list[RetrievalResult]
|
|
bm25: list[RetrievalResult]
|
|
graph: list[RetrievalResult]
|
|
temporal: list[RetrievalResult] | None
|
|
timings: dict[str, float] = field(default_factory=dict)
|
|
temporal_constraint: tuple | None = None # (start_date, end_date)
|
|
mpfp_timings: list[MPFPTimings] = field(default_factory=list) # MPFP sub-step timings per fact type
|
|
max_conn_wait: float = 0.0 # Maximum connection acquisition wait time across all methods
|
|
|
|
|
|
@dataclass
|
|
class MultiFactTypeRetrievalResult:
|
|
"""Result from retrieval across all fact types."""
|
|
|
|
# Results per fact type
|
|
results_by_fact_type: dict[str, ParallelRetrievalResult]
|
|
# Aggregate timings
|
|
timings: dict[str, float] = field(default_factory=dict)
|
|
# Max connection wait across all operations
|
|
max_conn_wait: float = 0.0
|
|
|
|
|
|
# Default graph retriever instance (can be overridden)
|
|
_default_graph_retriever: GraphRetriever | None = None
|
|
|
|
|
|
def get_default_graph_retriever() -> GraphRetriever:
|
|
"""Get or create the default graph retriever based on config."""
|
|
global _default_graph_retriever
|
|
if _default_graph_retriever is None:
|
|
config = get_config()
|
|
retriever_type = config.graph_retriever.lower()
|
|
if retriever_type == "mpfp":
|
|
_default_graph_retriever = MPFPGraphRetriever()
|
|
logger.info(
|
|
f"Using MPFP graph retriever (top_k_neighbors={_default_graph_retriever.config.top_k_neighbors})"
|
|
)
|
|
elif retriever_type == "bfs":
|
|
_default_graph_retriever = BFSGraphRetriever()
|
|
logger.info("Using BFS graph retriever")
|
|
elif retriever_type == "link_expansion":
|
|
_default_graph_retriever = LinkExpansionRetriever()
|
|
logger.info("Using LinkExpansion graph retriever")
|
|
else:
|
|
logger.warning(f"Unknown graph retriever '{retriever_type}', falling back to link_expansion")
|
|
_default_graph_retriever = LinkExpansionRetriever()
|
|
return _default_graph_retriever
|
|
|
|
|
|
def set_default_graph_retriever(retriever: GraphRetriever) -> None:
|
|
"""Set the default graph retriever (for configuration/testing)."""
|
|
global _default_graph_retriever
|
|
_default_graph_retriever = retriever
|
|
|
|
|
|
async def retrieve_semantic(
|
|
conn, query_emb_str: str, bank_id: str, fact_type: str, limit: int
|
|
) -> list[RetrievalResult]:
|
|
"""
|
|
Semantic retrieval via vector similarity.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
query_emb_str: Query embedding as string
|
|
agent_id: bank ID
|
|
fact_type: Fact type to filter
|
|
limit: Maximum results to return
|
|
|
|
Returns:
|
|
List of RetrievalResult objects
|
|
"""
|
|
results = await conn.fetch(
|
|
f"""
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
|
1 - (embedding <=> $1::vector) AS similarity
|
|
FROM {fq_table("memory_units")}
|
|
WHERE bank_id = $2
|
|
AND embedding IS NOT NULL
|
|
AND fact_type = $3
|
|
AND (1 - (embedding <=> $1::vector)) >= 0.3
|
|
ORDER BY embedding <=> $1::vector
|
|
LIMIT $4
|
|
""",
|
|
query_emb_str,
|
|
bank_id,
|
|
fact_type,
|
|
limit,
|
|
)
|
|
return [RetrievalResult.from_db_row(dict(r)) for r in results]
|
|
|
|
|
|
async def retrieve_bm25(conn, query_text: str, bank_id: str, fact_type: str, limit: int) -> list[RetrievalResult]:
|
|
"""
|
|
BM25 keyword retrieval via full-text search.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
query_text: Query text
|
|
agent_id: bank ID
|
|
fact_type: Fact type to filter
|
|
limit: Maximum results to return
|
|
|
|
Returns:
|
|
List of RetrievalResult objects
|
|
"""
|
|
import re
|
|
|
|
# Sanitize query text: remove special characters that have meaning in tsquery
|
|
# Keep only alphanumeric characters and spaces
|
|
sanitized_text = re.sub(r"[^\w\s]", " ", query_text.lower())
|
|
|
|
# Split and filter empty strings
|
|
tokens = [token for token in sanitized_text.split() if token]
|
|
|
|
if not tokens:
|
|
# If no valid tokens, return empty results
|
|
return []
|
|
|
|
# Convert query to tsquery using OR for more flexible matching
|
|
# This prevents empty results when some terms are missing
|
|
query_tsquery = " | ".join(tokens)
|
|
|
|
results = await conn.fetch(
|
|
f"""
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
|
ts_rank_cd(search_vector, to_tsquery('english', $1)) AS bm25_score
|
|
FROM {fq_table("memory_units")}
|
|
WHERE bank_id = $2
|
|
AND fact_type = $3
|
|
AND search_vector @@ to_tsquery('english', $1)
|
|
ORDER BY bm25_score DESC
|
|
LIMIT $4
|
|
""",
|
|
query_tsquery,
|
|
bank_id,
|
|
fact_type,
|
|
limit,
|
|
)
|
|
return [RetrievalResult.from_db_row(dict(r)) for r in results]
|
|
|
|
|
|
async def retrieve_semantic_bm25_combined(
|
|
conn,
|
|
query_emb_str: str,
|
|
query_text: str,
|
|
bank_id: str,
|
|
fact_types: list[str],
|
|
limit: int,
|
|
) -> dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]]:
|
|
"""
|
|
Combined semantic + BM25 retrieval for multiple fact types in a single query.
|
|
|
|
Uses CTEs with window functions to get top-N results per fact type per method,
|
|
all in one database round-trip.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
query_emb_str: Query embedding as string
|
|
query_text: Query text for BM25
|
|
bank_id: Bank ID
|
|
fact_types: List of fact types to retrieve
|
|
limit: Maximum results per method per fact type
|
|
|
|
Returns:
|
|
Dict mapping fact_type -> (semantic_results, bm25_results)
|
|
"""
|
|
import re
|
|
|
|
# Sanitize query text for BM25 (same as retrieve_bm25)
|
|
sanitized_text = re.sub(r"[^\w\s]", " ", query_text.lower())
|
|
tokens = [token for token in sanitized_text.split() if token]
|
|
|
|
# If no valid tokens for BM25, just run semantic
|
|
if not tokens:
|
|
results = await conn.fetch(
|
|
f"""
|
|
WITH semantic_ranked AS (
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
|
1 - (embedding <=> $1::vector) AS similarity,
|
|
NULL::float AS bm25_score,
|
|
'semantic' AS source,
|
|
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY embedding <=> $1::vector) AS rn
|
|
FROM {fq_table("memory_units")}
|
|
WHERE bank_id = $2
|
|
AND embedding IS NOT NULL
|
|
AND fact_type = ANY($3)
|
|
AND (1 - (embedding <=> $1::vector)) >= 0.3
|
|
)
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
|
similarity, bm25_score, source
|
|
FROM semantic_ranked
|
|
WHERE rn <= $4
|
|
""",
|
|
query_emb_str,
|
|
bank_id,
|
|
fact_types,
|
|
limit,
|
|
)
|
|
# Group by fact_type
|
|
result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {
|
|
ft: ([], []) for ft in fact_types
|
|
}
|
|
for r in results:
|
|
row = dict(r)
|
|
ft = row.get("fact_type")
|
|
row.pop("source", None)
|
|
if ft in result_dict:
|
|
result_dict[ft][0].append(RetrievalResult.from_db_row(row))
|
|
return result_dict
|
|
|
|
query_tsquery = " | ".join(tokens)
|
|
|
|
# Combined CTE query for both semantic and BM25 across all fact types
|
|
# Uses window functions to limit per fact_type per method
|
|
results = await conn.fetch(
|
|
f"""
|
|
WITH semantic_ranked AS (
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
|
1 - (embedding <=> $1::vector) AS similarity,
|
|
NULL::float AS bm25_score,
|
|
'semantic' AS source,
|
|
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY embedding <=> $1::vector) AS rn
|
|
FROM {fq_table("memory_units")}
|
|
WHERE bank_id = $2
|
|
AND embedding IS NOT NULL
|
|
AND fact_type = ANY($3)
|
|
AND (1 - (embedding <=> $1::vector)) >= 0.3
|
|
),
|
|
bm25_ranked AS (
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
|
NULL::float AS similarity,
|
|
ts_rank_cd(search_vector, to_tsquery('english', $5)) AS bm25_score,
|
|
'bm25' AS source,
|
|
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY ts_rank_cd(search_vector, to_tsquery('english', $5)) DESC) AS rn
|
|
FROM {fq_table("memory_units")}
|
|
WHERE bank_id = $2
|
|
AND fact_type = ANY($3)
|
|
AND search_vector @@ to_tsquery('english', $5)
|
|
),
|
|
semantic AS (
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
|
similarity, bm25_score, source
|
|
FROM semantic_ranked WHERE rn <= $4
|
|
),
|
|
bm25 AS (
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
|
similarity, bm25_score, source
|
|
FROM bm25_ranked WHERE rn <= $4
|
|
)
|
|
SELECT * FROM semantic
|
|
UNION ALL
|
|
SELECT * FROM bm25
|
|
""",
|
|
query_emb_str,
|
|
bank_id,
|
|
fact_types,
|
|
limit,
|
|
query_tsquery,
|
|
)
|
|
|
|
# Group results by fact_type and source
|
|
result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {ft: ([], []) for ft in fact_types}
|
|
for r in results:
|
|
row = dict(r)
|
|
source = row.pop("source", None)
|
|
ft = row.get("fact_type")
|
|
if ft in result_dict:
|
|
if source == "semantic":
|
|
result_dict[ft][0].append(RetrievalResult.from_db_row(row))
|
|
else:
|
|
result_dict[ft][1].append(RetrievalResult.from_db_row(row))
|
|
|
|
return result_dict
|
|
|
|
|
|
async def retrieve_temporal_combined(
|
|
conn,
|
|
query_emb_str: str,
|
|
bank_id: str,
|
|
fact_types: list[str],
|
|
start_date: datetime,
|
|
end_date: datetime,
|
|
budget: int,
|
|
semantic_threshold: float = 0.1,
|
|
) -> dict[str, list[RetrievalResult]]:
|
|
"""
|
|
Temporal retrieval for multiple fact types in a single query.
|
|
|
|
Batches the entry point query using window functions to get top-N per fact type,
|
|
then runs spreading for each fact type.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
query_emb_str: Query embedding as string
|
|
bank_id: Bank ID
|
|
fact_types: List of fact types to retrieve
|
|
start_date: Start of time range
|
|
end_date: End of time range
|
|
budget: Node budget for spreading per fact type
|
|
semantic_threshold: Minimum semantic similarity to include
|
|
|
|
Returns:
|
|
Dict mapping fact_type -> list of RetrievalResult
|
|
"""
|
|
from ..memory_engine import fq_table
|
|
|
|
# Ensure dates are timezone-aware
|
|
if start_date.tzinfo is None:
|
|
start_date = start_date.replace(tzinfo=UTC)
|
|
if end_date.tzinfo is None:
|
|
end_date = end_date.replace(tzinfo=UTC)
|
|
|
|
# Batch query: Get entry points for ALL fact types at once with window function
|
|
entry_points = await conn.fetch(
|
|
f"""
|
|
WITH ranked_entries AS (
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
|
1 - (embedding <=> $1::vector) AS similarity,
|
|
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC, embedding <=> $1::vector) AS rn
|
|
FROM {fq_table("memory_units")}
|
|
WHERE bank_id = $2
|
|
AND fact_type = ANY($3)
|
|
AND embedding IS NOT NULL
|
|
AND (
|
|
(occurred_start IS NOT NULL AND occurred_end IS NOT NULL
|
|
AND occurred_start <= $5 AND occurred_end >= $4)
|
|
OR
|
|
(mentioned_at IS NOT NULL AND mentioned_at BETWEEN $4 AND $5)
|
|
OR
|
|
(occurred_start IS NOT NULL AND occurred_start BETWEEN $4 AND $5)
|
|
OR
|
|
(occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
|
|
)
|
|
AND (1 - (embedding <=> $1::vector)) >= $6
|
|
)
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id, similarity
|
|
FROM ranked_entries
|
|
WHERE rn <= 10
|
|
""",
|
|
query_emb_str,
|
|
bank_id,
|
|
fact_types,
|
|
start_date,
|
|
end_date,
|
|
semantic_threshold,
|
|
)
|
|
|
|
if not entry_points:
|
|
return {ft: [] for ft in fact_types}
|
|
|
|
# Group entry points by fact type
|
|
entries_by_ft: dict[str, list] = {ft: [] for ft in fact_types}
|
|
for ep in entry_points:
|
|
ft = ep["fact_type"]
|
|
if ft in entries_by_ft:
|
|
entries_by_ft[ft].append(ep)
|
|
|
|
# Calculate shared temporal parameters
|
|
total_days = (end_date - start_date).total_seconds() / 86400
|
|
mid_date = start_date + (end_date - start_date) / 2
|
|
|
|
# Process each fact type (spreading needs to stay per fact type due to link filtering)
|
|
results_by_ft: dict[str, list[RetrievalResult]] = {}
|
|
|
|
for ft in fact_types:
|
|
ft_entry_points = entries_by_ft.get(ft, [])
|
|
if not ft_entry_points:
|
|
results_by_ft[ft] = []
|
|
continue
|
|
|
|
results = []
|
|
visited = set()
|
|
node_scores = {}
|
|
|
|
# Process entry points
|
|
for ep in ft_entry_points:
|
|
unit_id = str(ep["id"])
|
|
visited.add(unit_id)
|
|
|
|
# Calculate temporal proximity
|
|
best_date = None
|
|
if ep["occurred_start"] is not None and ep["occurred_end"] is not None:
|
|
best_date = ep["occurred_start"] + (ep["occurred_end"] - ep["occurred_start"]) / 2
|
|
elif ep["occurred_start"] is not None:
|
|
best_date = ep["occurred_start"]
|
|
elif ep["occurred_end"] is not None:
|
|
best_date = ep["occurred_end"]
|
|
elif ep["mentioned_at"] is not None:
|
|
best_date = ep["mentioned_at"]
|
|
|
|
if best_date:
|
|
days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
|
|
temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
|
else:
|
|
temporal_proximity = 0.5
|
|
|
|
ep_result = RetrievalResult.from_db_row(dict(ep))
|
|
ep_result.temporal_score = temporal_proximity
|
|
ep_result.temporal_proximity = temporal_proximity
|
|
results.append(ep_result)
|
|
node_scores[unit_id] = (ep["similarity"], 1.0)
|
|
|
|
# Spreading through temporal links (same as single-fact-type version)
|
|
frontier = list(node_scores.keys())
|
|
budget_remaining = budget - len(ft_entry_points)
|
|
batch_size = 20
|
|
|
|
while frontier and budget_remaining > 0:
|
|
batch_ids = frontier[:batch_size]
|
|
frontier = frontier[batch_size:]
|
|
|
|
neighbors = await conn.fetch(
|
|
f"""
|
|
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type, mu.document_id, mu.chunk_id,
|
|
ml.weight, ml.link_type, ml.from_unit_id,
|
|
1 - (mu.embedding <=> $1::vector) AS similarity
|
|
FROM {fq_table("memory_links")} ml
|
|
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
|
WHERE ml.from_unit_id = ANY($2::uuid[])
|
|
AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
|
|
AND ml.weight >= 0.1
|
|
AND mu.fact_type = $3
|
|
AND mu.embedding IS NOT NULL
|
|
AND (1 - (mu.embedding <=> $1::vector)) >= $4
|
|
ORDER BY ml.weight DESC
|
|
LIMIT $5
|
|
""",
|
|
query_emb_str,
|
|
batch_ids,
|
|
ft,
|
|
semantic_threshold,
|
|
batch_size * 10,
|
|
)
|
|
|
|
for n in neighbors:
|
|
neighbor_id = str(n["id"])
|
|
if neighbor_id in visited:
|
|
continue
|
|
|
|
visited.add(neighbor_id)
|
|
budget_remaining -= 1
|
|
|
|
parent_id = str(n["from_unit_id"])
|
|
_, parent_temporal_score = node_scores.get(parent_id, (0.5, 0.5))
|
|
|
|
neighbor_best_date = None
|
|
if n["occurred_start"] is not None and n["occurred_end"] is not None:
|
|
neighbor_best_date = n["occurred_start"] + (n["occurred_end"] - n["occurred_start"]) / 2
|
|
elif n["occurred_start"] is not None:
|
|
neighbor_best_date = n["occurred_start"]
|
|
elif n["occurred_end"] is not None:
|
|
neighbor_best_date = n["occurred_end"]
|
|
elif n["mentioned_at"] is not None:
|
|
neighbor_best_date = n["mentioned_at"]
|
|
|
|
if neighbor_best_date:
|
|
days_from_mid = abs((neighbor_best_date - mid_date).total_seconds() / 86400)
|
|
neighbor_temporal_proximity = (
|
|
1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
|
)
|
|
else:
|
|
neighbor_temporal_proximity = 0.3
|
|
|
|
link_type = n["link_type"]
|
|
if link_type in ("causes", "caused_by"):
|
|
causal_boost = 2.0
|
|
elif link_type in ("enables", "prevents"):
|
|
causal_boost = 1.5
|
|
else:
|
|
causal_boost = 1.0
|
|
|
|
propagated_temporal = parent_temporal_score * n["weight"] * causal_boost * 0.7
|
|
combined_temporal = max(neighbor_temporal_proximity, propagated_temporal)
|
|
|
|
neighbor_result = RetrievalResult.from_db_row(dict(n))
|
|
neighbor_result.temporal_score = combined_temporal
|
|
neighbor_result.temporal_proximity = neighbor_temporal_proximity
|
|
results.append(neighbor_result)
|
|
|
|
if budget_remaining > 0 and combined_temporal > 0.2:
|
|
node_scores[neighbor_id] = (n["similarity"], combined_temporal)
|
|
frontier.append(neighbor_id)
|
|
|
|
if budget_remaining <= 0:
|
|
break
|
|
|
|
results_by_ft[ft] = results
|
|
|
|
return results_by_ft
|
|
|
|
|
|
async def retrieve_temporal(
|
|
conn,
|
|
query_emb_str: str,
|
|
bank_id: str,
|
|
fact_type: str,
|
|
start_date: datetime,
|
|
end_date: datetime,
|
|
budget: int,
|
|
semantic_threshold: float = 0.1,
|
|
) -> list[RetrievalResult]:
|
|
"""
|
|
Temporal retrieval with spreading activation.
|
|
|
|
Strategy:
|
|
1. Find entry points (facts in date range with semantic relevance)
|
|
2. Spread through temporal links to related facts
|
|
3. Score by temporal proximity + semantic similarity + link weight
|
|
|
|
Args:
|
|
conn: Database connection
|
|
query_emb_str: Query embedding as string
|
|
agent_id: bank ID
|
|
fact_type: Fact type to filter
|
|
start_date: Start of time range
|
|
end_date: End of time range
|
|
budget: Node budget for spreading
|
|
semantic_threshold: Minimum semantic similarity to include
|
|
|
|
Returns:
|
|
List of RetrievalResult objects with temporal scores
|
|
"""
|
|
|
|
# Ensure start_date and end_date are timezone-aware (UTC) to match database datetimes
|
|
if start_date.tzinfo is None:
|
|
start_date = start_date.replace(tzinfo=UTC)
|
|
if end_date.tzinfo is None:
|
|
end_date = end_date.replace(tzinfo=UTC)
|
|
|
|
entry_points = await conn.fetch(
|
|
f"""
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
|
1 - (embedding <=> $1::vector) AS similarity
|
|
FROM {fq_table("memory_units")}
|
|
WHERE bank_id = $2
|
|
AND fact_type = $3
|
|
AND embedding IS NOT NULL
|
|
AND (
|
|
-- Match if occurred range overlaps with query range
|
|
(occurred_start IS NOT NULL AND occurred_end IS NOT NULL
|
|
AND occurred_start <= $5 AND occurred_end >= $4)
|
|
OR
|
|
-- Match if mentioned_at falls within query range
|
|
(mentioned_at IS NOT NULL AND mentioned_at BETWEEN $4 AND $5)
|
|
OR
|
|
-- Match if any occurred date is set and overlaps (even if only start or end is set)
|
|
(occurred_start IS NOT NULL AND occurred_start BETWEEN $4 AND $5)
|
|
OR
|
|
(occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
|
|
)
|
|
AND (1 - (embedding <=> $1::vector)) >= $6
|
|
ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC, (embedding <=> $1::vector) ASC
|
|
LIMIT 10
|
|
""",
|
|
query_emb_str,
|
|
bank_id,
|
|
fact_type,
|
|
start_date,
|
|
end_date,
|
|
semantic_threshold,
|
|
)
|
|
|
|
if not entry_points:
|
|
return []
|
|
|
|
# Calculate temporal scores for entry points
|
|
total_days = (end_date - start_date).total_seconds() / 86400
|
|
mid_date = start_date + (end_date - start_date) / 2 # Calculate once for all comparisons
|
|
results = []
|
|
visited = set()
|
|
|
|
for ep in entry_points:
|
|
unit_id = str(ep["id"])
|
|
visited.add(unit_id)
|
|
|
|
# Calculate temporal proximity using the most relevant date
|
|
# Priority: occurred_start/end (event time) > mentioned_at (mention time)
|
|
best_date = None
|
|
if ep["occurred_start"] is not None and ep["occurred_end"] is not None:
|
|
# Use midpoint of occurred range
|
|
best_date = ep["occurred_start"] + (ep["occurred_end"] - ep["occurred_start"]) / 2
|
|
elif ep["occurred_start"] is not None:
|
|
best_date = ep["occurred_start"]
|
|
elif ep["occurred_end"] is not None:
|
|
best_date = ep["occurred_end"]
|
|
elif ep["mentioned_at"] is not None:
|
|
best_date = ep["mentioned_at"]
|
|
|
|
# Temporal proximity score (closer to range center = higher score)
|
|
if best_date:
|
|
days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
|
|
temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
|
else:
|
|
temporal_proximity = 0.5 # Fallback if no dates (shouldn't happen due to WHERE clause)
|
|
|
|
# Create RetrievalResult with temporal scores
|
|
ep_result = RetrievalResult.from_db_row(dict(ep))
|
|
ep_result.temporal_score = temporal_proximity
|
|
ep_result.temporal_proximity = temporal_proximity
|
|
results.append(ep_result)
|
|
|
|
# Spread through temporal links using BATCHED neighbor fetching
|
|
# Map node_id -> (semantic_sim, temporal_score) for propagation
|
|
node_scores = {str(ep["id"]): (ep["similarity"], 1.0) for ep in entry_points}
|
|
frontier = list(node_scores.keys()) # Current batch of nodes to expand
|
|
budget_remaining = budget - len(entry_points)
|
|
batch_size = 20 # Process this many nodes per DB query
|
|
|
|
while frontier and budget_remaining > 0:
|
|
# Take a batch from frontier
|
|
batch_ids = frontier[:batch_size]
|
|
frontier = frontier[batch_size:]
|
|
|
|
# Batch fetch all neighbors for this batch of nodes
|
|
neighbors = await conn.fetch(
|
|
f"""
|
|
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type, mu.document_id, mu.chunk_id,
|
|
ml.weight, ml.link_type, ml.from_unit_id,
|
|
1 - (mu.embedding <=> $1::vector) AS similarity
|
|
FROM {fq_table("memory_links")} ml
|
|
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
|
WHERE ml.from_unit_id = ANY($2::uuid[])
|
|
AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
|
|
AND ml.weight >= 0.1
|
|
AND mu.fact_type = $3
|
|
AND mu.embedding IS NOT NULL
|
|
AND (1 - (mu.embedding <=> $1::vector)) >= $4
|
|
ORDER BY ml.weight DESC
|
|
LIMIT $5
|
|
""",
|
|
query_emb_str,
|
|
batch_ids,
|
|
fact_type,
|
|
semantic_threshold,
|
|
batch_size * 10, # Allow up to 10 neighbors per node in batch
|
|
)
|
|
|
|
for n in neighbors:
|
|
neighbor_id = str(n["id"])
|
|
if neighbor_id in visited:
|
|
continue
|
|
|
|
visited.add(neighbor_id)
|
|
budget_remaining -= 1
|
|
|
|
# Get parent's scores for propagation
|
|
parent_id = str(n["from_unit_id"])
|
|
_, parent_temporal_score = node_scores.get(parent_id, (0.5, 0.5))
|
|
|
|
# Calculate temporal score for neighbor using best available date
|
|
neighbor_best_date = None
|
|
if n["occurred_start"] is not None and n["occurred_end"] is not None:
|
|
neighbor_best_date = n["occurred_start"] + (n["occurred_end"] - n["occurred_start"]) / 2
|
|
elif n["occurred_start"] is not None:
|
|
neighbor_best_date = n["occurred_start"]
|
|
elif n["occurred_end"] is not None:
|
|
neighbor_best_date = n["occurred_end"]
|
|
elif n["mentioned_at"] is not None:
|
|
neighbor_best_date = n["mentioned_at"]
|
|
|
|
if neighbor_best_date:
|
|
days_from_mid = abs((neighbor_best_date - mid_date).total_seconds() / 86400)
|
|
neighbor_temporal_proximity = (
|
|
1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
|
)
|
|
else:
|
|
neighbor_temporal_proximity = 0.3 # Lower score if no temporal data
|
|
|
|
# Boost causal links (same as graph retrieval)
|
|
link_type = n["link_type"]
|
|
if link_type in ("causes", "caused_by"):
|
|
causal_boost = 2.0
|
|
elif link_type in ("enables", "prevents"):
|
|
causal_boost = 1.5
|
|
else:
|
|
causal_boost = 1.0
|
|
|
|
# Propagate temporal score through links (decay, with causal boost)
|
|
propagated_temporal = parent_temporal_score * n["weight"] * causal_boost * 0.7
|
|
|
|
# Combined temporal score
|
|
combined_temporal = max(neighbor_temporal_proximity, propagated_temporal)
|
|
|
|
# Create RetrievalResult with temporal scores
|
|
neighbor_result = RetrievalResult.from_db_row(dict(n))
|
|
neighbor_result.temporal_score = combined_temporal
|
|
neighbor_result.temporal_proximity = neighbor_temporal_proximity
|
|
results.append(neighbor_result)
|
|
|
|
# Track scores for propagation and add to frontier
|
|
if budget_remaining > 0 and combined_temporal > 0.2:
|
|
node_scores[neighbor_id] = (n["similarity"], combined_temporal)
|
|
frontier.append(neighbor_id)
|
|
|
|
if budget_remaining <= 0:
|
|
break
|
|
|
|
return results
|
|
|
|
|
|
async def retrieve_parallel(
|
|
pool,
|
|
query_text: str,
|
|
query_embedding_str: str,
|
|
bank_id: str,
|
|
fact_type: str,
|
|
thinking_budget: int,
|
|
question_date: datetime | None = None,
|
|
query_analyzer: Optional["QueryAnalyzer"] = None,
|
|
graph_retriever: GraphRetriever | None = None,
|
|
temporal_constraint: tuple | None = None, # Pre-extracted temporal constraint
|
|
) -> ParallelRetrievalResult:
|
|
"""
|
|
Run 3-way or 4-way parallel retrieval (adds temporal if detected).
|
|
|
|
Args:
|
|
pool: Database connection pool
|
|
query_text: Query text
|
|
query_embedding_str: Query embedding as string
|
|
bank_id: Bank ID
|
|
fact_type: Fact type to filter
|
|
thinking_budget: Budget for graph traversal and retrieval limits
|
|
question_date: Optional date when question was asked (for temporal filtering)
|
|
query_analyzer: Query analyzer to use (defaults to TransformerQueryAnalyzer)
|
|
graph_retriever: Graph retrieval strategy (defaults to configured retriever)
|
|
temporal_constraint: Pre-extracted temporal constraint (optional)
|
|
|
|
Returns:
|
|
ParallelRetrievalResult with semantic, bm25, graph, temporal results and timings
|
|
"""
|
|
retriever = graph_retriever or get_default_graph_retriever()
|
|
|
|
# Use optimized parallel path for MPFP and LinkExpansion (runs all methods truly in parallel)
|
|
# BFS uses legacy path that extracts temporal constraint upfront
|
|
if retriever.name in ("mpfp", "link_expansion"):
|
|
return await _retrieve_parallel_mpfp(
|
|
pool,
|
|
query_text,
|
|
query_embedding_str,
|
|
bank_id,
|
|
fact_type,
|
|
thinking_budget,
|
|
temporal_constraint,
|
|
retriever,
|
|
question_date,
|
|
query_analyzer,
|
|
)
|
|
else:
|
|
# For BFS, extract temporal constraint upfront (legacy path)
|
|
if temporal_constraint is None:
|
|
from .temporal_extraction import extract_temporal_constraint
|
|
|
|
temporal_constraint = extract_temporal_constraint(
|
|
query_text, reference_date=question_date, analyzer=query_analyzer
|
|
)
|
|
return await _retrieve_parallel_bfs(
|
|
pool, query_text, query_embedding_str, bank_id, fact_type, thinking_budget, temporal_constraint, retriever
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class _TimedResult:
|
|
"""Internal result with timing."""
|
|
|
|
results: list[RetrievalResult]
|
|
time: float
|
|
conn_wait: float = 0.0 # Connection acquisition wait time
|
|
|
|
|
|
async def _retrieve_parallel_mpfp(
|
|
pool,
|
|
query_text: str,
|
|
query_embedding_str: str,
|
|
bank_id: str,
|
|
fact_type: str,
|
|
thinking_budget: int,
|
|
temporal_constraint: tuple | None,
|
|
retriever: GraphRetriever,
|
|
question_date: datetime | None = None,
|
|
query_analyzer=None,
|
|
) -> ParallelRetrievalResult:
|
|
"""
|
|
MPFP retrieval with true parallelization.
|
|
|
|
All methods run independently in parallel:
|
|
- Semantic: vector similarity search
|
|
- BM25: keyword search
|
|
- Graph: MPFP traversal (does its own semantic seeds internally)
|
|
- Temporal: date extraction (if needed) + date-range search
|
|
|
|
Temporal extraction runs IN PARALLEL with other retrievals, so even if
|
|
dateparser is slow, it doesn't block semantic/BM25/graph.
|
|
"""
|
|
import time
|
|
|
|
async def run_semantic() -> _TimedResult:
|
|
"""Independent semantic retrieval."""
|
|
start = time.time()
|
|
acquire_start = time.time()
|
|
async with acquire_with_retry(pool) as conn:
|
|
conn_wait = time.time() - acquire_start
|
|
results = await retrieve_semantic(conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget)
|
|
return _TimedResult(results, time.time() - start, conn_wait)
|
|
|
|
async def run_bm25() -> _TimedResult:
|
|
"""Independent BM25 retrieval."""
|
|
start = time.time()
|
|
acquire_start = time.time()
|
|
async with acquire_with_retry(pool) as conn:
|
|
conn_wait = time.time() - acquire_start
|
|
results = await retrieve_bm25(conn, query_text, bank_id, fact_type, limit=thinking_budget)
|
|
return _TimedResult(results, time.time() - start, conn_wait)
|
|
|
|
async def run_graph() -> tuple[list[RetrievalResult], float, MPFPTimings | None]:
|
|
"""Independent graph retrieval - does its own semantic seeds."""
|
|
start = time.time()
|
|
|
|
# MPFP does its own semantic seeds via _find_semantic_seeds
|
|
# Note: temporal_seeds not used here to avoid dependency on temporal extraction
|
|
results, mpfp_timing = await retriever.retrieve(
|
|
pool=pool,
|
|
query_embedding_str=query_embedding_str,
|
|
bank_id=bank_id,
|
|
fact_type=fact_type,
|
|
budget=thinking_budget,
|
|
query_text=query_text,
|
|
semantic_seeds=None, # Let MPFP find its own seeds
|
|
temporal_seeds=None, # Don't wait for temporal extraction
|
|
)
|
|
return results, time.time() - start, mpfp_timing
|
|
|
|
@dataclass
|
|
class _TemporalWithConstraint:
|
|
"""Temporal results with the extracted constraint."""
|
|
|
|
results: list[RetrievalResult]
|
|
time: float
|
|
constraint: tuple | None
|
|
extraction_time: float # Time spent in query analyzer (dateparser)
|
|
conn_wait: float = 0.0 # Connection acquisition wait time
|
|
|
|
async def run_temporal_with_extraction() -> _TemporalWithConstraint:
|
|
"""
|
|
Extract temporal constraint AND run temporal retrieval.
|
|
|
|
This runs in parallel with semantic/BM25/graph, so dateparser
|
|
latency doesn't block other retrievals.
|
|
"""
|
|
start = time.time()
|
|
|
|
# Use pre-provided constraint if available
|
|
tc = temporal_constraint
|
|
extraction_time = 0.0
|
|
|
|
# Otherwise extract from query (this is the potentially slow dateparser call)
|
|
if tc is None:
|
|
from .temporal_extraction import extract_temporal_constraint
|
|
|
|
extraction_start = time.time()
|
|
tc = extract_temporal_constraint(query_text, reference_date=question_date, analyzer=query_analyzer)
|
|
extraction_time = time.time() - extraction_start
|
|
|
|
# If no temporal constraint found, return empty (but still report extraction time)
|
|
if tc is None:
|
|
return _TemporalWithConstraint([], time.time() - start, None, extraction_time, 0.0)
|
|
|
|
# Run temporal retrieval with the extracted constraint
|
|
tc_start, tc_end = tc
|
|
acquire_start = time.time()
|
|
async with acquire_with_retry(pool) as conn:
|
|
conn_wait = time.time() - acquire_start
|
|
results = await retrieve_temporal(
|
|
conn,
|
|
query_embedding_str,
|
|
bank_id,
|
|
fact_type,
|
|
tc_start,
|
|
tc_end,
|
|
budget=thinking_budget,
|
|
semantic_threshold=0.1,
|
|
)
|
|
return _TemporalWithConstraint(results, time.time() - start, tc, extraction_time, conn_wait)
|
|
|
|
# Run ALL methods in parallel (including temporal extraction!)
|
|
semantic_result, bm25_result, graph_result, temporal_result = await asyncio.gather(
|
|
run_semantic(),
|
|
run_bm25(),
|
|
run_graph(),
|
|
run_temporal_with_extraction(),
|
|
)
|
|
graph_results, graph_time, mpfp_timing = graph_result
|
|
|
|
# Compute max connection wait across all methods (graph handles its own connections)
|
|
max_conn_wait = max(semantic_result.conn_wait, bm25_result.conn_wait, temporal_result.conn_wait)
|
|
|
|
return ParallelRetrievalResult(
|
|
semantic=semantic_result.results,
|
|
bm25=bm25_result.results,
|
|
graph=graph_results,
|
|
temporal=temporal_result.results if temporal_result.results else None,
|
|
timings={
|
|
"semantic": semantic_result.time,
|
|
"bm25": bm25_result.time,
|
|
"graph": graph_time,
|
|
"temporal": temporal_result.time,
|
|
"temporal_extraction": temporal_result.extraction_time,
|
|
},
|
|
temporal_constraint=temporal_result.constraint,
|
|
mpfp_timings=[mpfp_timing] if mpfp_timing else [],
|
|
max_conn_wait=max_conn_wait,
|
|
)
|
|
|
|
|
|
async def _get_temporal_entry_points(
|
|
conn,
|
|
query_embedding_str: str,
|
|
bank_id: str,
|
|
fact_type: str,
|
|
start_date: datetime,
|
|
end_date: datetime,
|
|
limit: int = 20,
|
|
semantic_threshold: float = 0.1,
|
|
) -> list[RetrievalResult]:
|
|
"""Get temporal entry points (facts in date range with semantic relevance)."""
|
|
|
|
if start_date.tzinfo is None:
|
|
start_date = start_date.replace(tzinfo=UTC)
|
|
if end_date.tzinfo is None:
|
|
end_date = end_date.replace(tzinfo=UTC)
|
|
|
|
rows = await conn.fetch(
|
|
f"""
|
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
|
|
access_count, embedding, fact_type, document_id, chunk_id,
|
|
1 - (embedding <=> $1::vector) AS similarity
|
|
FROM {fq_table("memory_units")}
|
|
WHERE bank_id = $2
|
|
AND fact_type = $3
|
|
AND embedding IS NOT NULL
|
|
AND (
|
|
(occurred_start IS NOT NULL AND occurred_end IS NOT NULL
|
|
AND occurred_start <= $5 AND occurred_end >= $4)
|
|
OR (mentioned_at IS NOT NULL AND mentioned_at BETWEEN $4 AND $5)
|
|
OR (occurred_start IS NOT NULL AND occurred_start BETWEEN $4 AND $5)
|
|
OR (occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
|
|
)
|
|
AND (1 - (embedding <=> $1::vector)) >= $6
|
|
ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC,
|
|
(embedding <=> $1::vector) ASC
|
|
LIMIT $7
|
|
""",
|
|
query_embedding_str,
|
|
bank_id,
|
|
fact_type,
|
|
start_date,
|
|
end_date,
|
|
semantic_threshold,
|
|
limit,
|
|
)
|
|
|
|
results = []
|
|
total_days = max((end_date - start_date).total_seconds() / 86400, 1)
|
|
mid_date = start_date + (end_date - start_date) / 2
|
|
|
|
for row in rows:
|
|
result = RetrievalResult.from_db_row(dict(row))
|
|
|
|
# Calculate temporal proximity score
|
|
best_date = None
|
|
if row["occurred_start"] and row["occurred_end"]:
|
|
best_date = row["occurred_start"] + (row["occurred_end"] - row["occurred_start"]) / 2
|
|
elif row["occurred_start"]:
|
|
best_date = row["occurred_start"]
|
|
elif row["occurred_end"]:
|
|
best_date = row["occurred_end"]
|
|
elif row["mentioned_at"]:
|
|
best_date = row["mentioned_at"]
|
|
|
|
if best_date:
|
|
days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
|
|
result.temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0)
|
|
else:
|
|
result.temporal_proximity = 0.5
|
|
|
|
result.temporal_score = result.temporal_proximity
|
|
results.append(result)
|
|
|
|
return results
|
|
|
|
|
|
async def _retrieve_parallel_bfs(
|
|
pool,
|
|
query_text: str,
|
|
query_embedding_str: str,
|
|
bank_id: str,
|
|
fact_type: str,
|
|
thinking_budget: int,
|
|
temporal_constraint: tuple | None,
|
|
retriever: GraphRetriever,
|
|
) -> ParallelRetrievalResult:
|
|
"""BFS retrieval: all methods run in parallel (original behavior)."""
|
|
import time
|
|
|
|
async def run_semantic() -> _TimedResult:
|
|
start = time.time()
|
|
async with acquire_with_retry(pool) as conn:
|
|
results = await retrieve_semantic(conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget)
|
|
return _TimedResult(results, time.time() - start)
|
|
|
|
async def run_bm25() -> _TimedResult:
|
|
start = time.time()
|
|
async with acquire_with_retry(pool) as conn:
|
|
results = await retrieve_bm25(conn, query_text, bank_id, fact_type, limit=thinking_budget)
|
|
return _TimedResult(results, time.time() - start)
|
|
|
|
async def run_graph() -> _TimedResult:
|
|
start = time.time()
|
|
results, _ = await retriever.retrieve(
|
|
pool=pool,
|
|
query_embedding_str=query_embedding_str,
|
|
bank_id=bank_id,
|
|
fact_type=fact_type,
|
|
budget=thinking_budget,
|
|
query_text=query_text,
|
|
)
|
|
return _TimedResult(results, time.time() - start)
|
|
|
|
async def run_temporal(tc_start, tc_end) -> _TimedResult:
|
|
start = time.time()
|
|
async with acquire_with_retry(pool) as conn:
|
|
results = await retrieve_temporal(
|
|
conn,
|
|
query_embedding_str,
|
|
bank_id,
|
|
fact_type,
|
|
tc_start,
|
|
tc_end,
|
|
budget=thinking_budget,
|
|
semantic_threshold=0.1,
|
|
)
|
|
return _TimedResult(results, time.time() - start)
|
|
|
|
if temporal_constraint:
|
|
tc_start, tc_end = temporal_constraint
|
|
semantic_r, bm25_r, graph_r, temporal_r = await asyncio.gather(
|
|
run_semantic(),
|
|
run_bm25(),
|
|
run_graph(),
|
|
run_temporal(tc_start, tc_end),
|
|
)
|
|
return ParallelRetrievalResult(
|
|
semantic=semantic_r.results,
|
|
bm25=bm25_r.results,
|
|
graph=graph_r.results,
|
|
temporal=temporal_r.results,
|
|
timings={
|
|
"semantic": semantic_r.time,
|
|
"bm25": bm25_r.time,
|
|
"graph": graph_r.time,
|
|
"temporal": temporal_r.time,
|
|
},
|
|
temporal_constraint=temporal_constraint,
|
|
)
|
|
else:
|
|
semantic_r, bm25_r, graph_r = await asyncio.gather(
|
|
run_semantic(),
|
|
run_bm25(),
|
|
run_graph(),
|
|
)
|
|
return ParallelRetrievalResult(
|
|
semantic=semantic_r.results,
|
|
bm25=bm25_r.results,
|
|
graph=graph_r.results,
|
|
temporal=None,
|
|
timings={
|
|
"semantic": semantic_r.time,
|
|
"bm25": bm25_r.time,
|
|
"graph": graph_r.time,
|
|
},
|
|
temporal_constraint=None,
|
|
)
|
|
|
|
|
|
async def retrieve_all_fact_types_parallel(
|
|
pool,
|
|
query_text: str,
|
|
query_embedding_str: str,
|
|
bank_id: str,
|
|
fact_types: list[str],
|
|
thinking_budget: int,
|
|
question_date: datetime | None = None,
|
|
query_analyzer: Optional["QueryAnalyzer"] = None,
|
|
graph_retriever: GraphRetriever | None = None,
|
|
) -> MultiFactTypeRetrievalResult:
|
|
"""
|
|
Optimized retrieval for multiple fact types using batched queries.
|
|
|
|
This reduces database round-trips by:
|
|
1. Combining semantic + BM25 into one CTE query for ALL fact types (1 query instead of 2N)
|
|
2. Running graph retrieval per fact type in parallel (N parallel tasks)
|
|
3. Running temporal retrieval per fact type in parallel (N parallel tasks)
|
|
|
|
Args:
|
|
pool: Database connection pool
|
|
query_text: Query text
|
|
query_embedding_str: Query embedding as string
|
|
bank_id: Bank ID
|
|
fact_types: List of fact types to retrieve
|
|
thinking_budget: Budget for graph traversal and retrieval limits
|
|
question_date: Optional date when question was asked (for temporal filtering)
|
|
query_analyzer: Query analyzer to use (defaults to TransformerQueryAnalyzer)
|
|
graph_retriever: Graph retrieval strategy (defaults to configured retriever)
|
|
|
|
Returns:
|
|
MultiFactTypeRetrievalResult with results organized by fact type
|
|
"""
|
|
import time
|
|
|
|
retriever = graph_retriever or get_default_graph_retriever()
|
|
start_time = time.time()
|
|
timings: dict[str, float] = {}
|
|
|
|
# Step 1: Extract temporal constraint first (CPU work, no DB)
|
|
# Do this before DB queries so we know if we need temporal retrieval
|
|
temporal_extraction_start = time.time()
|
|
from .temporal_extraction import extract_temporal_constraint
|
|
|
|
temporal_constraint = extract_temporal_constraint(query_text, reference_date=question_date, analyzer=query_analyzer)
|
|
temporal_extraction_time = time.time() - temporal_extraction_start
|
|
timings["temporal_extraction"] = temporal_extraction_time
|
|
|
|
# Step 2: Run semantic + BM25 + temporal combined in ONE connection!
|
|
# This reduces connection usage from 2 to 1 for these operations
|
|
semantic_bm25_start = time.time()
|
|
temporal_results_by_ft: dict[str, list[RetrievalResult]] = {}
|
|
temporal_time = 0.0
|
|
|
|
async with acquire_with_retry(pool) as conn:
|
|
conn_wait = time.time() - semantic_bm25_start
|
|
|
|
# Semantic + BM25 combined
|
|
semantic_bm25_results = await retrieve_semantic_bm25_combined(
|
|
conn, query_embedding_str, query_text, bank_id, fact_types, thinking_budget
|
|
)
|
|
semantic_bm25_time = time.time() - semantic_bm25_start
|
|
|
|
# Temporal combined (if constraint detected) - same connection!
|
|
if temporal_constraint:
|
|
tc_start, tc_end = temporal_constraint
|
|
temporal_start = time.time()
|
|
temporal_results_by_ft = await retrieve_temporal_combined(
|
|
conn,
|
|
query_embedding_str,
|
|
bank_id,
|
|
fact_types,
|
|
tc_start,
|
|
tc_end,
|
|
budget=thinking_budget,
|
|
semantic_threshold=0.1,
|
|
)
|
|
temporal_time = time.time() - temporal_start
|
|
|
|
timings["semantic_bm25_combined"] = semantic_bm25_time
|
|
timings["temporal_combined"] = temporal_time
|
|
|
|
# Step 3: Run graph retrieval for each fact type in parallel
|
|
async def run_graph_for_fact_type(ft: str) -> tuple[str, list[RetrievalResult], float, MPFPTimings | None]:
|
|
graph_start = time.time()
|
|
results, mpfp_timing = await retriever.retrieve(
|
|
pool=pool,
|
|
query_embedding_str=query_embedding_str,
|
|
bank_id=bank_id,
|
|
fact_type=ft,
|
|
budget=thinking_budget,
|
|
query_text=query_text,
|
|
semantic_seeds=None,
|
|
temporal_seeds=None,
|
|
)
|
|
return ft, results, time.time() - graph_start, mpfp_timing
|
|
|
|
# Run graph for all fact types in parallel
|
|
graph_tasks = [run_graph_for_fact_type(ft) for ft in fact_types]
|
|
graph_results_list = await asyncio.gather(*graph_tasks)
|
|
|
|
# Organize results by fact type
|
|
results_by_fact_type: dict[str, ParallelRetrievalResult] = {}
|
|
max_conn_wait = conn_wait # Single connection for semantic+bm25+temporal
|
|
all_mpfp_timings: list[MPFPTimings] = []
|
|
|
|
for ft in fact_types:
|
|
# Get semantic + bm25 results for this fact type
|
|
semantic_results, bm25_results = semantic_bm25_results.get(ft, ([], []))
|
|
|
|
# Find graph results for this fact type
|
|
graph_results = []
|
|
graph_time = 0.0
|
|
mpfp_timing = None
|
|
for gr in graph_results_list:
|
|
if gr[0] == ft:
|
|
graph_results = gr[1]
|
|
graph_time = gr[2]
|
|
mpfp_timing = gr[3]
|
|
if mpfp_timing:
|
|
all_mpfp_timings.append(mpfp_timing)
|
|
break
|
|
|
|
# Get temporal results for this fact type from combined result
|
|
temporal_results = temporal_results_by_ft.get(ft) if temporal_constraint else None
|
|
if temporal_results is not None and len(temporal_results) == 0:
|
|
temporal_results = None
|
|
|
|
results_by_fact_type[ft] = ParallelRetrievalResult(
|
|
semantic=semantic_results,
|
|
bm25=bm25_results,
|
|
graph=graph_results,
|
|
temporal=temporal_results,
|
|
timings={
|
|
"semantic": semantic_bm25_time / 2, # Approximate split
|
|
"bm25": semantic_bm25_time / 2,
|
|
"graph": graph_time,
|
|
"temporal": temporal_time, # Same for all fact types (single query)
|
|
"temporal_extraction": temporal_extraction_time,
|
|
},
|
|
temporal_constraint=temporal_constraint,
|
|
mpfp_timings=[mpfp_timing] if mpfp_timing else [],
|
|
max_conn_wait=max_conn_wait,
|
|
)
|
|
|
|
total_time = time.time() - start_time
|
|
timings["total"] = total_time
|
|
|
|
return MultiFactTypeRetrievalResult(
|
|
results_by_fact_type=results_by_fact_type,
|
|
timings=timings,
|
|
max_conn_wait=max_conn_wait,
|
|
)
|