use rag_toolchain::chunkers::{Chunker, TokenChunker};
use rag_toolchain::clients::{AsyncEmbeddingClient, OpenAIEmbeddingClient};
use rag_toolchain::common::{Chunks, Embedding, OpenAIEmbeddingModel};
use rag_toolchain::stores::{EmbeddingStore, PostgresVectorStore};
#[tokio::main]
async fn main() {
const EMBEDDING_MODEL: OpenAIEmbeddingModel = OpenAIEmbeddingModel::TextEmbeddingAda002;
let text = std::fs::read_to_string("examples/pg_vector/example_text.txt").unwrap();
let chunker = TokenChunker::try_new(
std::num::NonZeroUsize::new(50).unwrap(),
25,
EMBEDDING_MODEL,
)
.unwrap();
let chunks: Chunks = chunker.generate_chunks(&text).unwrap();
let store: PostgresVectorStore = PostgresVectorStore::try_new("embeddings", EMBEDDING_MODEL)
.await
.unwrap();
let client: OpenAIEmbeddingClient = OpenAIEmbeddingClient::try_new(EMBEDDING_MODEL).unwrap();
let embeddings: Vec<Embedding> = client.generate_embeddings(chunks).await.unwrap();
store.store_batch(embeddings).await.unwrap();
}