small updatesd

Author: joachim-danswer

backend/alembic/versions/495cb26ce93e_create_knowlege_graph_tables.py

@@ -46,18 +46,21 @@ def upgrade() -> None:
         op.execute(
             text(
                 f"""
-            DO $$
-            BEGIN
-                IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN
-                    EXECUTE format('CREATE USER %I WITH PASSWORD %L', '{DB_READONLY_USER}', '{DB_READONLY_PASSWORD}');
-                    -- Explicitly revoke all privileges including CONNECT
-                    EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}');
-                    -- Grant only the CONNECT privilege
-                    EXECUTE format('GRANT CONNECT ON DATABASE %I TO %I', current_database(), '{DB_READONLY_USER}');
-                END IF;
-            END
-            $$;
-        """
+                DO $$
+                BEGIN
+                    -- Check if the read-only user already exists
+                    IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN
+                        -- Create the read-only user with the specified password
+                        EXECUTE format('CREATE USER %I WITH PASSWORD %L', '{DB_READONLY_USER}', '{DB_READONLY_PASSWORD}');
+                        -- First revoke all privileges to ensure a clean slate
+                        EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}');
+                        -- Grant only the CONNECT privilege to allow the user to connect to the database
+                        -- but not perform any operations without additional specific grants
+                        EXECUTE format('GRANT CONNECT ON DATABASE %I TO %I', current_database(), '{DB_READONLY_USER}');
+                    END IF;
+                END
+                $$;
+                """
             )
         )
 
@@ -448,28 +451,34 @@ def upgrade() -> None:
         "ON kg_entity_extraction_staging USING GIN (clustering_name gin_trgm_ops)"
     )
 
-    if not MULTI_TENANT:
-        # Create trigger to update clustering columns if entity w/ doc_id is created
-        alphanum_pattern = r"[^a-z0-9]+"
-        op.execute(
+    # Create trigger to update clustering columns if entity w/ doc_id is created
+    alphanum_pattern = r"[^a-z0-9]+"
+    op.execute(
+        text(
             f"""
             CREATE OR REPLACE FUNCTION update_kg_entity_clustering()
             RETURNS TRIGGER AS $$
             DECLARE
                 doc_semantic_id text;
                 cleaned_semantic_id text;
+                max_length integer := 1000; -- Limit length for performance
             BEGIN
                 -- Get semantic_id from document
                 SELECT semantic_id INTO doc_semantic_id
                 FROM document
                 WHERE id = NEW.document_id;
 
-                -- Clean the semantic_id with regex patterns
+                -- Clean the semantic_id with regex patterns and handle NULLs
                 cleaned_semantic_id = regexp_replace(
-                    lower(COALESCE(doc_semantic_id, NEW.name)),
+                    lower(COALESCE(doc_semantic_id, NEW.name, '')),
                     '{alphanum_pattern}', '', 'g'
                 );
 
+                -- Truncate if too long for performance
+                IF length(cleaned_semantic_id) > max_length THEN
+                    cleaned_semantic_id = left(cleaned_semantic_id, max_length);
+                END IF;
+
                 -- Set clustering_name to cleaned version and generate trigrams
                 NEW.clustering_name = cleaned_semantic_id;
                 NEW.clustering_trigrams = show_trgm(cleaned_semantic_id);
@@ -478,50 +487,56 @@ def upgrade() -> None:
             $$ LANGUAGE plpgsql;
             """
         )
-        op.execute(
+    )
+    op.execute(
+        text(
             """
             CREATE OR REPLACE FUNCTION update_kg_entity_extraction_clustering()
             RETURNS TRIGGER AS $$
             DECLARE
                 doc_semantic_id text;
             BEGIN
                 -- Get semantic_id from document
+                -- If no document is found, doc_semantic_id will be NULL and COALESCE will use NEW.name
                 SELECT semantic_id INTO doc_semantic_id
                 FROM document
                 WHERE id = NEW.document_id;
 
                 -- Set clustering_name to semantic_id
-                NEW.clustering_name = lower(COALESCE(doc_semantic_id, NEW.name));
+                NEW.clustering_name = lower(COALESCE(doc_semantic_id, NEW.name, ''));
                 RETURN NEW;
             END;
             $$ LANGUAGE plpgsql;
             """
         )
-        for table, function in (
-            ("kg_entity", "update_kg_entity_clustering"),
-            ("kg_entity_extraction_staging", "update_kg_entity_extraction_clustering"),
-        ):
-            trigger = f"{function}_trigger"
-            op.execute(f"DROP TRIGGER IF EXISTS {trigger} ON {table}")
-            op.execute(
-                f"""
-                CREATE TRIGGER {trigger}
-                    BEFORE INSERT
-                    ON {table}
-                    FOR EACH ROW
-                    EXECUTE FUNCTION {function}();
-                """
-            )
-
-        # Create trigger to update kg_entity clustering_name and its trigrams when document.clustering_name changes
+    )
+    for table, function in (
+        ("kg_entity", "update_kg_entity_clustering"),
+        ("kg_entity_extraction_staging", "update_kg_entity_extraction_clustering"),
+    ):
+        trigger = f"{function}_trigger"
+        op.execute(f"DROP TRIGGER IF EXISTS {trigger} ON {table}")
         op.execute(
+            f"""
+            CREATE TRIGGER {trigger}
+                BEFORE INSERT
+                ON {table}
+                FOR EACH ROW
+                EXECUTE FUNCTION {function}();
+            """
+        )
+
+    # Create trigger to update kg_entity clustering_name and its trigrams when document.clustering_name changes
+    op.execute(
+        text(
             f"""
             CREATE OR REPLACE FUNCTION update_kg_entity_clustering_from_doc()
             RETURNS TRIGGER AS $$
             DECLARE
                 cleaned_semantic_id text;
             BEGIN
                 -- Clean the semantic_id with regex patterns
+                -- If semantic_id is NULL, COALESCE will use empty string
                 cleaned_semantic_id = regexp_replace(
                     lower(COALESCE(NEW.semantic_id, '')),
                     '{alphanum_pattern}', '', 'g'
@@ -538,11 +553,15 @@ def upgrade() -> None:
             $$ LANGUAGE plpgsql;
             """
         )
-        op.execute(
+    )
+    op.execute(
+        text(
             """
             CREATE OR REPLACE FUNCTION update_kg_entity_extraction_clustering_from_doc()
             RETURNS TRIGGER AS $$
             BEGIN
+                -- Update clustering name for all entities in staging referencing this document
+                -- If semantic_id is NULL, COALESCE will use empty string
                 UPDATE kg_entity_extraction_staging
                 SET
                     clustering_name = lower(COALESCE(NEW.semantic_id, ''))
@@ -552,21 +571,22 @@ def upgrade() -> None:
             $$ LANGUAGE plpgsql;
             """
         )
-        for function in (
-            "update_kg_entity_clustering_from_doc",
-            "update_kg_entity_extraction_clustering_from_doc",
-        ):
-            trigger = f"{function}_trigger"
-            op.execute(f"DROP TRIGGER IF EXISTS {trigger} ON document")
-            op.execute(
-                f"""
-                CREATE TRIGGER {trigger}
-                    AFTER UPDATE OF semantic_id
-                    ON document
-                    FOR EACH ROW
-                    EXECUTE FUNCTION {function}();
-                """
-            )
+    )
+    for function in (
+        "update_kg_entity_clustering_from_doc",
+        "update_kg_entity_extraction_clustering_from_doc",
+    ):
+        trigger = f"{function}_trigger"
+        op.execute(f"DROP TRIGGER IF EXISTS {trigger} ON document")
+        op.execute(
+            f"""
+            CREATE TRIGGER {trigger}
+                AFTER UPDATE OF semantic_id
+                ON document
+                FOR EACH ROW
+                EXECUTE FUNCTION {function}();
+            """
+        )
 
 
 def downgrade() -> None:

backend/alembic_tenants/versions/3b9f09038764_add_read_only_kg_user.py

@@ -31,18 +31,21 @@ def upgrade() -> None:
         op.execute(
             text(
                 f"""
-            DO $$
-            BEGIN
-                IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN
-                    EXECUTE format('CREATE USER %I WITH PASSWORD %L', '{DB_READONLY_USER}', '{DB_READONLY_PASSWORD}');
-                    -- Explicitly revoke all privileges including CONNECT
-                    EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}');
-                    -- Grant only the CONNECT privilege
-                    EXECUTE format('GRANT CONNECT ON DATABASE %I TO %I', current_database(), '{DB_READONLY_USER}');
-                END IF;
-            END
-            $$;
-        """
+                DO $$
+                BEGIN
+                    -- Check if the read-only user already exists
+                    IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN
+                        -- Create the read-only user with the specified password
+                        EXECUTE format('CREATE USER %I WITH PASSWORD %L', '{DB_READONLY_USER}', '{DB_READONLY_PASSWORD}');
+                        -- First revoke all privileges to ensure a clean slate
+                        EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}');
+                        -- Grant only the CONNECT privilege to allow the user to connect to the database
+                        -- but not perform any operations without additional specific grants
+                        EXECUTE format('GRANT CONNECT ON DATABASE %I TO %I', current_database(), '{DB_READONLY_USER}');
+                    END IF;
+                END
+                $$;
+                """
             )
         )
 

backend/onyx/agents/agent_search/kb_search/nodes/a3_generate_simple_sql.py

@@ -31,14 +31,24 @@
 from onyx.prompts.kg_prompts import SIMPLE_SQL_CORRECTION_PROMPT
 from onyx.prompts.kg_prompts import SIMPLE_SQL_PROMPT
 from onyx.prompts.kg_prompts import SOURCE_DETECTION_PROMPT
-from onyx.prompts.kg_prompts import SQL_AGGREGATION_REMOVAL_PROMPT
 from onyx.utils.logger import setup_logger
 from onyx.utils.threadpool_concurrency import run_with_timeout
 
 
 logger = setup_logger()
 
 
+def _drop_temp_views(
+    allowed_docs_view_name: str, kg_relationships_view_name: str
+) -> None:
+    with get_session_with_current_tenant() as db_session:
+        drop_views(
+            db_session,
+            allowed_docs_view_name=allowed_docs_view_name,
+            kg_relationships_view_name=kg_relationships_view_name,
+        )
+
+
 def _build_entity_explanation_str(entity_normalization_map: dict[str, str]) -> str:
     """
     Build a string of contextualized entities to avoid the model not being aware of
@@ -57,44 +67,12 @@ def _sql_is_aggregate_query(sql_statement: str) -> bool:
     )
 
 
-def _remove_aggregation(sql_statement: str, llm: LLM) -> str:
-    """
-    Remove aggregate functions from the SQL statement.
-    """
-
-    sql_aggregation_removal_prompt = SQL_AGGREGATION_REMOVAL_PROMPT.replace(
-        "---sql_statement---", sql_statement
-    )
-
-    msg = [
-        HumanMessage(
-            content=sql_aggregation_removal_prompt,
-        )
-    ]
-
-    try:
-        llm_response = run_with_timeout(
-            KG_SQL_GENERATION_TIMEOUT,
-            llm.invoke,
-            prompt=msg,
-            timeout_override=25,
-            max_tokens=800,
-        )
-
-        cleaned_response = (
-            str(llm_response.content).replace("```json\n", "").replace("\n```", "")
-        )
-        sql_statement = cleaned_response.split("<sql>")[1].split("</sql>")[0].strip()
-        sql_statement = sql_statement.replace("sql", "").strip()
-
-    except Exception as e:
-        logger.error(f"Error in strategy generation: {e}")
-        raise e
-
-    return sql_statement
-
-
-def _get_source_documents(sql_statement: str, llm: LLM) -> str | None:
+def _get_source_documents(
+    sql_statement: str,
+    llm: LLM,
+    allowed_docs_view_name: str,
+    kg_relationships_view_name: str,
+) -> str | None:
     """
     Generate SQL to retrieve source documents based on the input sql statement.
     """
@@ -132,6 +110,11 @@ def _get_source_documents(sql_statement: str, llm: LLM) -> str | None:
             )
         else:
             logger.error(f"Could not generate source documents SQL: {e}")
+
+        _drop_temp_views(
+            allowed_docs_view_name=allowed_docs_view_name,
+            kg_relationships_view_name=kg_relationships_view_name,
+        )
         return None
 
     return sql_statement
@@ -282,6 +265,11 @@ def generate_simple_sql(
 
         except Exception as e:
             logger.error(f"Error in strategy generation: {e}")
+
+            _drop_temp_views(
+                allowed_docs_view_name=allowed_docs_view_name,
+                kg_relationships_view_name=kg_relationships_view_name,
+            )
             raise e
 
         logger.debug(f"A3 - sql_statement: {sql_statement}")
@@ -321,13 +309,24 @@ def generate_simple_sql(
             logger.error(
                 f"Error in generating the sql correction: {e}. Original model response: {cleaned_response}"
             )
+
+            _drop_temp_views(
+                allowed_docs_view_name=allowed_docs_view_name,
+                kg_relationships_view_name=kg_relationships_view_name,
+            )
+
             raise e
 
         logger.debug(f"A3 - sql_statement after correction: {sql_statement}")
 
         # Get SQL for source documents
 
-        source_documents_sql = _get_source_documents(sql_statement, llm=primary_llm)
+        source_documents_sql = _get_source_documents(
+            sql_statement,
+            llm=primary_llm,
+            allowed_docs_view_name=allowed_docs_view_name,
+            kg_relationships_view_name=kg_relationships_view_name,
+        )
 
         logger.info(f"A3 source_documents_sql: {source_documents_sql}")
 
@@ -373,12 +372,10 @@ def generate_simple_sql(
         else:
             source_document_results = None
 
-        with get_session_with_current_tenant() as db_session:
-            drop_views(
-                db_session,
-                allowed_docs_view_name=allowed_docs_view_name,
-                kg_relationships_view_name=kg_relationships_view_name,
-            )
+        _drop_temp_views(
+            allowed_docs_view_name=allowed_docs_view_name,
+            kg_relationships_view_name=kg_relationships_view_name,
+        )
 
         logger.info(f"A3 - Number of query_results: {len(query_results)}")