using_trieve_with_nu.nu

let dataset_id = $env.DATASET_ID
let api_key = $env.API_KEY
let clickhousePassword = $env.CLICKHOUSE_PASSWORD
let headers = {
    "TR-Dataset": $dataset_id
    "Authorization": $api_key
    "X-API-Version": "V2"
}

def clickhouse-query [query: string] {
    clickhouse client --host localhost --port 9001 --password $clickhousePassword --user default -q ($query | append "FORMAT CSVWithNames" | str join " ") | from csv 
}

def httppost [url: string, body: any] {
    http post --content-type application/json $url $body --headers $headers -e
}

###

let count = 1;
cat rag_queries.csv | from csv  | each {
    let body = {
        chunk_html: $in.user_message
        tracking_id: $in.id
        metadata: {
            hallucinations: $in.detected_hallucinations
        }
        upsert_by_tracking_id: true
    };
    httppost "https://api.trieve.ai/api/chunk" $body
}

let body = {
    query: "pool"
    search_type: "fulltext"
    page: 30
    page_size: 20
}
httppost $"https://api.trieve.ai/api/chunk/search" $body | get chunks | each {
    {
        id: $in.chunk.tracking_id
        query: $in.chunk.chunk_html
    }
} | to csv | save --append filtered.csv

cat filtered.csv | from csv | each {
    let llm_response = clickhouse-query $"select llm_response from rag_queries where id = '($in.id)'" | get llm_response | get 0;
    print $llm_response
    {
       id: $in.id
       query: $in.query
       ai_answer: $llm_response
    }
} | to tsv | save --append full.tsv