diff --git a/.env.example b/.env.example index 614d7a5..ac149f1 100644 --- a/.env.example +++ b/.env.example @@ -19,6 +19,11 @@ QDRANT_HOST=localhost QDRANT_PORT=6333 QDRANT_API_KEY= # Optional: set for Qdrant Cloud or when API key auth is enabled +FALKORDB_HOST=localhost +FALKORDB_PORT=6379 +FALKORDB_USERNAME= # Optional: leave empty if authentication is disabled +FALKORDB_PASSWORD= # Optional: leave empty if authentication is disabled + AIRFLOW_UID= echo $(id -u) _AIRFLOW_WWW_USER_USERNAME=airflow _AIRFLOW_WWW_USER_PASSWORD=airflow diff --git a/airflow_config/dags/knowledge_graph.py b/airflow_config/dags/knowledge_graph.py new file mode 100644 index 0000000..2122530 --- /dev/null +++ b/airflow_config/dags/knowledge_graph.py @@ -0,0 +1,49 @@ +from datetime import datetime, timedelta + +from airflow import DAG +from airflow.operators.python import PythonOperator +from notifier.notifications_template import ( + get_failure_notifier, + get_start_notifier, + get_success_notifier, +) + +from database import init_graph_schema, populate_graph_from_postgres + +default_args = { + "owner": "airflow", + "start_date": datetime(2025, 8, 1), + "retries": 1, + "retry_delay": timedelta(minutes=5), +} + +with DAG( + "KNOWLEDGE_GRAPH", + default_args=default_args, + schedule=None, + catchup=False, + max_active_runs=1, + description=( + "FalkorDB knowledge graph initialisation and back-fill DAG. " + "Initialises the graph schema and populates it from data already " + "stored in PostgreSQL (LEGI, JADE, BOFIP)." + ), + tags=["mediatech", "knowledge_graph", "falkordb", "graphrag"], +) as dag: + init_schema = PythonOperator( + task_id="init_graph_schema", + python_callable=init_graph_schema, + on_execute_callback=get_start_notifier(), + on_success_callback=get_success_notifier(), + on_failure_callback=get_failure_notifier(), + ) + + backfill_graph = PythonOperator( + task_id="populate_graph_from_postgres", + python_callable=populate_graph_from_postgres, + on_execute_callback=get_start_notifier(), + on_success_callback=get_success_notifier(), + on_failure_callback=get_failure_notifier(), + ) + + init_schema >> backfill_graph diff --git a/config/__init__.py b/config/__init__.py index 57327cf..211e05d 100644 --- a/config/__init__.py +++ b/config/__init__.py @@ -7,6 +7,11 @@ DATA_GOUV_DATASETS_CATALOG_DATA_FOLDER, DOLE_DATA_FOLDER, EMBEDDING_MODEL, + FALKORDB_GRAPH_NAME, + FALKORDB_HOST, + FALKORDB_PASSWORD, + FALKORDB_PORT, + FALKORDB_USERNAME, HF_TOKEN, LEGI_DATA_FOLDER, LLM_MODEL, diff --git a/config/config.py b/config/config.py index e386b44..a2a342a 100644 --- a/config/config.py +++ b/config/config.py @@ -13,6 +13,8 @@ postgres_port = "5432" qdrant_host = "qdrant" qdrant_port = "6333" + falkordb_host = "falkordb" + falkordb_port = "6379" else: # Locally, using relative paths base_path = "." @@ -20,6 +22,8 @@ postgres_port = os.getenv("POSTGRES_PORT", "5433") qdrant_host = os.getenv("QDRANT_HOST", "localhost") qdrant_port = os.getenv("QDRANT_PORT", "6333") + falkordb_host = os.getenv("FALKORDB_HOST", "localhost") + falkordb_port = os.getenv("FALKORDB_PORT", "6379") # PostgreSQL configuration @@ -35,6 +39,13 @@ QDRANT_URL = f"http://{qdrant_host}:{qdrant_port}" QDRANT_API_KEY = os.getenv("QDRANT_API_KEY", None) +# FalkorDB configuration +FALKORDB_HOST = falkordb_host +FALKORDB_PORT = int(falkordb_port) +FALKORDB_USERNAME = os.getenv("FALKORDB_USERNAME", None) or None +FALKORDB_PASSWORD = os.getenv("FALKORDB_PASSWORD", None) or None +FALKORDB_GRAPH_NAME = "frenchadmin" + BASE_PATH = base_path # Paths for configurations and data history diff --git a/database/__init__.py b/database/__init__.py index 12985b9..069bea3 100644 --- a/database/__init__.py +++ b/database/__init__.py @@ -11,3 +11,15 @@ split_legi_table, sync_obsolete_doc_ids, ) +from .graph_manage import ( + build_graphrag_knowledge_graph, + close_graph_connection, + init_graph_schema, + populate_graph_from_postgres, + upsert_bofip_chunk, + upsert_bofip_node, + upsert_jade_chunk, + upsert_jade_node, + upsert_legi_chunk, + upsert_legi_node, +) diff --git a/database/database_manage.py b/database/database_manage.py index b245ccf..01bec9c 100644 --- a/database/database_manage.py +++ b/database/database_manage.py @@ -445,6 +445,24 @@ def create_all_tables(model=EMBEDDING_MODEL, delete_existing: bool = False): ) """) + elif table_name.lower() == "bofip": + cursor.execute(f""" + CREATE TABLE BOFIP ( + chunk_id TEXT PRIMARY KEY, + doc_id TEXT NOT NULL, + chunk_index INTEGER NOT NULL, + chunk_xxh64 TEXT NOT NULL, + nature TEXT, + category TEXT, + title TEXT, + date TEXT, + text TEXT, + chunk_text TEXT, + "embeddings_{model_name}" vector({embedding_size}), + UNIQUE(chunk_id) + ) + """) + # Create HNSW index for vector similarity search try: if table_name.lower() not in CONFIG_TABLES: diff --git a/database/graph_manage.py b/database/graph_manage.py new file mode 100644 index 0000000..03d7357 --- /dev/null +++ b/database/graph_manage.py @@ -0,0 +1,838 @@ +""" +Knowledge Graph management module for FalkorDB. + +This module implements the FrenchAdmin knowledge graph on top of FalkorDB, +a graph database built on Redis. It provides the ontology, population +functions, and GraphRAG integration for LEGI, JADE, and BOFIP data. + +Graph Ontology +-------------- +Node labels: + - LegalText : LEGI articles / legislative texts + - JudicialDecision: JADE administrative court decisions + - TaxGuidance : BOFIP tax administration guidance + - LegalCode : French legal codes (Code civil, Code pénal, …) + - Ministry : French ministries + - Jurisdiction : Administrative courts / jurisdictions + - Chunk : Text chunks linked to a source document (GraphRAG) + +Relationship types: + - BELONGS_TO_CODE : (LegalText | TaxGuidance) → LegalCode + - ISSUED_BY : LegalText → Ministry + - REFERENCES : LegalText → LegalText (from LIENS metadata) + - DECIDED_BY : JudicialDecision → Jurisdiction + - PART_OF : Chunk → (LegalText | JudicialDecision | TaxGuidance) + +Design decision +--------------- +Nodes and relationships are populated **in parallel** with the PostgreSQL +inserts (Option B), i.e. during raw XML/CSV processing rather than from the +already-stored relational data. This avoids a second full-data pass, keeps +both stores consistent, and exposes richer structural information (e.g. the +LIENS cross-references available only in the raw XML). + +All graph operations are wrapped in try/except so that a FalkorDB outage or +misconfiguration never interrupts the main PostgreSQL pipeline. +""" + +import json + +from config import ( + FALKORDB_GRAPH_NAME, + FALKORDB_HOST, + FALKORDB_PASSWORD, + FALKORDB_PORT, + FALKORDB_USERNAME, + get_logger, +) + +logger = get_logger(__name__) + +# --------------------------------------------------------------------------- +# Lazy FalkorDB connection +# --------------------------------------------------------------------------- + +_graph = None + + +def _get_graph(): + """Return the shared FalkorDB graph handle, creating it on first call. + + Returns ``None`` if FalkorDB is unavailable or not configured, which + allows callers to skip graph operations gracefully. + """ + global _graph + if _graph is not None: + return _graph + try: + import falkordb # noqa: PLC0415 (imported lazily to keep startup fast) + + kwargs = {"host": FALKORDB_HOST, "port": FALKORDB_PORT} + if FALKORDB_USERNAME: + kwargs["username"] = FALKORDB_USERNAME + if FALKORDB_PASSWORD: + kwargs["password"] = FALKORDB_PASSWORD + + client = falkordb.FalkorDB(**kwargs) + _graph = client.select_graph(FALKORDB_GRAPH_NAME) + logger.info( + "Connected to FalkorDB at %s:%s, graph '%s'", + FALKORDB_HOST, + FALKORDB_PORT, + FALKORDB_GRAPH_NAME, + ) + return _graph + except Exception as exc: + logger.warning( + "FalkorDB is unavailable (%s). Knowledge-graph operations will be skipped.", + exc, + ) + return None + + +def close_graph_connection(): + """Reset the cached graph handle (useful for testing).""" + global _graph + _graph = None + + +# --------------------------------------------------------------------------- +# Schema initialisation +# --------------------------------------------------------------------------- + +_INDEXES = [ + ("LegalText", "doc_id"), + ("JudicialDecision", "doc_id"), + ("TaxGuidance", "doc_id"), + ("LegalCode", "name"), + ("Ministry", "name"), + ("Jurisdiction", "name"), + ("Chunk", "chunk_id"), +] + + +def init_graph_schema(): + """Create indexes required for efficient node lookups. + + Safe to call multiple times – existing indexes are silently ignored. + """ + graph = _get_graph() + if graph is None: + return + for label, prop in _INDEXES: + try: + graph.query(f"CREATE INDEX FOR (n:{label}) ON (n.{prop})") + logger.debug("Index created: (%s).%s", label, prop) + except Exception as exc: + if "already indexed" in str(exc).lower() or "equivalent index" in str(exc).lower(): + logger.debug("Index already exists: (%s).%s", label, prop) + else: + logger.warning("Could not create index (%s).%s: %s", label, prop, exc) + logger.info("FalkorDB graph schema initialised for graph '%s'", FALKORDB_GRAPH_NAME) + + +# --------------------------------------------------------------------------- +# Helper: fire-and-forget Cypher +# --------------------------------------------------------------------------- + +def _run(query: str, params: dict | None = None): + """Execute a Cypher query, logging any errors without raising.""" + graph = _get_graph() + if graph is None: + return + try: + graph.query(query, params or {}) + except Exception as exc: + logger.warning("Graph query failed: %s | params=%s | error=%s", query, params, exc) + + +# --------------------------------------------------------------------------- +# Node upsert helpers +# --------------------------------------------------------------------------- + +def _upsert_legal_code(name: str): + if not name: + return + _run( + "MERGE (n:LegalCode {name: $name})", + {"name": name}, + ) + + +def _upsert_ministry(name: str): + if not name: + return + _run( + "MERGE (n:Ministry {name: $name})", + {"name": name}, + ) + + +def _upsert_jurisdiction(name: str): + if not name: + return + _run( + "MERGE (n:Jurisdiction {name: $name})", + {"name": name}, + ) + + +# --------------------------------------------------------------------------- +# LEGI — LegalText nodes + relationships +# --------------------------------------------------------------------------- + +def upsert_legi_node( + doc_id: str, + nature: str | None, + category: str | None, + ministry: str | None, + status: str | None, + title: str | None, + full_title: str | None, + number: str | None, + start_date: str | None, + end_date: str | None, + links: list | None = None, +): + """Insert or update a :LegalText node for a LEGI article. + + Also creates :LegalCode and :Ministry nodes and the corresponding + BELONGS_TO_CODE, ISSUED_BY, and REFERENCES relationships. + + Args: + doc_id: Unique document identifier (CID). + nature: Legal nature of the text (e.g. "LOI", "DECRET"). + category: Legal code the article belongs to. + ministry: Ministry that issued the text. + status: Current status (e.g. "VIGUEUR"). + title: Short title. + full_title: Full title. + number: Article number. + start_date: Validity start date (YYYY-MM-DD). + end_date: Validity end date (YYYY-MM-DD). + links: List of cross-reference dicts from the LIENS XML element. + """ + graph = _get_graph() + if graph is None: + return + + # Upsert the LegalText node + _run( + """ + MERGE (n:LegalText {doc_id: $doc_id}) + SET n.nature = $nature, + n.category = $category, + n.ministry = $ministry, + n.status = $status, + n.title = $title, + n.full_title = $full_title, + n.number = $number, + n.start_date = $start_date, + n.end_date = $end_date + """, + { + "doc_id": doc_id, + "nature": nature, + "category": category, + "ministry": ministry, + "status": status, + "title": title, + "full_title": full_title, + "number": number, + "start_date": start_date, + "end_date": end_date, + }, + ) + + # LegalCode relationship + if category: + _upsert_legal_code(category) + _run( + """ + MATCH (lt:LegalText {doc_id: $doc_id}), (lc:LegalCode {name: $name}) + MERGE (lt)-[:BELONGS_TO_CODE]->(lc) + """, + {"doc_id": doc_id, "name": category}, + ) + + # Ministry relationship + if ministry: + _upsert_ministry(ministry) + _run( + """ + MATCH (lt:LegalText {doc_id: $doc_id}), (m:Ministry {name: $name}) + MERGE (lt)-[:ISSUED_BY]->(m) + """, + {"doc_id": doc_id, "name": ministry}, + ) + + # Cross-references (LIENS) + if links: + for link in links: + ref_doc_id = link.get("doc_id") + if not ref_doc_id: + continue + # Ensure the referenced node exists (minimal stub) + _run( + "MERGE (n:LegalText {doc_id: $ref_id})", + {"ref_id": ref_doc_id}, + ) + _run( + """ + MATCH (src:LegalText {doc_id: $src_id}), + (dst:LegalText {doc_id: $dst_id}) + MERGE (src)-[:REFERENCES {type: $link_type}]->(dst) + """, + { + "src_id": doc_id, + "dst_id": ref_doc_id, + "link_type": link.get("link_type", ""), + }, + ) + + +def upsert_legi_chunk(chunk_id: str, doc_id: str): + """Link a text chunk to its parent :LegalText node. + + Args: + chunk_id: Unique chunk identifier. + doc_id: Parent document identifier. + """ + graph = _get_graph() + if graph is None: + return + _run( + """ + MERGE (c:Chunk {chunk_id: $chunk_id}) + SET c.source_type = 'legi' + WITH c + MATCH (lt:LegalText {doc_id: $doc_id}) + MERGE (c)-[:PART_OF]->(lt) + """, + {"chunk_id": chunk_id, "doc_id": doc_id}, + ) + + +# --------------------------------------------------------------------------- +# JADE — JudicialDecision nodes + relationships +# --------------------------------------------------------------------------- + +def upsert_jade_node( + doc_id: str, + nature: str | None, + solution: str | None, + title: str | None, + number: str | None, + decision_date: str | None, + jurisdiction: str | None, + formation: str | None, +): + """Insert or update a :JudicialDecision node for a JADE decision. + + Also creates a :Jurisdiction node and the DECIDED_BY relationship. + + Args: + doc_id: Unique document identifier (CID). + nature: Nature of the decision. + solution: Decision solution (e.g. "Rejet", "Annulation"). + title: Title of the decision. + number: Decision number. + decision_date: Date of the decision (YYYY-MM-DD). + jurisdiction: Name of the issuing court. + formation: Court formation that rendered the decision. + """ + graph = _get_graph() + if graph is None: + return + + _run( + """ + MERGE (n:JudicialDecision {doc_id: $doc_id}) + SET n.nature = $nature, + n.solution = $solution, + n.title = $title, + n.number = $number, + n.decision_date = $decision_date, + n.jurisdiction = $jurisdiction, + n.formation = $formation + """, + { + "doc_id": doc_id, + "nature": nature, + "solution": solution, + "title": title, + "number": number, + "decision_date": decision_date, + "jurisdiction": jurisdiction, + "formation": formation, + }, + ) + + # Jurisdiction relationship + if jurisdiction: + _upsert_jurisdiction(jurisdiction) + _run( + """ + MATCH (jd:JudicialDecision {doc_id: $doc_id}), + (j:Jurisdiction {name: $name}) + MERGE (jd)-[:DECIDED_BY]->(j) + """, + {"doc_id": doc_id, "name": jurisdiction}, + ) + + +def upsert_jade_chunk(chunk_id: str, doc_id: str): + """Link a text chunk to its parent :JudicialDecision node. + + Args: + chunk_id: Unique chunk identifier. + doc_id: Parent document identifier. + """ + graph = _get_graph() + if graph is None: + return + _run( + """ + MERGE (c:Chunk {chunk_id: $chunk_id}) + SET c.source_type = 'jade' + WITH c + MATCH (jd:JudicialDecision {doc_id: $doc_id}) + MERGE (c)-[:PART_OF]->(jd) + """, + {"chunk_id": chunk_id, "doc_id": doc_id}, + ) + + +# --------------------------------------------------------------------------- +# BOFIP — TaxGuidance nodes + relationships +# --------------------------------------------------------------------------- + +def upsert_bofip_node( + doc_id: str, + nature: str | None, + category: str | None, + title: str | None, + date: str | None, +): + """Insert or update a :TaxGuidance node for a BOFIP document. + + Also creates a :LegalCode node and the BELONGS_TO_CODE relationship when + a category is provided. + + Args: + doc_id: Unique document identifier (CID). + nature: Nature of the guidance document. + category: Tax category / legal code reference. + title: Title of the guidance document. + date: Publication date (YYYY-MM-DD). + """ + graph = _get_graph() + if graph is None: + return + + _run( + """ + MERGE (n:TaxGuidance {doc_id: $doc_id}) + SET n.nature = $nature, + n.category = $category, + n.title = $title, + n.date = $date + """, + { + "doc_id": doc_id, + "nature": nature, + "category": category, + "title": title, + "date": date, + }, + ) + + if category: + _upsert_legal_code(category) + _run( + """ + MATCH (tg:TaxGuidance {doc_id: $doc_id}), + (lc:LegalCode {name: $name}) + MERGE (tg)-[:BELONGS_TO_CODE]->(lc) + """, + {"doc_id": doc_id, "name": category}, + ) + + +def upsert_bofip_chunk(chunk_id: str, doc_id: str): + """Link a text chunk to its parent :TaxGuidance node. + + Args: + chunk_id: Unique chunk identifier. + doc_id: Parent document identifier. + """ + graph = _get_graph() + if graph is None: + return + _run( + """ + MERGE (c:Chunk {chunk_id: $chunk_id}) + SET c.source_type = 'bofip' + WITH c + MATCH (tg:TaxGuidance {doc_id: $doc_id}) + MERGE (c)-[:PART_OF]->(tg) + """, + {"chunk_id": chunk_id, "doc_id": doc_id}, + ) + + +# --------------------------------------------------------------------------- +# Bulk population from PostgreSQL (alternative / back-fill path) +# --------------------------------------------------------------------------- + +def populate_graph_from_postgres(): + """Populate the knowledge graph from data already stored in PostgreSQL. + + This is the alternative to the parallel-population approach (Option A). + It is useful for a one-off back-fill when FalkorDB is introduced to an + existing deployment where PostgreSQL data has already been ingested. + + The function reads from the LEGI, JADE, and (if present) BOFIP tables and + re-creates the graph nodes and relationships. Large tables are streamed + with server-side cursors to avoid loading everything into memory. + """ + try: + import psycopg2 # noqa: PLC0415 + + from config import ( # noqa: PLC0415 + POSTGRES_DB, + POSTGRES_HOST, + POSTGRES_PASSWORD, + POSTGRES_PORT, + POSTGRES_USER, + ) + + conn = psycopg2.connect( + host=POSTGRES_HOST, + port=POSTGRES_PORT, + dbname=POSTGRES_DB, + user=POSTGRES_USER, + password=POSTGRES_PASSWORD, + ) + except Exception as exc: + logger.error("Cannot connect to PostgreSQL for graph back-fill: %s", exc) + return + + try: + _populate_legi_from_postgres(conn) + _populate_jade_from_postgres(conn) + _populate_bofip_from_postgres(conn) + finally: + conn.close() + logger.info("Graph back-fill from PostgreSQL completed.") + + +def _populate_legi_from_postgres(conn): + """Back-fill :LegalText nodes from the LEGI PostgreSQL table.""" + graph = _get_graph() + if graph is None: + return + try: + with conn.cursor(name="legi_cursor") as cur: + cur.execute( + """ + SELECT DISTINCT ON (doc_id) + doc_id, nature, category, ministry, status, + title, full_title, number, start_date, end_date, links + FROM LEGI + ORDER BY doc_id, chunk_index + """ + ) + count = 0 + for row in cur: + ( + doc_id, nature, category, ministry, status, + title, full_title, number, start_date, end_date, links_json, + ) = row + links = [] + if links_json: + try: + links = json.loads(links_json) if isinstance(links_json, str) else links_json + except (json.JSONDecodeError, TypeError): + links = [] + upsert_legi_node( + doc_id=doc_id, + nature=nature, + category=category, + ministry=ministry, + status=status, + title=title, + full_title=full_title, + number=number, + start_date=start_date, + end_date=end_date, + links=links, + ) + count += 1 + if count % 10000 == 0: + logger.info("LEGI back-fill: %d nodes upserted", count) + logger.info("LEGI back-fill complete: %d nodes total", count) + except Exception as exc: + logger.error("Error during LEGI graph back-fill: %s", exc) + + +def _populate_jade_from_postgres(conn): + """Back-fill :JudicialDecision nodes from the JADE PostgreSQL table.""" + graph = _get_graph() + if graph is None: + return + try: + with conn.cursor(name="jade_cursor") as cur: + cur.execute( + """ + SELECT DISTINCT ON (doc_id) + doc_id, nature, solution, title, number, + decision_date, jurisdiction, formation + FROM JADE + ORDER BY doc_id, chunk_index + """ + ) + count = 0 + for row in cur: + ( + doc_id, nature, solution, title, number, + decision_date, jurisdiction, formation, + ) = row + upsert_jade_node( + doc_id=doc_id, + nature=nature, + solution=solution, + title=title, + number=number, + decision_date=decision_date, + jurisdiction=jurisdiction, + formation=formation, + ) + count += 1 + if count % 10000 == 0: + logger.info("JADE back-fill: %d nodes upserted", count) + logger.info("JADE back-fill complete: %d nodes total", count) + except Exception as exc: + logger.error("Error during JADE graph back-fill: %s", exc) + + +def _populate_bofip_from_postgres(conn): + """Back-fill :TaxGuidance nodes from the BOFIP PostgreSQL table (if it exists).""" + graph = _get_graph() + if graph is None: + return + try: + with conn.cursor() as cur: + cur.execute( + """ + SELECT EXISTS ( + SELECT 1 FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name = 'bofip' + ) + """ + ) + if not cur.fetchone()[0]: + logger.info("BOFIP table not found in PostgreSQL, skipping back-fill.") + return + with conn.cursor(name="bofip_cursor") as cur: + cur.execute( + """ + SELECT DISTINCT ON (doc_id) + doc_id, nature, category, title, date + FROM BOFIP + ORDER BY doc_id, chunk_index + """ + ) + count = 0 + for row in cur: + doc_id, nature, category, title, date = row + upsert_bofip_node( + doc_id=doc_id, + nature=nature, + category=category, + title=title, + date=date, + ) + count += 1 + if count % 10000 == 0: + logger.info("BOFIP back-fill: %d nodes upserted", count) + logger.info("BOFIP back-fill complete: %d nodes total", count) + except Exception as exc: + logger.error("Error during BOFIP graph back-fill: %s", exc) + + +# --------------------------------------------------------------------------- +# GraphRAG — Knowledge Graph builder using graphrag-sdk +# --------------------------------------------------------------------------- + +def build_graphrag_knowledge_graph(llm_model: str | None = None): + """Build a GraphRAG Knowledge Graph on top of the FalkorDB graph. + + This function leverages the ``graphrag-sdk`` library to enrich the graph + with entity extraction and question-answering capabilities. + + The ontology passed to the SDK mirrors the hand-crafted one defined in this + module so that both the raw Cypher graph and the GraphRAG overlay share the + same node/relationship vocabulary. + + Args: + llm_model: LLM model identifier to use for entity extraction. When + ``None`` the value is taken from the ``LLM_MODEL`` configuration. + + Returns: + A ``KnowledgeGraph`` instance from ``graphrag_sdk`` (or ``None`` if the + SDK or FalkorDB is unavailable). + """ + try: + from graphrag_sdk import KnowledgeGraph, Ontology # noqa: PLC0415 + from graphrag_sdk.models.openai import OpenAiGenerativeModel # noqa: PLC0415 + from graphrag_sdk.ontology import Edge, Node, Property # noqa: PLC0415 + except ImportError: + logger.warning( + "graphrag-sdk is not installed. " + "Install it with: pip install graphrag-sdk==0.8.2" + ) + return None + + try: + from config import API_KEY, API_URL, LLM_MODEL # noqa: PLC0415 + except ImportError as exc: + logger.error("Missing dependency for GraphRAG KG builder: %s", exc) + return None + + model_name = llm_model or LLM_MODEL + + # ------------------------------------------------------------------ + # Define the ontology + # ------------------------------------------------------------------ + ontology = Ontology() + + # Nodes + legal_text_node = Node( + label="LegalText", + properties=[ + Property(name="doc_id", type="str", required=True, unique=True), + Property(name="nature", type="str"), + Property(name="category", type="str"), + Property(name="ministry", type="str"), + Property(name="status", type="str"), + Property(name="title", type="str"), + Property(name="full_title", type="str"), + Property(name="number", type="str"), + Property(name="start_date", type="str"), + Property(name="end_date", type="str"), + ], + ) + judicial_decision_node = Node( + label="JudicialDecision", + properties=[ + Property(name="doc_id", type="str", required=True, unique=True), + Property(name="nature", type="str"), + Property(name="solution", type="str"), + Property(name="title", type="str"), + Property(name="number", type="str"), + Property(name="decision_date", type="str"), + Property(name="jurisdiction", type="str"), + Property(name="formation", type="str"), + ], + ) + tax_guidance_node = Node( + label="TaxGuidance", + properties=[ + Property(name="doc_id", type="str", required=True, unique=True), + Property(name="nature", type="str"), + Property(name="category", type="str"), + Property(name="title", type="str"), + Property(name="date", type="str"), + ], + ) + legal_code_node = Node( + label="LegalCode", + properties=[ + Property(name="name", type="str", required=True, unique=True), + ], + ) + ministry_node = Node( + label="Ministry", + properties=[ + Property(name="name", type="str", required=True, unique=True), + ], + ) + jurisdiction_node = Node( + label="Jurisdiction", + properties=[ + Property(name="name", type="str", required=True, unique=True), + ], + ) + chunk_node = Node( + label="Chunk", + properties=[ + Property(name="chunk_id", type="str", required=True, unique=True), + Property(name="source_type", type="str"), + ], + ) + + for node in ( + legal_text_node, + judicial_decision_node, + tax_guidance_node, + legal_code_node, + ministry_node, + jurisdiction_node, + chunk_node, + ): + ontology.add_node(node) + + # Edges + ontology.add_edge( + Edge(relation="BELONGS_TO_CODE", source="LegalText", target="LegalCode") + ) + ontology.add_edge( + Edge(relation="BELONGS_TO_CODE", source="TaxGuidance", target="LegalCode") + ) + ontology.add_edge( + Edge(relation="ISSUED_BY", source="LegalText", target="Ministry") + ) + ontology.add_edge( + Edge(relation="REFERENCES", source="LegalText", target="LegalText", + properties=[Property(name="type", type="str")]) + ) + ontology.add_edge( + Edge(relation="DECIDED_BY", source="JudicialDecision", target="Jurisdiction") + ) + ontology.add_edge( + Edge(relation="PART_OF", source="Chunk", target="LegalText") + ) + ontology.add_edge( + Edge(relation="PART_OF", source="Chunk", target="JudicialDecision") + ) + ontology.add_edge( + Edge(relation="PART_OF", source="Chunk", target="TaxGuidance") + ) + + # ------------------------------------------------------------------ + # Instantiate the KnowledgeGraph + # ------------------------------------------------------------------ + try: + kwargs = {"host": FALKORDB_HOST, "port": FALKORDB_PORT} + if FALKORDB_USERNAME: + kwargs["username"] = FALKORDB_USERNAME + if FALKORDB_PASSWORD: + kwargs["password"] = FALKORDB_PASSWORD + + llm = OpenAiGenerativeModel(model=model_name, api_key=API_KEY, base_url=API_URL) + + kg = KnowledgeGraph( + name=FALKORDB_GRAPH_NAME, + ontology=ontology, + model=llm, + **kwargs, + ) + logger.info( + "GraphRAG KnowledgeGraph instantiated on graph '%s'", FALKORDB_GRAPH_NAME + ) + return kg + except Exception as exc: + logger.error("Failed to build GraphRAG KnowledgeGraph: %s", exc) + return None diff --git a/docker-compose.yml b/docker-compose.yml index a0200da..55b5d91 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -24,6 +24,10 @@ x-airflow-common: QDRANT_HOST: qdrant QDRANT_PORT: '6333' QDRANT_API_KEY: ${QDRANT_API_KEY:-} + FALKORDB_HOST: falkordb + FALKORDB_PORT: '6379' + FALKORDB_USERNAME: ${FALKORDB_USERNAME:-} + FALKORDB_PASSWORD: ${FALKORDB_PASSWORD:-} volumes: - ./airflow_config/dags:/opt/airflow/dags - ./airflow_config/logs:/opt/airflow/logs @@ -45,6 +49,8 @@ x-airflow-common: condition: service_healthy qdrant: condition: service_healthy + falkordb: + condition: service_healthy services: postgres: @@ -82,6 +88,23 @@ services: volumes: - qdrant_data:/qdrant/storage + falkordb: + image: falkordb/falkordb:v4.4.4 + container_name: falkordb_container + healthcheck: + test: ["CMD", "redis-cli", "-p", "6379", "ping"] + interval: 10s + retries: 5 + start_period: 10s + restart: always + environment: + FALKORDB_USERNAME: ${FALKORDB_USERNAME:-} + FALKORDB_PASSWORD: ${FALKORDB_PASSWORD:-} + ports: + - "127.0.0.1:${FALKORDB_PORT:-6379}:6379" + volumes: + - falkordb_data:/data + airflow-apiserver: <<: *airflow-common container_name: airflow-apiserver @@ -244,4 +267,5 @@ services: volumes: pgvector_data: - qdrant_data: \ No newline at end of file + qdrant_data: + falkordb_data: \ No newline at end of file diff --git a/download_and_processing/files_processing.py b/download_and_processing/files_processing.py index fbb8711..452c263 100644 --- a/download_and_processing/files_processing.py +++ b/download_and_processing/files_processing.py @@ -12,6 +12,14 @@ from config import BASE_PATH, EMBEDDING_MODEL, SOURCE_MAP, config_file_path, get_logger from database import insert_data, refresh_table, remove_data +from database import ( + upsert_bofip_chunk, + upsert_bofip_node, + upsert_jade_chunk, + upsert_jade_node, + upsert_legi_chunk, + upsert_legi_node, +) from utils import ( CheckpointManager, CorpusHandler, @@ -728,6 +736,23 @@ def _process_dila_xml_content(root: ET.Element, file_name: str, model: str): if data_to_insert: insert_data(data=data_to_insert, table_name=table_name) + # Populate the knowledge graph in parallel with PostgreSQL + upsert_legi_node( + doc_id=cid, + nature=nature, + category=category, + ministry=ministry, + status=status, + title=title, + full_title=full_title, + number=number, + start_date=start_date, + end_date=end_date, + links=links, + ) + for chunk_id, *_ in data_to_insert: + upsert_legi_chunk(chunk_id=chunk_id, doc_id=cid) + except Exception as e: logger.error(f"Error processing file {file_name}: {e}") raise e @@ -987,6 +1012,20 @@ def _process_dila_xml_content(root: ET.Element, file_name: str, model: str): if data_to_insert: insert_data(data=data_to_insert, table_name=table_name) + # Populate the knowledge graph in parallel with PostgreSQL + upsert_jade_node( + doc_id=cid, + nature=nature, + solution=solution, + title=title, + number=number, + decision_date=decision_date, + jurisdiction=jurisdiction, + formation=formation, + ) + for chunk_id, *_ in data_to_insert: + upsert_jade_chunk(chunk_id=chunk_id, doc_id=cid) + except Exception as e: logger.error(f"Error processing file {file_name}: {e}") raise e @@ -1357,6 +1396,98 @@ def _process_dila_xml_content(root: ET.Element, file_name: str, model: str): logger.error(f"Error processing file {file_name}: {e}") raise e + elif file_name.startswith("BOFIPTEXT") and file_name.endswith(".xml"): + table_name = "bofip" + try: + cid = root.find(".//ID").text + nature_elem = root.find(".//NATURE") + nature = nature_elem.text if nature_elem is not None else None + title_elem = root.find(".//TITRE") + title = title_elem.text if title_elem is not None else None + category_elem = root.find(".//TYPE") + category = category_elem.text if category_elem is not None else None + date_elem = root.find(".//DATE_PUBLI") or root.find(".//DATE_TEXTE") + try: + date = ( + datetime.strptime(date_elem.text, "%Y-%m-%d").strftime("%Y-%m-%d") + if date_elem is not None and date_elem.text + else None + ) + except ValueError: + date = date_elem.text if date_elem is not None else None + + contenu = root.find(".//BLOC_TEXTUEL/CONTENU") + text_content = [] + if contenu is not None: + content = ET.tostring(contenu, encoding="unicode", method="xml") + content = "".join(ET.fromstring(content).itertext()) + lines = content.splitlines() + cleaned_lines = [line for line in lines if line] + content = "\n".join(cleaned_lines) + text_content.append(content) + text_content = "\n".join(text_content) + + chunks = make_chunks( + text=text_content, + chunk_size=1024, + chunk_overlap=0, + length_function=model, + ) + data_to_insert = [] + + for k, text in enumerate(chunks): + try: + chunk_index = k + 1 + chunk_text = f"{title}\n{text}" if title else text + + chunk_xxh64 = xxhash.xxh64( + chunk_text.encode("utf-8"), seed=2025 + ).hexdigest() + + embeddings = generate_embeddings_with_retry( + data=chunk_text, attempts=5, model=model + )[0] + + chunk_id = f"{cid}_{chunk_index}" + + new_data = ( + chunk_id, + cid, + chunk_index, + chunk_xxh64, + nature, + category, + title, + date, + text, + chunk_text, + embeddings, + ) + data_to_insert.append(new_data) + except PermissionDeniedError as e: + logger.error( + f"PermissionDeniedError (API key issue) for chunk {chunk_index} of file {file_name}: {e}" + ) + raise e + + if data_to_insert: + insert_data(data=data_to_insert, table_name=table_name) + + # Populate the knowledge graph in parallel with PostgreSQL + upsert_bofip_node( + doc_id=cid, + nature=nature, + category=category, + title=title, + date=date, + ) + for chunk_id, *_ in data_to_insert: + upsert_bofip_chunk(chunk_id=chunk_id, doc_id=cid) + + except Exception as e: + logger.error(f"Error processing file {file_name}: {e}") + raise e + def _handle_dila_suppression_list(lines: list[str], table_name: str, source_name: str): """ @@ -1974,6 +2105,41 @@ def process_data(table_name: str, streaming: bool = True, model: str = EMBEDDING logger.debug( f"Folder: {current_dir} successfully removed after processing" ) + elif attributes.get("type") == "bofip": + # BOFIP archives follow the same DILA tgz format as dila_folder + logger.info(f"Processing BOFIP files located in: {base_folder}") + all_entities = sorted( + [f for f in os.listdir(base_folder) if f.endswith(".tgz") or f.endswith(".tar.gz")] + ) + all_entities = [os.path.join(base_folder, f) for f in all_entities] + + for entity in all_entities: + try: + with tarfile.open(entity, "r:gz") as tar: + for member in tar.getmembers(): + if member.isfile() and os.path.basename( + member.name + ).startswith("liste_suppression"): + file_object = tar.extractfile(member) + if file_object: + with file_object as f: + lines = f.read().decode("utf-8").splitlines() + _handle_dila_suppression_list( + lines=lines, + table_name=table_name, + source_name=entity, + ) + break + except Exception as e: + logger.error( + f"Error while finding suppression list from archive {entity}: {e}" + ) + continue + + process_dila_xml_files( + source_path=entity, streaming=True, model=model + ) + logger.info(f"BOFIP file: {entity} successfully processed") else: logger.error(f"Unknown base folder '{base_folder}' for processing data.") raise ValueError( diff --git a/pyproject.toml b/pyproject.toml index 00ab64e..ca4879a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,9 @@ dependencies = [ "apache-airflow-providers-apprise==2.1.2", "xxhash==3.5.0", "transformers==4.57.3", - "sentence-transformers==3.4.1" + "sentence-transformers==3.4.1", + "falkordb==1.6.0", + "graphrag-sdk==0.8.2", ] [project.urls]