portkey integration

.github/workflows/extension_ci.yml

@@ -111,6 +111,8 @@ jobs:
         env:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           CO_API_KEY: ${{ secrets.CO_API_KEY }}
+          PORTKEY_API_KEY: ${{ secrets.PORTKEY_API_KEY }}
+          PORTKEY_VIRTUAL_KEY_OPENAI: ${{ secrets.PORTKEY_VIRTUAL_KEY_OPENAI }}
         run: |
           cd ../core && cargo test
       - name: Restore cached binaries
@@ -135,6 +137,8 @@ jobs:
         env:
           HF_API_KEY: ${{ secrets.HF_API_KEY }}
           CO_API_KEY: ${{ secrets.CO_API_KEY }}
+          PORTKEY_API_KEY: ${{ secrets.PORTKEY_API_KEY }}
+          PORTKEY_VIRTUAL_KEY_OPENAI: ${{ secrets.PORTKEY_VIRTUAL_KEY_OPENAI }}
         run: |
           echo "\q" | make run
           make test-integration

core/src/transformers/mod.rs

@@ -1,5 +1,4 @@
 pub mod generic;
 pub mod http_handler;
-pub mod ollama;
 pub mod providers;
 pub mod types;

core/src/transformers/providers/mod.rs

@@ -1,6 +1,7 @@
 pub mod cohere;
 pub mod ollama;
 pub mod openai;
+pub mod portkey;
 pub mod vector_serve;
 
 use anyhow::Result;
@@ -51,6 +52,7 @@ pub fn get_provider(
     model_source: &ModelSource,
     api_key: Option<String>,
     url: Option<String>,
+    virtual_key: Option<String>,
 ) -> Result<Box<dyn EmbeddingProvider>, VectorizeError> {
     match model_source {
         ModelSource::OpenAI => Ok(Box::new(providers::openai::OpenAIProvider::new(
@@ -59,11 +61,42 @@ pub fn get_provider(
         ModelSource::Cohere => Ok(Box::new(providers::cohere::CohereProvider::new(
             url, api_key,
         ))),
+        ModelSource::Portkey => Ok(Box::new(providers::portkey::PortkeyProvider::new(
+            url,
+            api_key,
+            virtual_key,
+        ))),
         ModelSource::SentenceTransformers => Ok(Box::new(
             providers::vector_serve::VectorServeProvider::new(url, api_key),
         )),
-        ModelSource::Ollama | ModelSource::Tembo => Err(anyhow::anyhow!(
+        ModelSource::Ollama => Ok(Box::new(providers::ollama::OllamaProvider::new(url))),
+        ModelSource::Tembo => Err(anyhow::anyhow!(
             "Ollama/Tembo transformer not implemented yet"
         ))?,
     }
 }
+
+fn split_vector(vec: Vec<String>, chunk_size: usize) -> Vec<Vec<String>> {
+    vec.chunks(chunk_size).map(|chunk| chunk.to_vec()).collect()
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct ChatMessageRequest {
+    pub role: String,
+    pub content: String,
+}
+
+#[derive(Deserialize, Debug)]
+struct ChatResponse {
+    choices: Vec<Choice>,
+}
+
+#[derive(Deserialize, Debug)]
+struct Choice {
+    message: ResponseMessage,
+}
+
+#[derive(Deserialize, Debug)]
+struct ResponseMessage {
+    content: String,
+}

core/src/transformers/providers/ollama.rs

@@ -1,4 +1,6 @@
-use super::{EmbeddingProvider, GenericEmbeddingRequest, GenericEmbeddingResponse};
+use super::{
+    ChatMessageRequest, EmbeddingProvider, GenericEmbeddingRequest, GenericEmbeddingResponse,
+};
 use crate::errors::VectorizeError;
 use async_trait::async_trait;
 use ollama_rs::{generation::completion::request::GenerationRequest, Ollama};
@@ -8,20 +10,19 @@ use url::Url;
 pub const OLLAMA_BASE_URL: &str = "http://localhost:3001";
 
 pub struct OllamaProvider {
-    pub model_name: String,
     pub instance: Ollama,
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
 struct ModelInfo {
-    model: String,
     embedding_dimension: u32,
     max_seq_len: u32,
 }
 
 impl OllamaProvider {
-    pub fn new(model_name: String, url: String) -> Self {
-        let parsed_url = Url::parse(&url).unwrap_or_else(|_| panic!("invalid url: {}", url));
+    pub fn new(url: Option<String>) -> Self {
+        let url_in = url.unwrap_or_else(|| OLLAMA_BASE_URL.to_string());
+        let parsed_url = Url::parse(&url_in).unwrap_or_else(|_| panic!("invalid url: {}", url_in));
         let instance = Ollama::new(
             format!(
                 "{}://{}",
@@ -30,10 +31,7 @@ impl OllamaProvider {
             ),
             parsed_url.port().expect("parsed port missing"),
         );
-        OllamaProvider {
-            model_name,
-            instance,
-        }
+        OllamaProvider { instance }
     }
 }
 
@@ -44,10 +42,11 @@ impl EmbeddingProvider for OllamaProvider {
         request: &'a GenericEmbeddingRequest,
     ) -> Result<GenericEmbeddingResponse, VectorizeError> {
         let mut all_embeddings: Vec<Vec<f64>> = Vec::with_capacity(request.input.len());
+        let model_name = request.model.clone();
         for ipt in request.input.iter() {
             let embed = self
                 .instance
-                .generate_embeddings(self.model_name.clone(), ipt.clone(), None)
+                .generate_embeddings(model_name.clone(), ipt.clone(), None)
                 .await?;
             all_embeddings.push(embed.embeddings);
         }
@@ -66,9 +65,41 @@ impl EmbeddingProvider for OllamaProvider {
 }
 
 impl OllamaProvider {
-    pub async fn generate_response(&self, prompt_text: String) -> Result<String, VectorizeError> {
-        let req = GenerationRequest::new(self.model_name.clone(), prompt_text);
+    pub async fn generate_response(
+        &self,
+        model_name: String,
+        prompt_text: &[ChatMessageRequest],
+    ) -> Result<String, VectorizeError> {
+        let single_prompt: String = prompt_text
+            .iter()
+            .map(|x| x.content.clone())
+            .collect::<Vec<String>>()
+            .join("\n\n");
+        let req = GenerationRequest::new(model_name, single_prompt.to_owned());
         let res = self.instance.generate(req).await?;
         Ok(res.response)
     }
 }
+
+pub fn check_model_host(url: &str) -> Result<String, String> {
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .enable_io()
+        .enable_time()
+        .build()
+        .unwrap_or_else(|e| panic!("failed to initialize tokio runtime: {}", e));
+
+    runtime.block_on(async {
+        let response = reqwest::get(url).await.unwrap();
+        match response.status() {
+            reqwest::StatusCode::OK => Ok(format!("Success! {:?}", response)),
+            _ => Err(format!("Error! {:?}", response)),
+        }
+    })
+}
+
+pub fn ollama_embedding_dim(model_name: &str) -> i32 {
+    match model_name {
+        "llama2" => 5192,
+        _ => 1536,
+    }
+}

core/src/transformers/providers/openai.rs

@@ -1,9 +1,13 @@
 use reqwest::Client;
 use serde::{Deserialize, Serialize};
 
-use super::{EmbeddingProvider, GenericEmbeddingRequest, GenericEmbeddingResponse};
+use super::{
+    ChatMessageRequest, ChatResponse, EmbeddingProvider, GenericEmbeddingRequest,
+    GenericEmbeddingResponse,
+};
 use crate::errors::VectorizeError;
 use crate::transformers::http_handler::handle_response;
+use crate::transformers::providers;
 use crate::transformers::types::Inputs;
 use async_trait::async_trait;
 use std::env;
@@ -75,11 +79,10 @@ impl EmbeddingProvider for OpenAIProvider {
         request: &'a GenericEmbeddingRequest,
     ) -> Result<GenericEmbeddingResponse, VectorizeError> {
         let client = Client::new();
-
         let req = OpenAIEmbeddingBody::from(request.clone());
         let num_inputs = request.input.len();
         let todo_requests: Vec<OpenAIEmbeddingBody> = if num_inputs > 2048 {
-            split_vector(req.input, 2048)
+            providers::split_vector(req.input, 2048)
                 .iter()
                 .map(|chunk| OpenAIEmbeddingBody {
                     input: chunk.clone(),
@@ -128,8 +131,51 @@ pub fn openai_embedding_dim(model_name: &str) -> i32 {
     }
 }
 
-fn split_vector(vec: Vec<String>, chunk_size: usize) -> Vec<Vec<String>> {
-    vec.chunks(chunk_size).map(|chunk| chunk.to_vec()).collect()
+impl OpenAIProvider {
+    pub async fn generate_response(
+        &self,
+        model_name: String,
+        messages: &[ChatMessageRequest],
+    ) -> Result<String, VectorizeError> {
+        let client = Client::new();
+        let chat_url = format!("{}/chat/completions", self.url);
+        let message = serde_json::json!({
+            "model": model_name,
+            "messages": messages,
+        });
+        let response = client
+            .post(&chat_url)
+            .timeout(std::time::Duration::from_secs(120_u64))
+            .header("Accept", "application/json")
+            .header("Content-Type", "application/json")
+            .header("Authorization", &format!("Bearer {}", self.api_key))
+            .json(&message)
+            .send()
+            .await?;
+        let chat_response = handle_response::<ChatResponse>(response, "embeddings").await?;
+        Ok(chat_response.choices[0].message.content.clone())
+    }
+}
+
+// OpenAI embedding model has a limit of 8192 tokens per input
+// there can be a number of ways condense the inputs
+pub fn trim_inputs(inputs: &[Inputs]) -> Vec<String> {
+    inputs
+        .iter()
+        .map(|input| {
+            if input.token_estimate as usize > MAX_TOKEN_LEN {
+                // not example taking tokens, but naive way to trim input
+                let tokens: Vec<&str> = input.inputs.split_whitespace().collect();
+                tokens
+                    .into_iter()
+                    .take(MAX_TOKEN_LEN)
+                    .collect::<Vec<_>>()
+                    .join(" ")
+            } else {
+                input.inputs.clone()
+            }
+        })
+        .collect()
 }
 
 #[cfg(test)]
@@ -161,27 +207,6 @@ mod integration_tests {
     }
 }
 
-// OpenAI embedding model has a limit of 8192 tokens per input
-// there can be a number of ways condense the inputs
-pub fn trim_inputs(inputs: &[Inputs]) -> Vec<String> {
-    inputs
-        .iter()
-        .map(|input| {
-            if input.token_estimate as usize > MAX_TOKEN_LEN {
-                // not example taking tokens, but naive way to trim input
-                let tokens: Vec<&str> = input.inputs.split_whitespace().collect();
-                tokens
-                    .into_iter()
-                    .take(MAX_TOKEN_LEN)
-                    .collect::<Vec<_>>()
-                    .join(" ")
-            } else {
-                input.inputs.clone()
-            }
-        })
-        .collect()
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;