1032 lines
38 KiB
Julia
1032 lines
38 KiB
Julia
module llmUtil
|
||
|
||
export formatLLMtext, extractthink, checkAgentResponse_JSON, clean_json_response,
|
||
extract_column_metadata, generate_embedding_payloads, resolve_semantic_cluster,
|
||
harvest_entity_catalog, resolve_entity, harvest_db_undirected_schema_graph,
|
||
get_db_table_schema
|
||
|
||
using UUIDs, JSON, Dates, DataFrames, StringDistances, Graphs, LibPQ
|
||
using ..util
|
||
|
||
# ---------------------------------------------- 100 --------------------------------------------- #
|
||
|
||
function formatLLMtext_qwen3(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
|
||
|
||
|
||
""" Convert a vector of chat message dictionaries into LLM model instruct format.
|
||
|
||
# Arguments
|
||
- `messages::Vector{Dict{Symbol, T}}`
|
||
A vector of dictionaries where each dictionary contains the keys `:name` (the name of the message owner) and `:text` (the text of the message).
|
||
- `formatname::T`
|
||
The name of the format to be used for converting the chat messages.
|
||
# Return
|
||
- `formattedtext::String`
|
||
text formatted to model format
|
||
|
||
# Example
|
||
```jldoctest
|
||
julia> using Revise
|
||
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 = 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"
|
||
```
|
||
"""
|
||
function formatLLMtext(messages::Vector{Dict{Symbol, T}}, formatname::String
|
||
)::String where {T<:AbstractString}
|
||
f =
|
||
if formatname == "llama3instruct"
|
||
formatLLMtext_llama3instruct
|
||
elseif formatname == "mistral"
|
||
# not define yet
|
||
elseif formatname == "phi3instruct"
|
||
# not define yet
|
||
elseif formatname == "qwen"
|
||
formatLLMtext_qwen
|
||
elseif formatname == "qwen3"
|
||
formatLLMtext_qwen3
|
||
elseif formatname == "phi4"
|
||
formatLLMtext_phi4
|
||
elseif formatname == "granite3"
|
||
formatLLMtext_granite3
|
||
else
|
||
error("$formatname template not define yet")
|
||
end
|
||
|
||
str = ""
|
||
for (i, t) in enumerate(messages)
|
||
if i < length(messages)
|
||
str *= f(t[:name], t[:text])
|
||
else
|
||
str *= f(t[:name], t[:text]; assistantStarter=true)
|
||
end
|
||
end
|
||
|
||
return str
|
||
end
|
||
|
||
|
||
function extractthink(text::String)
|
||
think = nothing
|
||
str = nothing
|
||
if occursin("<think>", text)
|
||
r = GeneralUtils.extractTextBetweenString(text, "<think>", "</think>")
|
||
if r[:success]
|
||
think = r[:text]
|
||
end
|
||
str = string(split(text, "</think>")[2])
|
||
else
|
||
str = text
|
||
end
|
||
return think, str
|
||
end
|
||
|
||
|
||
|
||
""" Validate that an agent's JSON response contains all required keys and no extra keys.
|
||
|
||
The function checks if `responsedict` contains exactly the keys specified in `requiredKeys`
|
||
with no duplicates and no missing keys. It is designed to validate structured agent responses
|
||
against an expected schema.
|
||
|
||
# Arguments
|
||
- `responsedict::Dict`
|
||
A dictionary containing the agent's JSON response. Must be a plain `Dict` or
|
||
dictionary-like object with String keys.
|
||
- `requiredKeys::T` where `T<:Array{String}`
|
||
An array of required key names that must be present in `responsedict`.
|
||
|
||
# Return
|
||
- `Tuple{Bool, Union{String, Nothing}}`
|
||
A tuple where the first element indicates whether validation passed (`true`) or failed (`false`),
|
||
and the second element contains an error message if validation failed, or `nothing` if it passed.
|
||
|
||
# Details
|
||
The validation logic checks:
|
||
1. **Duplicate keys**: If `responsedict` contains more keys than `requiredKeys`, validation fails
|
||
because the agent included extra/unexpected keys.
|
||
2. **Missing keys**: If any key in `requiredKeys` is absent from `responsedict`, validation fails
|
||
and the specific missing keys are listed in the error message.
|
||
3. **Valid response**: If all required keys are present and no extra keys exist, validation passes.
|
||
|
||
# Example
|
||
|
||
```julia
|
||
julia> using GeneralUtils
|
||
julia> requiredKeys = ["wine_name", "price", "rating"]
|
||
julia> response = Dict("wine_name"=>"Château Margaux", "price"=>250.0, "rating"=>98)
|
||
julia> GeneralUtils.checkAgentResponse_JSON(response, requiredKeys)
|
||
(true, nothing)
|
||
|
||
julia> response_missing = Dict("wine_name"=>"Château Margaux", "price"=>250.0)
|
||
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> GeneralUtils.checkAgentResponse_JSON(response_extra, requiredKeys)
|
||
(false, "Your previous attempt has duplicated points according to the required response format")
|
||
```
|
||
"""
|
||
function checkAgentResponse_JSON(responsedict::T1, requiredKeys::T2
|
||
)::Tuple where {T1<:AbstractDict, T2<:Array{String}}
|
||
_responsedictKey = keys(responsedict)
|
||
responsedictKey = [i for i in _responsedictKey] # convert into a list
|
||
is_requiredKeys_in_responsedictKey = [i ∈ responsedictKey for i in requiredKeys]
|
||
ispass = false
|
||
errormsg = nothing
|
||
if length(is_requiredKeys_in_responsedictKey) > length(requiredKeys)
|
||
errormsg = "Your previous attempt has duplicated points according to the required response format"
|
||
ispass = false
|
||
elseif !all(is_requiredKeys_in_responsedictKey)
|
||
zeroind = findall(x -> x == 0, is_requiredKeys_in_responsedictKey)
|
||
missingkeys = [requiredKeys[i] for i in zeroind]
|
||
errormsg = "$missingkeys are missing from your previous response"
|
||
ispass = false
|
||
else
|
||
ispass = true
|
||
end
|
||
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)
|
||
and wraps them in proper JSON braces to create a valid JSON string. It cleans the input
|
||
by removing common formatting artifacts like braces, code block markers, and language
|
||
specifiers before wrapping the content.
|
||
|
||
# Arguments
|
||
- `text::String`
|
||
A string containing key-value pairs, typically in the format `key: value, key: value`.
|
||
May contain leading/trailing braces, code block markers (```), or language specifiers
|
||
(e.g., `json`) that will be removed.
|
||
|
||
# Return
|
||
- `String`
|
||
A JSON-formatted string with the key-value pairs wrapped in `{}` braces.
|
||
|
||
# Notes
|
||
- The function removes `{`, `}`, `````, and `json` from the input before wrapping.
|
||
- The output is always wrapped in curly braces to create valid JSON structure.
|
||
- This is typically used to sanitize LLM responses that contain key-value data
|
||
but may include formatting artifacts.
|
||
|
||
# Examples
|
||
```jldoctest
|
||
julia> text = "thought: Hello, action: CHATBOX, input: How can I help?"
|
||
julia> clean_json_response(text)
|
||
"{thought: Hello, action: CHATBOX, input: How can I help?}"
|
||
|
||
julia> text = "{thought: Hello, action: CHATBOX}"
|
||
julia> clean_json_response(text)
|
||
"{thought: Hello, action: CHATBOX}"
|
||
|
||
julia> text = "```json{thought: Hello, action: CHATBOX}```"
|
||
julia> clean_json_response(text)
|
||
"{thought: Hello, action: CHATBOX}"
|
||
```
|
||
"""
|
||
function clean_json_response(text::String)
|
||
removelist = ["{", "}", "```", "json", "\n"]
|
||
return '{' * removestring(text, removelist) * '}'
|
||
end
|
||
|
||
|
||
|
||
""" Harvest database schema as undirected graph from PostgreSQL.
|
||
|
||
Extracts table structure and relationships from a PostgreSQL database by querying
|
||
system catalogs to build a graph representation of tables and their relationships.
|
||
|
||
# Arguments
|
||
- `pg_conn_str::String`
|
||
PostgreSQL connection string in LibPQ format (e.g., "host=hostname port=5432 dbname=database user=username password=secret")
|
||
|
||
# Return
|
||
- `g::SimpleGraph`: An undirected graph where nodes represent tables and edges represent
|
||
relationships (foreign keys or shared ID column patterns).
|
||
- `id_to_table::Dict{Int, String}`: Maps numerical node IDs (1..n) to actual table names.
|
||
- `table_to_id::Dict{String, Int}`: Reverse mapping from table names to graph node IDs.
|
||
|
||
# Details
|
||
The function extracts schema information using two strategies:
|
||
1. **Explicit Foreign Keys**: Queries `pg_constraint` for actual foreign key relationships
|
||
2. **Fallback Name Matching**: Infers relationships from shared column naming patterns
|
||
(e.g., `seller_id`, `product_id` columns across tables)
|
||
|
||
# Example
|
||
```julia
|
||
julia> using GeneralUtils
|
||
julia> pg_conn_str = "host=localhost port=5432 dbname=winedb user=admin password=secret"
|
||
julia> g, id_to_table, table_to_id = GeneralUtils.harvest_db_undirected_schema_graph(pg_conn_str)
|
||
julia> vertices(g)
|
||
10
|
||
```
|
||
|
||
# Integration Guide: Finding Related Tables for User Questions
|
||
This function is designed to work together with `extract_column_metadata` and `resolve_semantic_cluster` to answer user questions by identifying related tables:
|
||
|
||
```julia
|
||
# Step 1: Extract column metadata and generate embeddings for semantic search
|
||
pg_conn_str = "host=localhost port=5432 dbname=winedb user=admin password=secret"
|
||
metadata_df = GeneralUtils.extract_column_metadata(pg_conn_str)
|
||
embedding_ready = GeneralUtils.generate_embedding_payloads(metadata_df)
|
||
|
||
# Use only text content for embedding
|
||
embedding_ready_2 = [i["text_content"] for i in embedding_ready]
|
||
table_embedding = get_embedding(embedding_ready_2)
|
||
|
||
# Get embedding for user question
|
||
_user_question_embedding = get_embedding([question])
|
||
user_question_embedding = Float64.(_user_question_embedding["data"][1]["embedding"])
|
||
|
||
# Calculate similarity between question and all columns
|
||
user_question_similarity = []
|
||
for i in table_embedding["data"]
|
||
i_data = i["embedding"]
|
||
i_float = Float64.(i_data)
|
||
r = 1 - Distances.cosine_dist(i_float, user_question_embedding)
|
||
push!(user_question_similarity, r)
|
||
end
|
||
|
||
# Step 2: Find top related tables
|
||
new_df = hcat(metadata_df, DataFrame(user_question_similarity = user_question_similarity))
|
||
sorted_df = sort(new_df, :user_question_similarity, rev=true)
|
||
_top_20_tables = unique(sorted_df[1:20, :table_name])
|
||
top_20_tables = [i for i in _top_20_tables]
|
||
|
||
# Step 3: Build schema graph and resolve table relationships
|
||
g, id_to_table, table_to_id = GeneralUtils.harvest_db_undirected_schema_graph(pg_conn_str)
|
||
table_relationship = GeneralUtils.resolve_semantic_cluster(top_20_tables, g, table_to_id, id_to_table)
|
||
|
||
# table_relationship now contains tables in the order they should be joined
|
||
```
|
||
"""
|
||
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 eachindex(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,
|
||
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 in LibPQ format (e.g., "host=hostname port=5432 dbname=database user=username password=secret")
|
||
|
||
# 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 = "host=localhost port=5432 dbname=winedb user=admin password=secret"
|
||
julia> df = GeneralUtils.extract_column_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
|
||
```
|
||
|
||
# Integration Guide
|
||
This function is the first step in the semantic table discovery pipeline. Use it with `generate_embedding_payloads` to prepare data for vector similarity search:
|
||
|
||
```julia
|
||
# Step 1: Extract metadata
|
||
metadata_df = GeneralUtils.extract_column_metadata(pg_conn_str)
|
||
|
||
# Step 2: Generate embedding payloads for semantic search
|
||
embedding_ready = GeneralUtils.generate_embedding_payloads(metadata_df)
|
||
|
||
# Use text_content field for embedding
|
||
text_content_list = [i["text_content"] for i in embedding_ready]
|
||
embeddings = get_embedding(text_content_list)
|
||
```
|
||
"""
|
||
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 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 = LibPQ.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_column_metadata` into a vector of
|
||
dictionaries structured for vector embedding storage and retrieval.
|
||
|
||
# Arguments
|
||
- `df::DataFrame`
|
||
A DataFrame with columns from `extract_column_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 = "host=localhost port=5432 dbname=winedb user=admin password=secret"
|
||
julia> df = GeneralUtils.extract_column_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"))
|
||
```
|
||
|
||
# Integration Guide
|
||
This function transforms column metadata into a format suitable for semantic search.
|
||
Use `text_content` field to generate embeddings, then compare against user question embeddings:
|
||
|
||
```julia
|
||
# 1. Extract metadata
|
||
metadata_df = GeneralUtils.extract_column_metadata(pg_conn_str)
|
||
|
||
# 2. Generate payloads
|
||
embedding_ready = GeneralUtils.generate_embedding_payloads(metadata_df)
|
||
|
||
# 3. Extract text content for embedding (only this field is used for similarity)
|
||
text_content_list = [i["text_content"] for i in embedding_ready]
|
||
|
||
# 4. Generate embeddings for all columns
|
||
column_embeddings = get_embedding(text_content_list)
|
||
|
||
# 5. Generate embedding for user question
|
||
question_embedding = get_embedding([user_question])
|
||
|
||
# 6. Calculate cosine similarity to find most relevant columns/tables
|
||
# (See harvest_db_undirected_schema_graph for full integration example)
|
||
```
|
||
"""
|
||
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"]
|
||
```
|
||
|
||
# Integration Guide: Understanding Table Relationships
|
||
This function determines how tables from vector search results are connected in the database schema.
|
||
The returned list represents the optimal join order for constructing SQL queries.
|
||
|
||
**How tables are linked:**
|
||
- **Explicit Foreign Keys**: Direct relationships defined by `FOREIGN KEY` constraints in PostgreSQL
|
||
- **Implicit Relationships**: Tables sharing similar column naming patterns (e.g., `user_id` in both `users` and `orders` tables)
|
||
|
||
**Example output interpretation:**
|
||
```julia
|
||
# Input: Top 20 tables from semantic search
|
||
top_tables = ["users", "products", "payments"]
|
||
|
||
# Output: Tables in join order
|
||
table_relationship = ["users", "orders", "payments"]
|
||
# Interpretation:
|
||
# 1. Start with 'users' table
|
||
# 2. Join 'orders' via foreign key (likely users.id -> orders.user_id)
|
||
# 3. Join 'payments' via foreign key (likely orders.id -> payments.order_id)
|
||
```
|
||
|
||
**Full workflow example:**
|
||
```julia
|
||
# Step 1: Get column embeddings and find semantically related tables
|
||
metadata_df = GeneralUtils.extract_column_metadata(pg_conn_str)
|
||
embedding_ready = GeneralUtils.generate_embedding_payloads(metadata_df)
|
||
text_content = [i["text_content"] for i in embedding_ready]
|
||
table_embedding = get_embedding(text_content)
|
||
|
||
# Step 2: Calculate similarity with user question
|
||
question_embedding = Float64.(get_embedding([question])["data"][1]["embedding"])
|
||
similarities = [1 - Distances.cosine_dist(Float64.(i["embedding"]), question_embedding)
|
||
for i in table_embedding["data"]]
|
||
|
||
# Step 3: Extract top related tables
|
||
top_tables = unique(metadata_df[sortperm(similarities)[1:20], :table_name])
|
||
|
||
# Step 4: Build schema graph and resolve relationships
|
||
g, id_to_table, table_to_id = GeneralUtils.harvest_db_undirected_schema_graph(pg_conn_str)
|
||
related_tables = GeneralUtils.resolve_semantic_cluster(top_tables, g, table_to_id, id_to_table)
|
||
|
||
# related_tables now contains tables in optimal join order for SQL query construction
|
||
```
|
||
"""
|
||
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
|
||
|
||
|
||
""" Harvest entity catalog from database column.
|
||
|
||
Extracts unique, non-null values from a specific column to build a local index for
|
||
semantic search or entity resolution.
|
||
|
||
# Arguments
|
||
- `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
|
||
- `column::String`
|
||
Column name containing entity values
|
||
|
||
# Return
|
||
- `Vector{String}`
|
||
A vector of unique, stripped strings from the specified column. Empty strings
|
||
are removed via `strip()`.
|
||
|
||
# Details
|
||
The function:
|
||
1. Connects to PostgreSQL database
|
||
2. Executes `SELECT DISTINCT column FROM table WHERE column IS NOT NULL`
|
||
3. Converts result to DataFrame
|
||
4. Strips whitespace from each value and converts to String
|
||
5. Returns clean vector of unique entity values
|
||
|
||
# Example
|
||
```julia
|
||
julia> using GeneralUtils
|
||
julia> conn = "host=localhost port=5432 dbname=winedb user=admin password=secret"
|
||
julia> fruits = GeneralUtils.harvest_entity_catalog(conn, "products", "fruit_name")
|
||
["Apple", "Banana", "Orange", "Mango"]
|
||
```
|
||
|
||
# Integration Guide: Finding Closest Entity Match in Database
|
||
When users type queries, they often make typos or use abbreviations. This function finds
|
||
the closest matching entity from a database column, handling common input errors.
|
||
|
||
**Example:**
|
||
```julia
|
||
# Database contains: ["Hand Old Bar & Grill", "Hand Old", "Wine Cellar"]
|
||
# User types: "HandOld" (compressed words, missing space)
|
||
|
||
catalog = GeneralUtils.harvest_entity_catalog(pg_conn_str, "venues", "name")
|
||
closest = GeneralUtils.resolve_entity("HandOld", catalog; threshold=0.9)
|
||
# Returns: "Hand Old Bar & Grill" (closest match)
|
||
```
|
||
|
||
**Step 1: Extract all valid values from a column**
|
||
```julia
|
||
# Get all wine names from the database
|
||
wine_names = GeneralUtils.harvest_entity_catalog(pg_conn_str, "wine", "wine_name")
|
||
# Returns: ["Château Margaux", "Château Lafite", "Domaine Leroy", ...]
|
||
```
|
||
|
||
**Step 2: Find closest match for user input**
|
||
```julia
|
||
# User types a wine name with a typo
|
||
catalog = GeneralUtils.harvest_entity_catalog(pg_conn_str, "wine", "wine_name")
|
||
closest_wine = GeneralUtils.resolve_entity("Chateu Margaux", catalog; threshold=0.9)
|
||
# Returns: "Château Margaux" (closest match)
|
||
```
|
||
|
||
**Step 3: Resolve all fields in a response**
|
||
```julia
|
||
# Agent produces a response with potentially misspelled values
|
||
responsedict = Dict("wine_name" => "Chateu Margaux", "region" => "Bord")
|
||
|
||
# For each field, find the closest match in the database
|
||
for (k, v) in responsedict
|
||
catalog = GeneralUtils.harvest_entity_catalog(pg_conn_str, "wine", k)
|
||
resolved = GeneralUtils.resolve_entity(v, catalog; threshold=0.9)
|
||
responsedict[k] = resolved # Replace messy input with database entity
|
||
end
|
||
```
|
||
|
||
**Threshold guidance:**
|
||
- `threshold=0.5` - Lenient, may return false positives
|
||
- `threshold=0.7` - Moderate balance
|
||
- `threshold=0.9` - Strict, only high-confidence matches
|
||
- Returns `""` when no match meets threshold
|
||
"""
|
||
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::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;"
|
||
|
||
try
|
||
df = DataFrame(LibPQ.execute(conn, query))
|
||
# Return as a clean array of strings
|
||
return String.(strip.(df[:, 1]))
|
||
catch
|
||
return String[]
|
||
finally
|
||
close(conn)
|
||
end
|
||
end
|
||
|
||
|
||
""" Resolve entity name from messy input using fuzzy string matching.
|
||
|
||
Matches user-provided text against a reference catalog using Jaro-Winkler similarity
|
||
and returns the closest matching exact string from the database catalog.
|
||
|
||
# Arguments
|
||
- `messy_input::String`
|
||
The user input text that may contain typos, compressed words, or variations.
|
||
- `catalog::Vector{String}`
|
||
A vector of valid, exact entity strings from the database.
|
||
|
||
# Keyword Arguments
|
||
- `threshold::Float64` (default: `0.5`)
|
||
Minimum similarity score (0.0 to 1.0) required to return a match. Lower values
|
||
allow more lenient matching; higher values require closer matches.
|
||
|
||
# Return
|
||
- `String`
|
||
The exact matching string from `catalog` if similarity score ≥ threshold,
|
||
otherwise an empty string `""`.
|
||
|
||
# Details
|
||
The function:
|
||
1. Normalizes input to lowercase and strips whitespace
|
||
2. Computes Jaro-Winkler similarity score against each catalog entry
|
||
3. Applies substring fallback: if compressed words match (e.g., "HandOld" → "Hand Old Bar & Grill"),
|
||
boosts score to 0.85
|
||
4. Returns the highest-scoring catalog entry if score ≥ threshold, else empty string
|
||
|
||
# Example
|
||
```julia
|
||
julia> using GeneralUtils
|
||
julia> catalog = ["Hand Old Bar & Grill", "Hand Old", "Wine Cellar"]
|
||
julia> GeneralUtils.resolve_entity("HandOld", catalog; threshold=0.5)
|
||
"Hand Old Bar & Grill"
|
||
|
||
julia> GeneralUtils.resolve_entity("Wine Cellar", catalog; threshold=0.5)
|
||
"Wine Cellar"
|
||
|
||
julia> GeneralUtils.resolve_entity("Unknown Place", catalog; threshold=0.5)
|
||
""
|
||
```
|
||
|
||
# Integration Guide: Finding Closest Entity Match in Database
|
||
When users type queries, they often make typos or use abbreviations. This function finds
|
||
the closest matching entity from a database column, handling common input errors.
|
||
|
||
**Example:**
|
||
```julia
|
||
# Database contains: ["Hand Old Bar & Grill", "Hand Old", "Wine Cellar"]
|
||
# User types: "HandOld" (compressed words, missing space)
|
||
|
||
catalog = GeneralUtils.harvest_entity_catalog(pg_conn_str, "venues", "name")
|
||
closest = GeneralUtils.resolve_entity("HandOld", catalog; threshold=0.9)
|
||
# Returns: "Hand Old Bar & Grill" (closest match)
|
||
```
|
||
|
||
**Step 1: Extract all valid values from a column**
|
||
```julia
|
||
# Get all wine names from the database
|
||
wine_names = GeneralUtils.harvest_entity_catalog(pg_conn_str, "wine", "wine_name")
|
||
# Returns: ["Château Margaux", "Château Lafite", "Domaine Leroy", ...]
|
||
```
|
||
|
||
**Step 2: Find closest match for user input**
|
||
```julia
|
||
# User types a wine name with a typo
|
||
catalog = GeneralUtils.harvest_entity_catalog(pg_conn_str, "wine", "wine_name")
|
||
closest_wine = GeneralUtils.resolve_entity("Chateu Margaux", catalog; threshold=0.9)
|
||
# Returns: "Château Margaux" (closest match)
|
||
```
|
||
|
||
**Step 3: Resolve all fields in a response**
|
||
```julia
|
||
# Agent produces a response with potentially misspelled values
|
||
responsedict = Dict("wine_name" => "Chateu Margaux", "region" => "Bord")
|
||
|
||
# For each field, find the closest match in the database
|
||
for (k, v) in responsedict
|
||
catalog = GeneralUtils.harvest_entity_catalog(pg_conn_str, "wine", k)
|
||
resolved = GeneralUtils.resolve_entity(v, catalog; threshold=0.9)
|
||
responsedict[k] = resolved # Replace messy input with database entity
|
||
end
|
||
```
|
||
|
||
**Threshold guidance:**
|
||
- `threshold=0.5` - Lenient, may return false positives
|
||
- `threshold=0.7` - Moderate balance
|
||
- `threshold=0.9` - Strict, only high-confidence matches
|
||
- Returns `""` when no match meets threshold
|
||
"""
|
||
function resolve_entity(messy_input::String, catalog::Vector{String}; threshold=0.5)::String
|
||
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
|
||
|
||
|
||
function get_db_table_schema(pg_conn_str::String, table_name::String)::DataFrame
|
||
conn = LibPQ.Connection(pg_conn_str)
|
||
return get_db_table_schema(conn, table_name)
|
||
end
|
||
|
||
function get_db_table_schema(conn::LibPQ.Connection, table_name::String)::DataFrame
|
||
# This direct SQL query pulls the column specifications along with column-level descriptions
|
||
query = """
|
||
SELECT
|
||
a.attname AS column_name,
|
||
format_type(a.atttypid, a.atttypmod) AS data_type,
|
||
COALESCE(d.description, '') AS column_comment,
|
||
CASE
|
||
WHEN p.contype = 'p' THEN 'PRIMARY KEY'
|
||
WHEN p.contype = 'u' THEN 'UNIQUE'
|
||
WHEN p.contype = 'f' THEN 'FOREIGN KEY'
|
||
WHEN p.contype = 'c' THEN 'CHECK'
|
||
ELSE ''
|
||
END AS constraint_type,
|
||
COALESCE(p.conname, '') AS constraint_name
|
||
FROM
|
||
pg_catalog.pg_attribute a
|
||
JOIN
|
||
pg_catalog.pg_class c ON a.attrelid = c.oid
|
||
JOIN
|
||
pg_catalog.pg_namespace n ON c.relnamespace = n.oid
|
||
LEFT JOIN
|
||
pg_catalog.pg_description d ON d.objoid = c.oid AND d.objsubid = a.attnum
|
||
LEFT JOIN
|
||
pg_catalog.pg_constraint p ON p.conrelid = c.oid AND a.attnum = ANY(p.conkey)
|
||
WHERE
|
||
c.relname = '$table_name' -- <-- Put your table name here
|
||
AND n.nspname = 'public' -- <-- Change schema if not 'public'
|
||
AND a.attnum > 0
|
||
AND NOT a.attisdropped
|
||
ORDER BY
|
||
a.attnum;
|
||
"""
|
||
|
||
try
|
||
# Execute and format into a clean DataFrame
|
||
result = LibPQ.execute(conn, query)
|
||
df = DataFrame(result)
|
||
|
||
return df
|
||
finally
|
||
close(conn)
|
||
end
|
||
end
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
end # module llmUtil |