This commit is contained in:
2026-07-13 21:23:21 +07:00
parent d658d9a25b
commit f28405f3f1
-355
View File
@@ -1,355 +0,0 @@
using LibPQ
using DataFrames
"""
extract_vector_metadata(pg_conn_str::String) -> DataFrame
Queries PostgreSQL system catalogs to extract a rich semantic text map of every
column in the database. Returns a DataFrame designed for vector embedding generation.
"""
function extract_vector_metadata(pg_conn_str::String)
conn = LibPQ.Connection(pg_conn_str)
# This direct SQL query pulls the column specifications along with column-level descriptions
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
result = execute(conn, query)
df = DataFrame(result)
# Create a unique document ID for each vector row
df.vector_id = ["col_\$(row.table_name)_\$(row.column_name)" for row in eachrow(df)]
return df
finally
close(conn)
end
end
"""
generate_embedding_payloads(df::DataFrame) -> Vector{Dict}
Transforms the metadata DataFrame into structured text strings optimal for
vector space mapping.
"""
function generate_embedding_payloads(df::DataFrame)
payloads = Dict[]
for row in eachrow(df)
# 1. Build a rich text description summarizing the column's role
text_payload = "Table: $(row.table_name) | Column: $(row.column_name) | Type: $(row.data_type)"
if row.is_primary_key
text_payload *= " [PRIMARY KEY]"
end
if row.is_foreign_key
text_payload *= " [FOREIGN KEY RELATIONAL LINK]"
end
# Append business descriptions if they exist in the database comments
if !isempty(strip(row.column_description))
text_payload *= " | Description: $(row.column_description)"
else
text_payload *= " | Description: Represents $(row.column_name) data fields within the $(row.table_name) architecture."
end
# 2. Package everything neatly to be passed to your vector store client
push!(payloads, Dict(
"id" => row.vector_id,
"text_content" => text_payload,
"metadata" => Dict(
"table" => row.table_name,
"column" => row.column_name,
"type" => row.data_type
)
))
end
return payloads
end
"""
resolve_semantic_cluster(vector_hits::Vector{String}, g::SimpleGraph, table_to_id::Dict{String, Int}, id_to_table::Dict{Int, String}) -> Vector{String}
Takes a scattered array of semantically matched tables from Stage 1, navigates
the undirected network structure, and isolates the minimum interconnected subgraph
required to weave ALL hits into a single valid SQL query.
"""
function resolve_semantic_cluster(
vector_hits::Vector{String},
g::SimpleGraph,
table_to_id::Dict{String, Int},
id_to_table::Dict{Int, String}
)
# Filter out hits that don't exist in our actual database graph mapping
valid_node_ids = Int[]
for hit in vector_hits
if haskey(table_to_id, hit)
push!(valid_node_ids, table_to_id[hit])
else
@warn "Vector hit '$hit' does not map to an existing database table."
end
end
unique!(valid_node_ids)
# Edge Case Handlers
if isempty(valid_node_ids)
return String[]
elseif length(valid_node_ids) == 1
return [id_to_table[valid_node_ids[1]]]
end
# The Isolated Subgraph Set to build our final context
schema_subgraph_nodes = Set{Int}()
# Phase A: Select an initial anchor component. We use the highest-ranked vector hit.
anchor_node = valid_node_ids[1]
push!(schema_subgraph_nodes, anchor_node)
# Phase B: Sequentially route paths to all other semantic coordinates
for target_node in valid_node_ids[2:end]
# Skip if an earlier loop trajectory already naturally absorbed this table
if target_node in schema_subgraph_nodes
continue
end
# Calculate the shortest path tree from the CURRENT state of our subgraph
# We find the shortest path from the target back to ANY node currently in our tree
shortest_paths = dijkstra_shortest_paths(g, target_node)
# Find which node currently in our subgraph is closest to the target node
closest_subgraph_node = 0
min_distance = Inf
for subgraph_node in schema_subgraph_nodes
dist = shortest_paths.dists[subgraph_node]
if dist < min_distance
min_distance = dist
closest_subgraph_node = subgraph_node
end
end
# Reconstruct the path from the target node to the closest point on our existing tree
if closest_subgraph_node != 0
curr = closest_subgraph_node
while curr != 0
push!(schema_subgraph_nodes, curr)
curr = shortest_paths.parents[curr]
if curr == target_node
push!(schema_subgraph_nodes, target_node)
break
end
end
end
end
# Map the unique structural nodes back to clean table names
return [id_to_table[node_id] for node_id in schema_subgraph_nodes]
end
function get_embedding(nats_conn::NATS.Connection, text::AbstractArray{String})
documents_dict = Dict("documents" => text)
payloads = [("documents", documents_dict, "dictionary")]
_, msg_envelope_json_str = msghandler.smartpack(
config["externalservice"]["servicesloadbalancer"]["nats"],
payloads;
msg_purpose="embedding",
broker_url=config["nats_server_info"]["url"],
fileserver_url=config["externalservice"]["fileserver"]["url"])
reply = NATS.request(nats_conn,
config["externalservice"]["servicesloadbalancer"]["nats"],
msg_envelope_json_str, timeout=120)
incoming_env_json_str = String(reply.payload)
incoming_env = msghandler.smartunpack(incoming_env_json_str)
embedding_response = incoming_env["payloads"][1][2]
return embedding_response
end
nats_conn = NATS.connect(config["nats_server_info"]["url"])
# Run the extractor
metadata_df = extract_vector_metadata(pg_conn_str)
embedding_ready = generate_embedding_payloads(metadata_df)
println(embedding_ready[1]["text_content"])
# Output: "Table: join_table | Column: seller_id | Type: integer [PRIMARY KEY] [FOREIGN KEY RELATIONAL LINK] | Description: Links unique sellers to their corresponding product items."
embedding_ready_2 = [i["text_content"] for i in embedding_ready]
table_embedding = get_embedding(nats_conn, embedding_ready_2)
user_question =
"""
Retrieves ["winery", "wine_name", "wine_id", "vintage", "region", "country", "wine_type", "grape", "serving_temperature", "sweetness", "intensity", "tannin", "acidity", "tasting_notes", "price", "currency", "image_url", "retailer_name", "retailer_id"] of wines that match the following criteria - {wine_name: Montrachet Grand Cru, winery: Domaine Jacques Prieur, region: Montrachet, country: France, , retailer_name: Yiem Wines Ltd, retailerid: f54eab6b-7650-4448-b009-c53f3efbcc3b}
"""
user_question_embedding = get_embedding(nats_conn, [user_question])
using Distances
similarity = 1 - cosine_dist(Float64.(table_embedding["data"][1]["embedding"]),
Float64.(user_question_embedding["data"][1]["embedding"])
)
user_question_embedding = Float64.(user_question_embedding["data"][1]["embedding"])
user_question_similarity = []
for i in table_embedding["data"]
i_data = i["embedding"]
i_float = Float64.(i_data)
r = 1 - cosine_dist(i_float, user_question_embedding)
push!(user_question_similarity, r)
end
new_df = hcat(metadata_df, DataFrame(user_question_similarity = user_question_similarity))
sorted_df = sort(new_df, :user_question_similarity, rev=true) # sort max to min
# top 20 of sorted_df get this tables
vector_hits = ["retailer_wine", "wine", "wine_food", "retailer"]
g, id_to_table, table_to_id = harvest_undirected_schema_graph(pg_conn_str)
# tables that I should put schema in LLM context
optimized_context = resolve_semantic_cluster(vector_hits, g, table_to_id, id_to_table)
function related_tables_for_user_question()
end
# ---------------------------------------------- 100 --------------------------------------------- #
# Agent 3 (The Entity Resolver): Instantly runs a fast, local token search (like BM25) to map messy user text (like HandOld) to the exact database string (Hand Old Bar & Grill) before the SQL is drafted.
using StringDistances
"""
harvest_entity_catalog(conn_str::String, table::String, column::String) -> Vector{String}
Pulls unique, clean text strings from a specific entity column to build a local index.
"""
function harvest_entity_catalog(conn_str::String, table::String, column::String)
conn = LibPQ.Connection(conn_str)
# 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;"
try
df = DataFrame(execute(conn, query))
# Return as a clean array of strings
return String.(strip.(df[:, 1]))
finally
close(conn)
end
end
"""
resolve_entity(messy_input::String, catalog::Vector{String}; threshold=0.6) -> String
Parses user text, matches it against the real database catalog, and returns
the exact string found in the database. Returns an empty string if no confident match.
"""
function resolve_entity(messy_input::String, catalog::Vector{String}; threshold=0.5)
best_match = ""
highest_score = 0.0
# Normalize input text to ensure case-insensitive matching
clean_input = lowercase(strip(messy_input))
for real_string in catalog
clean_real = lowercase(real_string)
# Calculate phonetic/structural similarity score (0.0 to 1.0)
# JaroWinkler is optimized for short strings, names, and partial acronyms
score = compare(clean_real, clean_input, JaroWinkler())
# Substring/Token fallback: handle cases like "HandOld" matching "Hand Old Bar & Grill"
# We strip spaces to check if the user just compressed words together
if contains(replace(clean_real, " " => ""), clean_input)
score = max(score, 0.85)
end
if score > highest_score
highest_score = score
best_match = real_string
end
end
# Only return if we cross our safety confidence barrier
if highest_score >= threshold
return best_match
end
return "" # No confident match found
end
winery_catalog = harvest_entity_catalog(conn_str, "wine", "winery")
# Let's assume the catalog contains: ["Hand Old Bar & Grill", "Bangkok Diner", "Phuket Seafood"]
# 2. The user asks a messy question with a typo and compressed text
user_question = "What are the total sales at HandOld last week?"
# 3. Agent 3 isolates potential nouns or scans the question against the index
# We look for words that don't match standard english dictionary tokens, or check the full string segments
detected_entity = "Jacob"
# 4. Run the resolution engine
exact_db_string = resolve_entity(detected_entity, winery_catalog)
# "United States"
println("Messy Input: ", detected_entity)
println("Resolved Engine Value: ", exact_db_string)
# Output: "Hand Old Bar & Grill"