Using new openai crate to generate embeddings

tokenizer/Cargo.toml

@@ -11,4 +11,4 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 mongodb = "2.0"
 ingestion = { path = "../ingestion" }
-openai = "1.0.0-alpha.13"
+openai-rust = "0.5.1"

tokenizer/src/main.rs

@@ -6,8 +6,8 @@ async fn main() -> std::io::Result<()> {
     mongo.connect().await?;
 
     // set collection to tokenize
-    let collection = "github_data"
-    let oai_key = "generate_2023"
+    let collection = "github_data";
+    let oai_key = "generate_2023";
 
     // TODO: create a new collection for each repo, insert documents into sub collections 
     let tokenizer: OpenAIClient = OpenAIClient::new(oai_key, mongo, collection);

tokenizer/src/oai.rs

@@ -1,12 +1,12 @@
 use ingestion::mongo_utils::MongoDriver;
 use mongodb::bson::{doc, Document};
-use openai::{Language, OpenaiClient};
+use openai_rust::Client;
 use std::error::Error;
 use std::env;
 
 
 pub struct OpenAIClient {
-    oai_client: OpenaiClient,
+    oai_client: Client,
     mongo_model: &MongoDriver,
     model: Language,
 }
@@ -15,7 +15,7 @@ impl OpenAIClient {
 
     pub fn new(openai_api_key: &str, mongo_model: &MongoDriver) -> Self {
         OpenAIClient {
-            oai_client: OpenaiClient::new(openai_api_key),
+            oai_client: Client::new(openai_api_key),
             mongo_model: mongo_model,
             model: Language::English,
         }
@@ -32,10 +32,11 @@ impl OpenAIClient {
             match result {
                 Ok(document) => {
                     if let Some(text) = document.get_str("text") {
-
-                        let tokens = openai.tokenize(&language_model, text)?;
+                        let args: openai_rust::embeddings::EmbeddingsArguments = openai_rust::embeddings::EmbeddingsArguments::new("text-embedding-ada-002", text.to_owned());
+                        let embedding: Vec<openai_rust::embeddings::EmbeddingsData> = self.oai_client.create_embeddings(args).await.unwrap().data;
+
                         let update_doc = doc! {
-                            "$set": { "tokens": tokens }
+                            "$set": { "embedding": embedding }
                         };
 
                         col.update_one(document, update_doc, None).await?;