a-agmon

use async_trait::async_trait;
use chrono::Local;
use graph_flow::{
    Context, ExecutionStatus, FlowRunner, GraphBuilder, GraphStorage, InMemoryGraphStorage,
    InMemorySessionStorage, NextAction, Session, SessionStorage, Task, TaskResult,
};
use image::{ImageBuffer, Rgb};
use minifb::{Key, Window, WindowOptions};
use nokhwa::{
    pixel_format::RgbFormat,

a-agmon

from utils.database import get_db_connection, TABLE_NAME
from llama_index.embeddings.huggingface import HuggingFaceEmbedding


def zscore_search(query_text, z_threshold=2.5, top_k=1000):
    """
    Perform z-score calibrated vector search.
    
    Args:
        query_text: Text to search for

a-agmon

a-agmon

use std::{sync::Arc, time::Instant};

use datafusion::{
    arrow::{array::RecordBatch, json::ArrayWriter},
    prelude::SessionContext,
};
use iceberg::{io::FileIO, table::StaticTable, TableIdent};
use iceberg_datafusion::IcebergTableProvider;
use lambda_runtime::{
    run, service_fn,

a-agmon

async fn init_db_writer_task(
    reciever: Receiver<EmbeddingEntry>,
    db_uri: &str,
    table_name: &str,
    buffer_size: usize,
) -> anyhow::Result<JoinHandle<()>> {
    let db = storage::VecDB::connect(db_uri, table_name).await?;
    let task_handle = tokio::spawn(async move {
        let mut embeddings_buffer = Vec::new();

a-agmon

let lance_table = connection.open_table("table_name").execute().await?;
let batches = vec![Ok(RecordBatch::try_new(
    schema.clone(),
    vec![Arc::new(key_array), Arc::new(vectors_array)],
)?)
.map_err(|e| ArrowError::from_external_error(e.into()))];
let batch_iterator = RecordBatchIterator::new(batches, schema);
let boxed_batches = Box::new(batch_iterator);
lance_table.add(boxed_batches).execute().await?;

a-agmon

  Arc::new(Schema::new(vec![
      Field::new("filename", DataType::Utf8, false),
      Field::new(
          "vector",
          DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 384),
          true,
      ),
  ]))

a-agmon

const EMB_MODEL_ID: &str = "sentence-transformers/all-MiniLM-L6-v2";
const EMB_MODEL_REV: &str = "refs/pr/21";

thread_local! {
    static BERT_MODEL: Rc<BertModelWrapper> =   {
        info!("Loading a model on thread: {:?}", std::thread::current().id());
        let model = BertModelWrapper::new(candle::Device::Cpu, EMB_MODEL_ID, EMB_MODEL_REV);
        match model {
            Ok(model) => Rc::new(model),
            Err(e) => {

a-agmon

fn process_text_file(sender: Sender<EmbeddingEntry>, filename: &str) -> anyhow::Result<()> {
    let bert_model = models::bert::get_model_reference(EMB_MODEL_ID, EMB_MODEL_REV)?;
    let text_chunks: Vec<&str> = read_file_in_chunks(filename, 256)?;
    let file_vector = embed_multiple_sentences(&text_chunks, false, &bert_model)?;
    sender.send(EmbeddingEntry {
        filename: filename.to_string(),
        embedding: file_vector[0].clone(),
    })?;
    Ok(())
}

a-agmon

async fn main() -> anyhow::Result<()> {
    // init the sync channel 
    let (sender, reciever) = std::sync::mpsc::channel::<EmbeddingEntry>();
    // start the write task loop and get a handle to it
    let db_writer_task = init_writer_task(reciever,db_uri.as_str());

    // list the files in the directory to be embedded
    let files_dir = fs::read_dir(cli_args.input_directory)?;
    files_dir
        .into_iter()