Skip to content

feat: KG improvements #5048

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jul 18, 2025
Merged

feat: KG improvements #5048

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e7e564c
improvements
joachim-danswer Jul 18, 2025
c1c301f
drop views if SQL fails
joachim-danswer Jul 18, 2025
af254f7
mypy fix
joachim-danswer Jul 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 43 additions & 5 deletions backend/onyx/agents/agent_search/kb_search/nodes/a3_generate_simple_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -203,6 +203,8 @@ def generate_simple_sql(
if state.kg_entity_temp_view_name is None:
raise ValueError("kg_entity_temp_view_name is not set")

sql_statement_display: str | None = None

## STEP 3 - articulate goals

stream_write_step_activities(writer, _KG_STEP_NR)
Expand Down Expand Up @@ -381,7 +383,18 @@ def generate_simple_sql(

raise e

logger.debug(f"A3 - sql_statement after correction: {sql_statement}")
# display sql statement with view names replaced by general view names
sql_statement_display = sql_statement.replace(
state.kg_doc_temp_view_name, "<your_allowed_docs_view_name>"
)
sql_statement_display = sql_statement_display.replace(
state.kg_rel_temp_view_name, "<your_relationship_view_name>"
)
sql_statement_display = sql_statement_display.replace(
state.kg_entity_temp_view_name, "<your_entity_view_name>"
)

logger.debug(f"A3 - sql_statement after correction: {sql_statement_display}")

# Get SQL for source documents

Expand Down Expand Up @@ -409,7 +422,20 @@ def generate_simple_sql(
"relationship_table", rel_temp_view
)

logger.debug(f"A3 source_documents_sql: {source_documents_sql}")
if source_documents_sql:
source_documents_sql_display = source_documents_sql.replace(
state.kg_doc_temp_view_name, "<your_allowed_docs_view_name>"
)
source_documents_sql_display = source_documents_sql_display.replace(
state.kg_rel_temp_view_name, "<your_relationship_view_name>"
)
source_documents_sql_display = source_documents_sql_display.replace(
state.kg_entity_temp_view_name, "<your_entity_view_name>"
)
else:
source_documents_sql_display = "(No source documents SQL generated)"

logger.debug(f"A3 source_documents_sql: {source_documents_sql_display}")

scalar_result = None
query_results = None
Expand All @@ -435,7 +461,13 @@ def generate_simple_sql(
rows = result.fetchall()
query_results = [dict(row._mapping) for row in rows]
except Exception as e:
# TODO: raise error on frontend
logger.error(f"Error executing SQL query: {e}")
drop_views(
allowed_docs_view_name=doc_temp_view,
kg_relationships_view_name=rel_temp_view,
kg_entity_view_name=ent_temp_view,
)

raise e

Expand All @@ -459,8 +491,14 @@ def generate_simple_sql(
for source_document_result in query_source_document_results
]
except Exception as e:
# No stopping here, the individualized SQL query is not mandatory
# TODO: raise error on frontend

drop_views(
allowed_docs_view_name=doc_temp_view,
kg_relationships_view_name=rel_temp_view,
kg_entity_view_name=ent_temp_view,
)

logger.error(f"Error executing Individualized SQL query: {e}")

else:
Expand Down Expand Up @@ -493,11 +531,11 @@ def generate_simple_sql(
if reasoning:
stream_write_step_answer_explicit(writer, step_nr=_KG_STEP_NR, answer=reasoning)

if main_sql_statement:
if sql_statement_display:
stream_write_step_answer_explicit(
writer,
step_nr=_KG_STEP_NR,
answer=f" \n Generated SQL: {main_sql_statement}",
answer=f" \n Generated SQL: {sql_statement_display}",
)

stream_close_step_answer(writer, _KG_STEP_NR)
Expand Down
9 changes: 6 additions & 3 deletions backend/onyx/db/kg_temp_view.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import random

from sqlalchemy import text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session
Expand All @@ -17,10 +19,11 @@ def get_user_view_names(user_email: str, tenant_id: str) -> KGViewNames:
user_email_cleaned = (
user_email.replace("@", "__").replace(".", "_").replace("+", "_")
)
random_suffix_str = str(random.randint(1000000, 9999999))
return KGViewNames(
allowed_docs_view_name=f'"{tenant_id}".{KG_TEMP_ALLOWED_DOCS_VIEW_NAME_PREFIX}_{user_email_cleaned}',
kg_relationships_view_name=f'"{tenant_id}".{KG_TEMP_KG_RELATIONSHIPS_VIEW_NAME_PREFIX}_{user_email_cleaned}',
kg_entity_view_name=f'"{tenant_id}".{KG_TEMP_KG_ENTITIES_VIEW_NAME_PREFIX}_{user_email_cleaned}',
allowed_docs_view_name=f'"{tenant_id}".{KG_TEMP_ALLOWED_DOCS_VIEW_NAME_PREFIX}_{user_email_cleaned}_{random_suffix_str}',
kg_relationships_view_name=f'"{tenant_id}".{KG_TEMP_KG_RELATIONSHIPS_VIEW_NAME_PREFIX}_{user_email_cleaned}_{random_suffix_str}',
kg_entity_view_name=f'"{tenant_id}".{KG_TEMP_KG_ENTITIES_VIEW_NAME_PREFIX}_{user_email_cleaned}_{random_suffix_str}',
)


Expand Down
Loading