update
This commit is contained in:
+126
-33
@@ -177,6 +177,7 @@ function checkAgentResponse_JSON(responsedict::T1, requiredKeys::T2
|
||||
return (ispass, errormsg)
|
||||
end
|
||||
|
||||
|
||||
""" Convert a plain text string containing key-value pairs into a JSON-formatted string.
|
||||
|
||||
This function takes text containing key-value pairs (typically extracted from LLM responses)
|
||||
@@ -221,6 +222,96 @@ function clean_json_response(text::String)
|
||||
end
|
||||
|
||||
|
||||
|
||||
"""
|
||||
harvest_db_undirected_schema_graph(pg_conn_str::String)
|
||||
|
||||
Connects to a PostgreSQL instance, queries its metadata catalogs, and returns:
|
||||
1. `g::SimpleDiGraph`: The structural graph where nodes are tables.
|
||||
2. `id_to_table::Dict{Int, String}`: Maps numerical node IDs to real table names.
|
||||
3. `table_to_id::Dict{String, Int}`: Maps table names back to graph node IDs.
|
||||
"""
|
||||
function harvest_db_undirected_schema_graph(pg_conn_str)
|
||||
conn = LibPQ.Connection(pg_conn_str)
|
||||
try
|
||||
# 1. Fetch all user tables
|
||||
table_query = """
|
||||
SELECT c.relname AS table_name, c.oid
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE n.nspname = 'public' AND c.relkind = 'r';
|
||||
"""
|
||||
table_df = DataFrame(execute(conn, table_query))
|
||||
table_names = table_df.table_name
|
||||
|
||||
num_tables = length(table_names)
|
||||
table_to_id = Dict{String, Int}(name => i for (i, name) in enumerate(table_names))
|
||||
id_to_table = Dict{Int, String}(i => name for (i, name) in enumerate(table_names))
|
||||
|
||||
# CRITICAL: Use SimpleGraph (Undirected) so pathfinding can traverse both ways
|
||||
g = SimpleGraph(num_tables)
|
||||
|
||||
# 2. Strategy A: Extract Explicit Foreign Keys
|
||||
fk_query = """
|
||||
SELECT
|
||||
conrelid::regclass::text AS source_table,
|
||||
confrelid::regclass::text AS target_table
|
||||
FROM pg_constraint c
|
||||
JOIN pg_namespace n ON n.oid = c.connamespace
|
||||
WHERE c.contype = 'f' AND n.nspname = 'public';
|
||||
"""
|
||||
fk_df = DataFrame(execute(conn, fk_query))
|
||||
|
||||
for row in eachrow(fk_df)
|
||||
src = split(replace(row.source_table, "\"" => ""), '.')[end]
|
||||
tgt = split(replace(row.target_table, "\"" => ""), '.')[end]
|
||||
|
||||
if haskey(table_to_id, src) && haskey(table_to_id, tgt)
|
||||
add_edge!(g, table_to_id[src], table_to_id[tgt])
|
||||
end
|
||||
end
|
||||
|
||||
# 3. Strategy B: Fallback Name Matching (For missing explicit constraints)
|
||||
# Fetch all columns for all tables
|
||||
col_query = """
|
||||
SELECT a.attname AS column_name, c.relname AS table_name
|
||||
FROM pg_attribute a
|
||||
JOIN pg_class c ON c.oid = a.attrelid
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE n.nspname = 'public' AND c.relkind = 'r' AND a.attnum > 0 AND NOT a.attisdropped;
|
||||
"""
|
||||
col_df = DataFrame(execute(conn, col_query))
|
||||
|
||||
# Group columns by name to see who shares what fields
|
||||
for gdf in groupby(col_df, :column_name)
|
||||
col_name = gdf.column_name[1]
|
||||
|
||||
# We only infer connections on ID fields (e.g., seller_id, product_id)
|
||||
# Avoid matching generic names like 'id', 'created_at', or 'name'
|
||||
if endswith(lowercase(col_name), "_id") && col_name != "id"
|
||||
sharing_tables = gdf.table_name
|
||||
|
||||
# Connect all tables that share this ID column
|
||||
for i in 1:length(sharing_tables), j in (i+1):length(sharing_tables)
|
||||
t1 = sharing_tables[i]
|
||||
t2 = sharing_tables[j]
|
||||
|
||||
if haskey(table_to_id, t1) && haskey(table_to_id, t2)
|
||||
add_edge!(g, table_to_id[t1], table_to_id[t2])
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
return g, id_to_table, table_to_id
|
||||
|
||||
finally
|
||||
close(conn)
|
||||
end
|
||||
end
|
||||
|
||||
|
||||
|
||||
""" Extract vector metadata from PostgreSQL database.
|
||||
|
||||
Queries PostgreSQL system catalogs to extract column metadata including table names,
|
||||
@@ -261,36 +352,38 @@ products user_id integer Reference to user false true
|
||||
"""
|
||||
function extract_column_metadata(pg_conn_str::String)::DataFrame
|
||||
conn = LibPQ.Connection(pg_conn_str)
|
||||
|
||||
return extract_column_metadata(conn)
|
||||
end
|
||||
|
||||
function extract_column_metadata(conn::LibPQ.Connection)::DataFrame
|
||||
# This direct SQL query pulls the column specifications along with column-level descriptions
|
||||
query =
|
||||
"""
|
||||
SELECT
|
||||
c.relname AS column_name,
|
||||
format_type(a.atttypid, a.atttypmod) AS data_type,
|
||||
COALESCE(d.description, '') AS column_description,
|
||||
CASE WHEN pk.contype = 'p' THEN true ELSE false END AS is_primary_key,
|
||||
CASE WHEN AS table_name,
|
||||
a.attnamefk.contype = 'f' THEN true ELSE false END AS is_foreign_key
|
||||
FROM pg_attribute a
|
||||
JOIN pg_class c ON c.oid = a.attrelid
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
-- Join to fetch column comments/descriptions
|
||||
LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = a.attnum
|
||||
-- Check if column is part of a Primary Key
|
||||
LEFT JOIN pg_constraint pk ON pk.conrelid = c.oid
|
||||
AND pk.contype = 'p'
|
||||
AND a.attnum = ANY(pk.conkey)
|
||||
-- Check if column is part of a Foreign Key
|
||||
LEFT JOIN pg_constraint fk ON fk.conrelid = c.oid
|
||||
AND fk.contype = 'f'
|
||||
AND a.attnum = ANY(fk.conkey)
|
||||
WHERE
|
||||
n.nspname = 'public' -- Only user schemas
|
||||
AND c.relkind = 'r' -- Only standard tables
|
||||
AND a.attnum > 0 -- Skip system hidden columns
|
||||
AND NOT a.attisdropped; -- Skip dropped columns
|
||||
"""
|
||||
query = """
|
||||
SELECT
|
||||
c.relname AS table_name,
|
||||
a.attname AS column_name,
|
||||
format_type(a.atttypid, a.atttypmod) AS data_type,
|
||||
COALESCE(d.description, '') AS column_description,
|
||||
CASE WHEN pk.contype = 'p' THEN true ELSE false END AS is_primary_key,
|
||||
CASE WHEN fk.contype = 'f' THEN true ELSE false END AS is_foreign_key
|
||||
FROM pg_attribute a
|
||||
JOIN pg_class c ON c.oid = a.attrelid
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
-- Join to fetch column comments/descriptions
|
||||
LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = a.attnum
|
||||
-- Check if column is part of a Primary Key
|
||||
LEFT JOIN pg_constraint pk ON pk.conrelid = c.oid
|
||||
AND pk.contype = 'p'
|
||||
AND a.attnum = ANY(pk.conkey)
|
||||
-- Check if column is part of a Foreign Key
|
||||
LEFT JOIN pg_constraint fk ON fk.conrelid = c.oid
|
||||
AND fk.contype = 'f'
|
||||
AND a.attnum = ANY(fk.conkey)
|
||||
WHERE
|
||||
n.nspname = 'public' -- Only user schemas
|
||||
AND c.relkind = 'r' -- Only standard tables
|
||||
AND a.attnum > 0 -- Skip system hidden columns
|
||||
AND NOT a.attisdropped; -- Skip dropped columns
|
||||
"""
|
||||
|
||||
try
|
||||
# Execute and format into a clean DataFrame
|
||||
@@ -503,7 +596,7 @@ Extracts unique, non-null values from a specific column to build a local index f
|
||||
semantic search or entity resolution.
|
||||
|
||||
# Arguments
|
||||
- `conn_str::String`
|
||||
- `pg_conn_str::String`
|
||||
PostgreSQL connection string in LibPQ format (e.g., "host=hostname port=5432 dbname=database user=username password=secret")
|
||||
- `table::String`
|
||||
Table name to query
|
||||
@@ -531,12 +624,12 @@ julia> fruits = GeneralUtils.harvest_entity_catalog(conn, "products", "fruit_nam
|
||||
["Apple", "Banana", "Orange", "Mango"]
|
||||
```
|
||||
"""
|
||||
function harvest_entity_catalog(conn_str::String, table::String, column::String)::Vector{String}
|
||||
conn = LibPQ.Connection(conn_str)
|
||||
function harvest_entity_catalog(pg_conn_str::String, table::String, column::String)::Vector{String}
|
||||
conn = LibPQ.Connection(pg_conn_str)
|
||||
return harvest_entity_catalog(conn, table, column)
|
||||
end
|
||||
|
||||
function harvest_entity_catalog(conn, table::String, column::String)::Vector{String}
|
||||
function harvest_entity_catalog(conn::LibPQ.Connection, table::String, column::String)::Vector{String}
|
||||
|
||||
# We only care about unique, non-null values to keep the index fast and dense
|
||||
query = "SELECT DISTINCT $(column) FROM $(table) WHERE $(column) IS NOT NULL;"
|
||||
|
||||
Reference in New Issue
Block a user