[ADD] RRF-like clustering

edouardsilvestre · edouardsilvestre · commit bd91a337202d · 2026-01-12T13:45:19.000+01:00
diff --git a/server/clustering/cluster_scoring.py b/server/clustering/cluster_scoring.py
@@ -0,0 +1,26 @@
+from typing import Dict, List
+
+def compute_entity_similarity(a: Dict, b: Dict) -> float:
+    score = 0.0
+
+    if a["subject"] and a["subject"] == b["subject"]:
+        score += 1.0
+
+    if a["event"] and a["event"] == b["event"]:
+        score += 0.5
+
+    orgs_a = set(a.get("orgs", []))
+    orgs_b = set(b.get("orgs", []))
+
+    if orgs_a and orgs_b:
+        score += 0.2 * len(orgs_a & orgs_b)
+
+    return score
+
+def compute_final_score(
+    semantic_score: float,
+    entity_score: float,
+    w_sem: float = 0.6,
+    w_ent: float = 0.4,
+) -> float:
+    return w_sem * semantic_score + w_ent * entity_score
diff --git a/server/clustering/cluster_thresholds.py b/server/clustering/cluster_thresholds.py
@@ -0,0 +1,2 @@
+SEMANTIC_SIMILARITY_THRESHOLD = 0.70
+FINAL_SCORE_THRESHOLD = 0.75
diff --git a/server/database.py b/server/database.py
@@ -2,26 +2,32 @@
 
 from contextlib import contextmanager
 from datetime import UTC, datetime
-from typing import Dict, List, Optional, Tuple, Any
+from typing import Dict, Optional, Any, List
+import json
 
 import numpy as np
 import psycopg
 from pgvector.psycopg import register_vector
 from psycopg.rows import dict_row
 
+from clustering.cluster_scoring import (
+    compute_entity_similarity,
+    compute_final_score,
+)
+from clustering.cluster_thresholds import (
+    SEMANTIC_SIMILARITY_THRESHOLD,
+    FINAL_SCORE_THRESHOLD,
+)
 
-class DatabaseManager:
-    """Manages unified database operations using PostgreSQL."""
 
+class DatabaseManager:
     def __init__(self, db_url: str, embedding_dimension: int = 1536):
-        """Initialize the database manager and ensure schema is ready."""
         self.db_url = db_url
         self.embedding_dimension = embedding_dimension
         self.setup_database()
 
     @contextmanager
     def get_connection(self):
-        """Context manager for PostgreSQL connections with pgvector registered."""
         conn = psycopg.connect(self.db_url, row_factory=dict_row, autocommit=False)
         register_vector(conn)
         try:
@@ -34,10 +40,8 @@ def get_connection(self):
             conn.close()
 
     def setup_database(self):
-        """Initialize database and create tables (idempotent)."""
         with self.get_connection() as conn:
             cur = conn.cursor()
-
             cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
 
             cur.execute(
@@ -55,20 +59,14 @@ def setup_database(self):
                     published_date TIMESTAMPTZ,
                     item_type TEXT,
                     created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP,
-                    updated_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
+                    updated_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP,
+                    subject TEXT,
+                    organization_list JSONB,
+                    event_type TEXT,
+                    cluster_id INTEGER
                 )
                 """
             )
-            
-            cur.execute(
-                """
-                ALTER TABLE articles 
-                ADD COLUMN IF NOT EXISTS subject TEXT,
-                ADD COLUMN IF NOT EXISTS organization_list JSONB, 
-                ADD COLUMN IF NOT EXISTS event_type TEXT,
-                ADD COLUMN IF NOT EXISTS cluster_id INTEGER 
-                """
-            )
 
             cur.execute(
                 f"""
@@ -82,136 +80,89 @@ def setup_database(self):
                 """
             )
 
-            cur.execute(
-                """
-                CREATE TABLE IF NOT EXISTS sync_history (
-                    id SERIAL PRIMARY KEY,
-                    source_site TEXT NOT NULL,
-                    sync_mode TEXT NOT NULL,
-                    last_sync_time TIMESTAMPTZ,
-                    items_processed INTEGER DEFAULT 0,
-                    created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
-                )
-                """
-            )
-
-            cur.execute("CREATE INDEX IF NOT EXISTS idx_content_hash ON articles(content_hash)")
-            cur.execute("CREATE INDEX IF NOT EXISTS idx_source_site ON articles(source_site)")
             cur.execute("CREATE INDEX IF NOT EXISTS idx_cluster_id ON articles(cluster_id)")
 
-    def _compute_content_hash(self, item: Dict) -> str:
-        import hashlib
-        full_content = item.get("full_content", item.get("description", ""))
-        return hashlib.sha256(full_content.encode()).hexdigest()
-
-    def save_article(self, item: Dict, conn: Optional[psycopg.Connection] = None) -> bool:
-        content_hash = self._compute_content_hash(item)
-        if conn is None:
-            with self.get_connection() as owned_conn:
-                return self.save_article(item, owned_conn)
-        cur = conn.cursor()
-        cur.execute(
-            """
-            INSERT INTO articles
-            (id, source_site, title, description, full_content, content_hash, author_info, keywords, content_url, published_date, item_type, created_at)
-            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
-            ON CONFLICT DO NOTHING
-            """,
-            (
-                item["id"],
-                item["source_site"],
-                item["title"],
-                item.get("description", ""),
-                item.get("full_content", item.get("description", "")),
-                content_hash,
-                item.get("author_info", ""),
-                item.get("keywords", ""),
-                item["content_url"],
-                item.get("published_date", datetime.now(UTC)),
-                item.get("item_type", "article"),
-                datetime.now(UTC),
-            ),
-        )
-        return cur.rowcount > 0
-
-    def article_exists(self, article_id: str) -> bool:
-        with self.get_connection() as conn:
-            cur = conn.cursor()
-            cur.execute("SELECT 1 FROM articles WHERE id = %s", (article_id,))
-            return cur.fetchone() is not None
-
-    def article_exists_by_hash(self, content_hash: str) -> bool:
+    def assign_cluster_with_similarity(
+        self,
+        article_id: str,
+        subject: str,
+        orgs: List[str],
+        event: str,
+    ):
         with self.get_connection() as conn:
             cur = conn.cursor()
-            cur.execute("SELECT 1 FROM articles WHERE content_hash = %s", (content_hash,))
-            return cur.fetchone() is not None
 
-    def save_embedding(self, article_id: str, embedding: np.ndarray, model: str = "default") -> bool:
-        with self.get_connection() as conn:
-            cur = conn.cursor()
             cur.execute(
                 """
-                INSERT INTO embeddings (article_id, embedding, embedding_model)
-                VALUES (%s, %s, %s)
-                ON CONFLICT (article_id) DO UPDATE
-                SET embedding = EXCLUDED.embedding,
-                    embedding_model = EXCLUDED.embedding_model,
-                    created_at = CURRENT_TIMESTAMP
+                SELECT embedding FROM embeddings WHERE article_id = %s
                 """,
-                (article_id, embedding.tolist(), model),
+                (article_id,)
             )
-            return cur.rowcount > 0
+            row = cur.fetchone()
+            if not row:
+                return
+
+            article_embedding = row["embedding"]
 
-    def assign_cluster_by_entities(self, article_id: str, subject: str, orgs: str, event: str):
-        with self.get_connection() as conn:
-            cur = conn.cursor()
             cur.execute(
                 """
-                SELECT cluster_id FROM articles 
-                WHERE subject = %s AND organization_list = %s::jsonb AND event_type = %s 
-                AND cluster_id IS NOT NULL LIMIT 1
+                SELECT a.cluster_id, a.subject, a.organization_list, a.event_type,
+                       1 - (e.embedding <=> %s) AS semantic_similarity
+                FROM articles a
+                JOIN embeddings e ON a.id = e.article_id
+                WHERE a.cluster_id IS NOT NULL
+                  AND (a.subject = %s OR a.event_type = %s)
                 """,
-                (subject, orgs, event)
+                (article_embedding, subject, event),
             )
-            row = cur.fetchone()
-            if row:
-                cid = row['cluster_id']
-            else:
-                cur.execute("SELECT COALESCE(MAX(cluster_id), 0) + 1 as next_id FROM articles")
-                cid = cur.fetchone()['next_id']
-            
+
+            best_cluster_id = None
+            best_score = 0.0
+
+            for row in cur.fetchall():
+                semantic_score = row["semantic_similarity"]
+
+                if semantic_score < SEMANTIC_SIMILARITY_THRESHOLD:
+                    continue
+
+                entity_score = compute_entity_similarity(
+                    {
+                        "subject": subject,
+                        "orgs": orgs,
+                        "event": event,
+                    },
+                    {
+                        "subject": row["subject"],
+                        "orgs": row["organization_list"] or [],
+                        "event": row["event_type"],
+                    },
+                )
+
+                final_score = compute_final_score(semantic_score, entity_score)
+
+                if final_score > best_score:
+                    best_score = final_score
+                    best_cluster_id = row["cluster_id"]
+
+            if best_cluster_id is None or best_score < FINAL_SCORE_THRESHOLD:
+                cur.execute("SELECT COALESCE(MAX(cluster_id), 0) + 1 AS next_id FROM articles")
+                best_cluster_id = cur.fetchone()["next_id"]
+
             cur.execute(
                 """
-                UPDATE articles SET subject=%s, organization_list=%s::jsonb, 
-                event_type=%s, cluster_id=%s, updated_at=CURRENT_TIMESTAMP 
-                WHERE id=%s
+                UPDATE articles
+                SET subject = %s,
+                    organization_list = %s::jsonb,
+                    event_type = %s,
+                    cluster_id = %s,
+                    updated_at = CURRENT_TIMESTAMP
+                WHERE id = %s
                 """,
-                (subject, orgs, event, cid, article_id)
+                (
+                    subject,
+                    json.dumps(orgs),
+                    event,
+                    best_cluster_id,
+                    article_id,
+                ),
             )
-
-    def get_stats(self) -> Dict:
-        with self.get_connection() as conn:
-            cur = conn.cursor()
-            cur.execute("SELECT COUNT(*) AS count FROM articles")
-            total_articles = cur.fetchone()["count"]
-            cur.execute("SELECT COUNT(*) AS count FROM embeddings")
-            total_embeddings = cur.fetchone()["count"]
-            cur.execute("SELECT COUNT(DISTINCT cluster_id) AS count FROM articles WHERE cluster_id IS NOT NULL")
-            total_clusters = cur.fetchone()["count"]
-            cur.execute("SELECT source_site, COUNT(*) as count FROM articles GROUP BY source_site")
-            articles_by_source = {row["source_site"]: row["count"] for row in cur.fetchall()}
-            return {
-                "total_articles": total_articles,
-                "total_embeddings": total_embeddings,
-                "articles_by_source": articles_by_source,
-                "articles_without_embeddings": total_articles - total_embeddings,
-                "total_clusters": total_clusters,
-            }
-
-    def record_sync(self, source: str, mode: str, items_processed: int = 0):
-        with self.get_connection() as conn:
-            cur = conn.cursor()
-            cur.execute(
-                "INSERT INTO sync_history (source_site, sync_mode, last_sync_time, items_processed) VALUES (%s, %s, %s, %s)",
-                (source, mode, datetime.now(UTC), items_processed),
-            )
diff --git a/server/entity_llm_processor.py b/server/entity_llm_processor.py
@@ -1,7 +1,7 @@
 import logging
 import json
 from typing import Dict, Any
-from openai import OpenAI, APIError
+from openai import OpenAI
 
 logger = logging.getLogger(__name__)
 
@@ -15,19 +15,15 @@ def __init__(self, model: str = "gpt-4o-mini"):
             "JSON format. Focus on high-level subjects, involved organizations, "
             "and the type of event the article describes."
         )
-        self.entity_schema = {
-            "type": "object",
-            "properties": {
-                "subject": {"type": "string"},
-                "organization_list": {"type": "array", "items": {"type": "string"}},
-                "event_type": {"type": "string"},
-            },
-            "required": ["subject", "organization_list", "event_type"]
-        }
 
     def process(self, article: Dict, db_manager: Any) -> bool:
         article_id = article["id"]
-        content = f"Title: {article.get('title', '')}\nDescription: {article.get('description', '')}\nContent: {article.get('full_content', '')[:500]}..."
+        content = (
+            f"Title: {article.get('title', '')}\n"
+            f"Description: {article.get('description', '')}\n"
+            f"Content: {article.get('full_content', '')[:500]}..."
+        )
+
         try:
             response = self.client.chat.completions.create(
                 model=self.model,
@@ -37,14 +33,17 @@ def process(self, article: Dict, db_manager: Any) -> bool:
                 ],
                 response_format={"type": "json_object"},
             )
+
             data = json.loads(response.choices[0].message.content)
-            db_manager.assign_cluster_by_entities(
-                article_id,
-                subject=data.get("subject"),
-                orgs=json.dumps(data.get("organization_list", [])),
-                event=data.get("event_type")
+
+            db_manager.assign_cluster_with_similarity(
+                article_id=article_id,
+                subject=data.get("subject", ""),
+                orgs=data.get("organization_list", []),
+                event=data.get("event_type", "")
             )
             return True
+
         except Exception as e:
             logger.error(f"Error for {article_id}: {e}")
-            return False
+            return False

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+SEMANTIC_SIMILARITY_THRESHOLD = 0.70`
	`2`	`+FINAL_SCORE_THRESHOLD = 0.75`