fix(retrieval): Apply proof_count boost to link_expansion retrieval

Author: Abdulkadirklc

hindsight-api-slim/hindsight_api/engine/search/link_expansion_retrieval.py

@@ -59,7 +59,7 @@ async def _find_semantic_seeds(
     rows = await conn.fetch(
         f"""
         SELECT id, text, context, event_date, occurred_start, occurred_end,
-               mentioned_at, fact_type, document_id, chunk_id, tags,
+               mentioned_at, fact_type, document_id, chunk_id, tags, proof_count,
                1 - (embedding <=> $1::vector) AS similarity
         FROM {fq_table("memory_units")}
         WHERE bank_id = $2
@@ -274,7 +274,7 @@ async def _expand_combined(
                 -- Score = COUNT(DISTINCT shared entities), mapped to [0,1] via tanh.
                 SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
                        mu.occurred_end, mu.mentioned_at,
-                       mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
+                       mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
                        COUNT(DISTINCT ue_seed.entity_id)::float AS score,
                        'entity'::text AS source
                 FROM {ue} ue_seed
@@ -298,14 +298,14 @@ async def _expand_combined(
                 SELECT
                     id, text, context, event_date, occurred_start,
                     occurred_end, mentioned_at,
-                    fact_type, document_id, chunk_id, tags,
+                    fact_type, document_id, chunk_id, tags, proof_count,
                     MAX(weight) AS score,
                     'semantic'::text AS source
                 FROM (
                     SELECT
                         mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
                         mu.occurred_end, mu.mentioned_at,
-                        mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
+                        mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
                         ml.weight
                     FROM {ml} ml
                     JOIN {mu} mu ON mu.id = ml.to_unit_id
@@ -317,7 +317,7 @@ async def _expand_combined(
                     SELECT
                         mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
                         mu.occurred_end, mu.mentioned_at,
-                        mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
+                        mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
                         ml.weight
                     FROM {ml} ml
                     JOIN {mu} mu ON mu.id = ml.from_unit_id
@@ -328,7 +328,7 @@ async def _expand_combined(
                 ) sem_raw
                 GROUP BY id, text, context, event_date, occurred_start,
                          occurred_end, mentioned_at,
-                         fact_type, document_id, chunk_id, tags
+                         fact_type, document_id, chunk_id, tags, proof_count
                 ORDER BY score DESC
                 LIMIT $3
             ),
@@ -339,7 +339,7 @@ async def _expand_combined(
                 SELECT DISTINCT ON (mu.id)
                     mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
                     mu.occurred_end, mu.mentioned_at,
-                    mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
+                    mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
                     ml.weight AS score,
                     'causal'::text AS source
                 FROM {ml} ml
@@ -429,7 +429,7 @@ async def _expand_observations(
             SELECT
                 mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
                 mu.occurred_end, mu.mentioned_at,
-                mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
+                mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
                 (SELECT COUNT(DISTINCT s) FROM unnest(mu.source_memory_ids) s WHERE s = ANY(ca.source_ids))::float AS score
             FROM {fq_table("memory_units")} mu, connected_array ca
             WHERE mu.fact_type = 'observation'
@@ -453,35 +453,35 @@ async def _expand_observations(
                 SELECT
                     id, text, context, event_date, occurred_start,
                     occurred_end, mentioned_at,
-                    fact_type, document_id, chunk_id, tags,
+                    fact_type, document_id, chunk_id, tags, proof_count,
                     MAX(weight) AS score,
                     'semantic'::text AS source
                 FROM (
                     SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
                            mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
-                           mu.chunk_id, mu.tags, ml.weight
+                           mu.chunk_id, mu.tags, mu.proof_count, ml.weight
                     FROM {ml} ml JOIN {mu} mu ON mu.id = ml.to_unit_id
                     WHERE ml.from_unit_id = ANY($1::uuid[])
                       AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
                       AND mu.id != ALL($1::uuid[])
                     UNION ALL
                     SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
                            mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
-                           mu.chunk_id, mu.tags, ml.weight
+                           mu.chunk_id, mu.tags, mu.proof_count, ml.weight
                     FROM {ml} ml JOIN {mu} mu ON mu.id = ml.from_unit_id
                     WHERE ml.to_unit_id = ANY($1::uuid[])
                       AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
                       AND mu.id != ALL($1::uuid[])
                 ) sem_raw
                 GROUP BY id, text, context, event_date, occurred_start, occurred_end,
-                         mentioned_at, fact_type, document_id, chunk_id, tags
+                         mentioned_at, fact_type, document_id, chunk_id, tags, proof_count
                 ORDER BY score DESC LIMIT $2
             ),
             causal_expanded AS (
                 SELECT DISTINCT ON (mu.id)
                     mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
                     mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
-                    mu.chunk_id, mu.tags, ml.weight AS score, 'causal'::text AS source
+                    mu.chunk_id, mu.tags, mu.proof_count, ml.weight AS score, 'causal'::text AS source
                 FROM {ml} ml JOIN {mu} mu ON ml.to_unit_id = mu.id
                 WHERE ml.from_unit_id = ANY($1::uuid[])
                   AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')