From edad4422427071a52024e76713cfb2615189564d Mon Sep 17 00:00:00 2001 From: narawat Date: Mon, 13 Jul 2026 21:06:09 +0700 Subject: [PATCH] update --- etc.jl | 355 ++++++++++++++++++++++++++++ src/llmUtil.jl | 614 ++++++++++++++++++++++--------------------------- 2 files changed, 631 insertions(+), 338 deletions(-) create mode 100644 etc.jl diff --git a/etc.jl b/etc.jl new file mode 100644 index 0000000..894b9a3 --- /dev/null +++ b/etc.jl @@ -0,0 +1,355 @@ + + +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" + + + + + + + + + + + + + + diff --git a/src/llmUtil.jl b/src/llmUtil.jl index 95dda24..be57fba 100644 --- a/src/llmUtil.jl +++ b/src/llmUtil.jl @@ -1,103 +1,12 @@ module llmUtil -export formatLLMtext, formatLLMtext_llama3instruct, jsoncorrection, deFormatLLMtext, extractthink, +export formatLLMtext, extractthink, checkAgentResponse_JSON, clean_json_response -using UUIDs, JSON, Dates +using UUIDs, JSON, Dates, DataFrames using GeneralUtils # ---------------------------------------------- 100 --------------------------------------------- # -#[PENDING] update code to use JSON - -""" Convert a single chat dictionary into LLM model instruct format. - -# Llama 3 instruct format example - <|begin_of_text|> - <|start_header_id|>system<|end_header_id|> - You are a helpful assistant. - <|eot_id|> - <|start_header_id|>user<|end_header_id|> - Get me an icecream. - <|eot_id|> - <|start_header_id|>assistant<|end_header_id|> - Go buy it yourself at 7-11. - <|eot_id|> - -# Arguments - - `name::T` - message owner name e.f. "system", "user" or "assistant" - - `text::T` - -# Return - - `formattedtext::String` - text formatted to model format - -# Example -```jldoctest -julia> using Revise -julia> using YiemAgent -julia> d = Dict(:name=> "system",:text=> "You are a helpful, respectful and honest assistant.",) -julia> formattedtext = YiemAgent.formatLLMtext_llama3instruct(d[:name], d[:text]) -"<|begin_of_text|>\n <|start_header_id|>system<|end_header_id|>\n You are a helpful, respectful and honest assistant.\n <|eot_id|>\n" -``` - -Signature -""" -function formatLLMtext_llama3instruct(name::T, text::T; - assistantStarter::Bool=false) where {T<:AbstractString} - formattedtext = - if name == "system" - """ - <|start_header_id|>$name<|end_header_id|> - $text - <|eot_id|> - """ - else - """ - <|start_header_id|>$name<|end_header_id|> - $text - <|eot_id|> - """ - end - - if assistantStarter - formattedtext *= - """ - <|start_header_id|>assistant<|end_header_id|> - """ - end - - return formattedtext -end - - -function formatLLMtext_qwen(name::T, text::T; - assistantStarter::Bool=false) where {T<:AbstractString} - formattedtext = - if name == "system" - """ - <|im_start|>$name - $text - <|im_end|> - """ - else - """ - <|im_start|>$name - $text - <|im_end|> - """ - end - - if assistantStarter - formattedtext *= - """ - <|im_start|>assistant - """ - end - - return formattedtext -end - function formatLLMtext_qwen3(name::T, text::T; assistantStarter::Bool=false) where {T<:AbstractString} @@ -127,59 +36,6 @@ function formatLLMtext_qwen3(name::T, text::T; end -function formatLLMtext_phi4(name::T, text::T; - assistantStarter::Bool=false) where {T<:AbstractString} - formattedtext = - if name == "system" - """ - <|system|> - $text - <|end|> - """ - else - """ - <|assistant|> - $text - <|end|> - """ - end - - if assistantStarter - formattedtext *= - """ - <|assistant|> - """ - end - - return formattedtext -end - - -function formatLLMtext_granite3(name::T, text::T; - assistantStarter::Bool=false) where {T<:AbstractString} - formattedtext = - if name == "system" - """ - <|start_of_role|>system<|end_of_role|>{$text}<|end_of_text|> - """ - else - """ - <|start_of_role|>$name<|end_of_role|>{$text}<|end_of_text|> - """ - end - - if assistantStarter - formattedtext *= - """ - <|start_of_role|>assistant<|end_of_role|>{ - """ - end - - return formattedtext -end - - - """ Convert a vector of chat message dictionaries into LLM model instruct format. # Arguments @@ -194,13 +50,13 @@ end # Example ```jldoctest julia> using Revise -julia> using YiemAgent +julia> using GeneralUtils julia> chatmessage = [ Dict(:name=> "system",:text=> "You are a helpful, respectful and honest assistant.",), Dict(:name=> "user",:text=> "list me all planets in our solar system.",), Dict(:name=> "assistant",:text=> "I'm sorry. I don't know. You tell me.",), ] -julia> formattedtext = YiemAgent.formatLLMtext(chatmessage, "llama3instruct") +julia> formattedtext = GeneralUtils.formatLLMtext(chatmessage, "llama3instruct") "<|begin_of_text|>\n <|start_header_id|>system<|end_header_id|>\n You are a helpful, respectful and honest assistant.\n <|eot_id|>\n <|start_header_id|>user<|end_header_id|>\n list me all planets in our solar system.\n <|eot_id|>\n <|start_header_id|>assistant<|end_header_id|>\n I'm sorry. I don't know. You tell me.\n <|eot_id|>\n" ``` """ @@ -237,192 +93,6 @@ function formatLLMtext(messages::Vector{Dict{Symbol, T}}, formatname::String return str end -""" Revert LLM-format response back into regular text. - -# Arguments - - `text::String` - The LLM formatted string to be converted. - -# Return - - `normalText::String` - The original plain text extracted from the given LLM-formatted string. - -# Example -```jldoctest -julia> using Revise -julia> using YiemAgent -julia> response = "<|begin_of_text|>This is a sample system instruction.<|eot_id|>" -julia> normalText = YiemAgent.deFormatLLMtext(response, "granite3") -"This is a sample system instruction." -``` -""" -function deFormatLLMtext(text::String, formatname::String; includethink::Bool=false - )::String - f = - if formatname == "granite3" - deFormatLLMtext_granite3 - elseif formatname == "qwen3" - deFormatLLMtext_qwen3 - else - error("$formatname template not define yet") - end - - r = f(text) - result = r === nothing ? text : r - return result -end - - -""" Revert LLM-format response back into regular text for Granite 3 format. - -# Arguments - - `text::String` - The LLM formatted string to be converted. - -# Return - - `normalText::Union{Nothing, String}` - The original plain text extracted from the given LLM-formatted string. - Returns nothing if the text is not in Granite 3 format. - -# Example -```jldoctest -julia> using Revise -julia> using YiemAgent -julia> response = "{This is a sample LLM response.}" -julia> normalText = YiemAgent.deFormatLLMtext(response, "granite3") -"This is a sample LLM response." -""" -function deFormatLLMtext_granite3(text::String)::Union{Nothing, String} - # check if '{' and '}' are in the text because it's a special format for the LLM response - if contains(text, "<|im_start|>assistant") - # get the text between '{' and '}' - text_between_braces = GeneralUtils.extractTextBetweenCharacter(text, '{', '}')[1] - return text_between_braces - elseif text[end] == '}' - text = "{$text" - text_between_braces = GeneralUtils.extractTextBetweenCharacter(text, '{', '}')[1] - else - return nothing - end -end - - -function deFormatLLMtext_qwen3(text::String)::Union{Nothing, String} - return text -end - -# function deFormatLLMtext_qwen3(text::String; includethink::Bool=false)::Union{Nothing, String} -# think = nothing -# str = nothing - -# if occursin("", text) -# r = GeneralUtils.extractTextBetweenString(text, "", "") -# if r[:success] -# think = r[:text] -# end -# str = string(split(text, "")[2]) -# end - -# if includethink == true && occursin("", text) -# result = "ModelThought: $think $str" -# return result -# elseif includethink == false && occursin("", text) -# result = str -# return result -# else -# return text -# end -# end - - -""" Attemp to correct LLM response's incorrect JSON response. - -# Arguments - - `a::T1` - one of Yiem's agent - - `input::T2` - text to be send to virtual wine customer - -# Return - - `correctjson::String` - corrected json string - -# Example -```jldoctest -julia> -``` - -# Signature -""" -function jsoncorrection(config::T1, input::T2, correctJsonExample::T3; - maxattempt::Integer=3 - ) where {T1<:AbstractDict, T2<:AbstractString, T3<:AbstractString} - - incorrectjson = deepcopy(input) - correctjson = nothing - - for attempt in 1:maxattempt - try - d = copy(JSON3.read(incorrectjson)) - correctjson = incorrectjson - return correctjson - catch e - @warn "Attempting to correct JSON string. Attempt $attempt" - e = """$e""" - if occursin("EOF", e) - e = split(e, "EOF")[1] * "EOF" - end - incorrectjson = deepcopy(input) - _prompt = - """ - Your goal are: - 1) Use the expected JSON format as a guideline to check why the given JSON string failed to load and provide a corrected version that can be loaded by Python's json.load function. - 2) Provide Corrected JSON string only. Do not provide any other info. - - $correctJsonExample - - Let's begin! - Given JSON string: $incorrectjson - The given JSON string failed to load previously because: $e - Corrected JSON string: - """ - - # apply LLM specific instruct format - externalService = config[:externalservice][:text2textinstruct] - llminfo = externalService[:llminfo] - prompt = - if llminfo[:name] == "llama3instruct" - formatLLMtext_llama3instruct("system", _prompt) - else - error("llm model name is not defied yet $(@__LINE__)") - end - - # send formatted input to user using GeneralUtils.sendReceiveMqttMsg - msgMeta = GeneralUtils.generate_msgMeta( - externalService[:mqtttopic], - senderName= "jsoncorrection", - senderId= uuid4snakecase(), - receiverName= "text2textinstruct", - mqttBroker= config[:mqttServerInfo][:broker], - mqttBrokerPort= config[:mqttServerInfo][:port], - ) - - outgoingMsg = Dict( - :msgMeta=> msgMeta, - :payload=> Dict( - :text=> prompt, - :kwargs=> Dict( - :max_tokens=> 512, - :stop=> ["<|eot_id|>"], - ) - ) - ) - result = GeneralUtils.sendReceiveMqttMsg(outgoingMsg; timeout=120) - incorrectjson = result[:response][:text] - end - end -end - function extractthink(text::String) think = nothing @@ -470,18 +140,18 @@ The validation logic checks: # Example ```julia -julia> using YiemAgent +julia> using GeneralUtils julia> requiredKeys = ["wine_name", "price", "rating"] julia> response = Dict("wine_name"=>"Château Margaux", "price"=>250.0, "rating"=>98) -julia> YiemAgent.checkAgentResponse_JSON(response, requiredKeys) +julia> GeneralUtils.checkAgentResponse_JSON(response, requiredKeys) (true, nothing) julia> response_missing = Dict("wine_name"=>"Château Margaux", "price"=>250.0) -julia> YiemAgent.checkAgentResponse_JSON(response_missing, requiredKeys) +julia> GeneralUtils.checkAgentResponse_JSON(response_missing, requiredKeys) (false, "rating are missing from your previous response") julia> response_extra = Dict("wine_name"=>"Château Margaux", "price"=>250.0, "rating"=>98, "extra_field"=>"data") -julia> YiemAgent.checkAgentResponse_JSON(response_extra, requiredKeys) +julia> GeneralUtils.checkAgentResponse_JSON(response_extra, requiredKeys) (false, "Your previous attempt has duplicated points according to the required response format") ``` """ @@ -550,12 +220,280 @@ function clean_json_response(text::String) end +""" Extract vector metadata from PostgreSQL database. + +Queries PostgreSQL system catalogs to extract column metadata including table names, +column names, data types, descriptions, and constraint information (primary/foreign keys). +Returns a DataFrame designed for vector embedding generation. + +# Arguments +- `pg_conn_str::String` + PostgreSQL connection string (e.g., "postgresql://user:pass@host:port/dbname") + +# Return +- `DataFrame` + A DataFrame with columns: + - `table_name`: Name of the table + - `column_name`: Name of the column + - `data_type`: PostgreSQL data type with typemod + - `column_description`: Column's comment/description (empty string if none) + - `is_primary_key`: Boolean indicating if column is part of primary key + - `is_foreign_key`: Boolean indicating if column is part of foreign key + - `vector_id`: Generated unique identifier in format `col_{table_name}_{column_name}` + +# Example +```julia +julia> using GeneralUtils +julia> pg_conn = "postgresql://user:pass@localhost:5432/mydb" +julia> df = GeneralUtils.extract_vector_metadata(pg_conn) +DataFrame +6 rows × 7 columns +table_name column_name data_type column_description is_primary_key is_foreign_key vector_id +─────────────┬───────────┬───────────┬───────────────────┬───────────────┬───────────────┬──────────────────────── +users id integer User ID true false col_users_id +users name text User name false false col_users_name +users email text User email false false col_users_email +products id integer Product ID true false col_products_id +products price numeric Product price false false col_products_price +products user_id integer Reference to user false true col_products_user_id +``` +""" +function extract_vector_metadata(pg_conn_str::String)::DataFrame + conn = LibPQ.Connection(pg_conn_str) + + # 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 + """ + + 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 from vector metadata DataFrame. + +Transforms the metadata DataFrame from `extract_vector_metadata` into a vector of +dictionaries structured for vector embedding storage and retrieval. + +# Arguments +- `df::DataFrame` + A DataFrame with columns from `extract_vector_metadata`: `table_name`, `column_name`, + `data_type`, `column_description`, `is_primary_key`, `is_foreign_key`, `vector_id`. + +# Return +- `Vector{Dict}` + A vector of dictionaries with keys: + - `id`: The vector ID from `vector_id` + - `text_content`: A structured text string combining table, column, type, constraint + markers, and description + - `metadata`: A dictionary containing `table`, `column`, and `type` + +# Details +The function constructs rich text payloads by: +1. Building a base string with table name, column name, and data type +2. Appending constraint markers `[PRIMARY KEY]` or `[FOREIGN KEY RELATIONAL LINK]` +3. Adding column description if available, otherwise generating a default description + +# Example +```julia +julia> using GeneralUtils +julia> pg_conn = "postgresql://user:pass@localhost:5432/mydb" +julia> df = GeneralUtils.extract_vector_metadata(pg_conn) +julia> payloads = GeneralUtils.generate_embedding_payloads(df) +3-element Vector{Dict}: + Dict("id" => "col_users_id", "text_content" => "Table: users | Column: id | Type: integer [PRIMARY KEY] | Description: User ID", "metadata" => Dict("table" => "users", "column" => "id", "type" => "integer")) + Dict("id" => "col_users_name", "text_content" => "Table: users | Column: name | Type: text | Description: User name", "metadata" => Dict("table" => "users", "column" => "name", "type" => "text")) + Dict("id" => "col_users_email", "text_content" => "Table: users | Column: email | Type: text | Description: User email", "metadata" => Dict("table" => "users", "column" => "email", "type" => "text")) +``` +""" +function generate_embedding_payloads(df::DataFrame)::Vector{Dict} + 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 from vector hits using graph traversal. +Finds the minimum interconnected subgraph that connects all semantically matched tables +from vector search results, enabling construction of valid SQL queries across related tables. +# Arguments +- `vector_hits::Vector{String}` + A vector of table names matched semantically from Stage 1 vector search. +- `g::SimpleGraph` + An undirected graph representing table relationships (nodes=tables, edges=foreign key relations). +- `table_to_id::Dict{String, Int}` + Mapping from table names to node IDs in the graph. +- `id_to_table::Dict{Int, String}` + Reverse mapping from node IDs to table names. + +# Return +- `Vector{String}` + A vector of table names representing the minimum subgraph that connects all input vector hits. + The order reflects traversal path from anchor node to connected components. + +# Details +The algorithm: +1. Validates vector hits against the graph's table mapping +2. Handles edge cases: empty hits, single table (returns as-is) +3. Uses the highest-ranked vector hit as anchor node +4. For each remaining hit, finds shortest path to current subgraph using Dijkstra's algorithm +5. Builds minimal connected subgraph containing all hits +6. Returns table names in traversal order + +# Example +```julia +julia> using GeneralUtils, Graphs +julia> g = SimpleGraph(5) +julia> add_edge!(g, 1, 2) +julia> add_edge!(g, 2, 3) +julia> add_edge!(g, 3, 4) +julia> table_to_id = Dict("users" => 1, "orders" => 2, "payments" => 3, "products" => 4, "inventory" => 5) +julia> id_to_table = Dict(1 => "users", 2 => "orders", 3 => "payments", 4 => "products", 5 => "inventory") +julia> vector_hits = ["users", "products"] +julia> GeneralUtils.resolve_semantic_cluster(vector_hits, g, table_to_id, id_to_table) +["users", "orders", "products"] +``` +""" +function resolve_semantic_cluster( + vector_hits::Vector{String}, + g::SimpleGraph, + table_to_id::Dict{String, Int}, + id_to_table::Dict{Int, String} + )::Vector{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