Move shared types into their own crate (#41)

Signed-off-by: José Ulises Niño Rivera <junr03@users.noreply.github.com>

envoyfilter/Cargo.lock

@@ -918,6 +918,7 @@ dependencies = [
  "open-message-format-embeddings",
  "proxy-wasm",
  "proxy-wasm-test-framework",
+ "public-types",
  "serde",
  "serde_json",
  "serde_yaml",
@@ -1377,6 +1378,14 @@ dependencies = [
  "cc",
 ]
 
+[[package]]
+name = "public-types"
+version = "0.1.0"
+dependencies = [
+ "open-message-format-embeddings",
+ "serde",
+]
+
 [[package]]
 name = "quanta"
 version = "0.12.3"

envoyfilter/Cargo.toml

@@ -15,6 +15,7 @@ serde_yaml = "0.9.34"
 serde_json = "1.0"
 md5 = "0.7.0"
 open-message-format-embeddings = { path = "../open-message-format/clients/omf-embeddings-rust" }
+public-types = { path = "../public-types" }
 http = "1.1.0"
 governor = "0.6.3"
 

envoyfilter/src/common_types.rs

@@ -1,89 +0,0 @@
-use crate::configuration::PromptTarget;
-use open_message_format_embeddings::models::CreateEmbeddingRequest;
-use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct EmbeddingRequest {
-    pub create_embedding_request: CreateEmbeddingRequest,
-    pub prompt_target: PromptTarget,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct VectorPoint {
-    pub id: String,
-    pub payload: HashMap<String, String>,
-    pub vector: Vec<f64>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct StoreVectorEmbeddingsRequest {
-    pub points: Vec<VectorPoint>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-#[allow(clippy::large_enum_variant)]
-pub enum CallContext {
-    EmbeddingRequest(EmbeddingRequest),
-    StoreVectorEmbeddings(StoreVectorEmbeddingsRequest),
-    CreateVectorCollection(String),
-}
-
-// https://api.qdrant.tech/master/api-reference/search/points
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct SearchPointsRequest {
-    pub vector: Vec<f64>,
-    pub limit: i32,
-    pub with_payload: bool,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct SearchPointResult {
-    pub id: String,
-    pub version: i32,
-    pub score: f64,
-    pub payload: HashMap<String, String>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct SearchPointsResponse {
-    pub result: Vec<SearchPointResult>,
-    pub status: String,
-    pub time: f64,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct NERRequest {
-    pub input: String,
-    pub labels: Vec<String>,
-    pub model: String,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct Entity {
-    pub text: String,
-    pub label: String,
-    pub score: f64,
-}
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct NERResponse {
-    pub data: Vec<Entity>,
-    pub model: String,
-}
-
-pub mod open_ai {
-    use serde::{Deserialize, Serialize};
-
-    #[derive(Debug, Clone, Serialize, Deserialize)]
-    pub struct ChatCompletions {
-        #[serde(default)]
-        pub model: String,
-        pub messages: Vec<Message>,
-    }
-
-    #[derive(Debug, Clone, Serialize, Deserialize)]
-    pub struct Message {
-        pub role: String,
-        pub content: Option<String>,
-    }
-}

envoyfilter/src/configuration.rs

@@ -1,157 +0,0 @@
-use serde::{Deserialize, Serialize};
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct Configuration {
-    pub default_prompt_endpoint: String,
-    pub load_balancing: LoadBalancing,
-    pub timeout_ms: u64,
-    pub embedding_provider: EmbeddingProviver,
-    pub llm_providers: Vec<LlmProvider>,
-    pub system_prompt: Option<String>,
-    pub prompt_targets: Vec<PromptTarget>,
-    pub ratelimits: Option<Vec<Ratelimit>>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct Ratelimit {
-    pub provider: String,
-    pub selector: Header,
-    pub limit: Limit,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct Limit {
-    pub tokens: u32,
-    pub unit: TimeUnit,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub enum TimeUnit {
-    #[serde(rename = "second")]
-    Second,
-    #[serde(rename = "minute")]
-    Minute,
-    #[serde(rename = "hour")]
-    Hour,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
-pub struct Header {
-    pub key: String,
-    pub value: Option<String>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub enum LoadBalancing {
-    #[serde(rename = "round_robin")]
-    RoundRobin,
-    #[serde(rename = "random")]
-    Random,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-//TODO: use enum for model, but if there is a new model, we need to update the code
-pub struct EmbeddingProviver {
-    pub name: String,
-    pub model: String,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-//TODO: use enum for model, but if there is a new model, we need to update the code
-pub struct LlmProvider {
-    pub name: String,
-    pub api_key: Option<String>,
-    pub model: String,
-    pub default: Option<bool>,
-    pub endpoint: Option<EnpointType>,
-}
-#[derive(Debug, Clone, Serialize, Deserialize)]
-#[serde(untagged)]
-pub enum EnpointType {
-    String(String),
-    Struct(Endpoint),
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct Endpoint {
-    pub cluster: String,
-    pub path: Option<String>,
-    pub method: Option<String>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct Entity {
-    pub name: String,
-    pub required: Option<bool>,
-    pub description: Option<String>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct PromptTarget {
-    #[serde(rename = "type")]
-    pub prompt_type: String,
-    pub name: String,
-    pub few_shot_examples: Vec<String>,
-    pub entities: Option<Vec<Entity>>,
-    pub endpoint: Option<Endpoint>,
-    pub system_prompt: Option<String>,
-}
-
-#[cfg(test)]
-mod test {
-    pub const CONFIGURATION: &str = r#"
-default_prompt_endpoint: "127.0.0.1"
-load_balancing: "round_robin"
-timeout_ms: 5000
-
-embedding_provider:
-  name: "SentenceTransformer"
-  model: "all-MiniLM-L6-v2"
-
-llm_providers:
-  - name: "open-ai-gpt-4"
-    api_key: "$OPEN_AI_API_KEY"
-    model: gpt-4
-
-system_prompt: |
-  You are a helpful weather forecaster. Please following following guidelines when responding to user queries:
-  - Use farenheight for temperature
-  - Use miles per hour for wind speed
-
-prompt_targets:
-  - type: context_resolver
-    name: weather_forecast
-    few_shot_examples:
-      - what is the weather in New York?
-    endpoint:
-      cluster: weatherhost
-      path: /weather
-    entities:
-      - name: location
-        required: true
-        description: "The location for which the weather is requested"
-
-  - type: context_resolver
-    name: weather_forecast_2
-    few_shot_examples:
-      - what is the weather in New York?
-    endpoint:
-      cluster: weatherhost
-      path: /weather
-    entities:
-      - name: city
-
-ratelimits:
-  - provider: open-ai-gpt-4
-    selector:
-      key: x-katanemo-openai-limit-id
-    limit:
-      tokens: 100
-      unit: minute
-  "#;
-
-    #[test]
-    fn test_deserialize_configuration() {
-        let _: super::Configuration = serde_yaml::from_str(CONFIGURATION).unwrap();
-    }
-}

envoyfilter/src/filter_context.rs

@@ -1,7 +1,3 @@
-use crate::common_types::{
-    CallContext, EmbeddingRequest, StoreVectorEmbeddingsRequest, VectorPoint,
-};
-use crate::configuration::{Configuration, PromptTarget};
 use crate::consts::DEFAULT_EMBEDDING_MODEL;
 use crate::ratelimit;
 use crate::stats::{Gauge, RecordingMetric};
@@ -13,6 +9,10 @@ use open_message_format_embeddings::models::{
 };
 use proxy_wasm::traits::*;
 use proxy_wasm::types::*;
+use public_types::common_types::{
+    CallContext, EmbeddingRequest, StoreVectorEmbeddingsRequest, VectorPoint,
+};
+use public_types::configuration::{Configuration, PromptTarget};
 use serde_json::to_string;
 use std::collections::HashMap;
 use std::time::Duration;

envoyfilter/src/lib.rs

@@ -2,8 +2,6 @@ use filter_context::FilterContext;
 use proxy_wasm::traits::*;
 use proxy_wasm::types::*;
 
-mod common_types;
-mod configuration;
 mod consts;
 mod filter_context;
 mod ratelimit;

envoyfilter/src/ratelimit.rs

@@ -1,6 +1,6 @@
-use crate::configuration;
-use crate::configuration::{Limit, Ratelimit, TimeUnit};
 use governor::{DefaultKeyedRateLimiter, InsufficientCapacity, Quota};
+use public_types::configuration;
+use public_types::configuration::{Limit, Ratelimit, TimeUnit};
 use std::num::{NonZero, NonZeroU32};
 use std::sync::RwLock;
 use std::{collections::HashMap, sync::OnceLock};
@@ -376,8 +376,8 @@ fn different_provider_can_have_different_limits_with_the_same_keys() {
 #[cfg(test)]
 mod test {
     use super::ratelimits;
-    use crate::configuration;
     use configuration::{Limit, Ratelimit, TimeUnit};
+    use public_types::configuration;
     use std::num::NonZero;
     use std::thread;
 

envoyfilter/src/stream_context.rs

@@ -1,8 +1,3 @@
-use crate::common_types::{
-    open_ai::{ChatCompletions, Message},
-    NERRequest, NERResponse, SearchPointsRequest, SearchPointsResponse,
-};
-use crate::configuration::{Entity, PromptTarget};
 use crate::consts::{
     DEFAULT_COLLECTION_NAME, DEFAULT_EMBEDDING_MODEL, DEFAULT_NER_MODEL, DEFAULT_NER_THRESHOLD,
     DEFAULT_PROMPT_TARGET_THRESHOLD, SYSTEM_ROLE, USER_ROLE,
@@ -16,6 +11,11 @@ use open_message_format_embeddings::models::{
 };
 use proxy_wasm::traits::*;
 use proxy_wasm::types::*;
+use public_types::common_types::{
+    open_ai::{ChatCompletions, Message},
+    NERRequest, NERResponse, SearchPointsRequest, SearchPointsResponse,
+};
+use public_types::configuration::{Entity, PromptTarget};
 use std::collections::HashMap;
 use std::time::Duration;
 

envoyfilter/tests/integration.rs

@@ -1,6 +1,7 @@
 use http::StatusCode;
 use proxy_wasm_test_framework::tester;
 use proxy_wasm_test_framework::types::{Action, BufferType, MapType, MetricType, ReturnType};
+use public_types::common_types::Entity;
 use serial_test::serial;
 use std::path::Path;
 
@@ -177,3 +178,92 @@ fn bad_request_to_open_ai_chat_completions() {
         .execute_and_expect(ReturnType::Action(Action::Pause))
         .unwrap();
 }
+
+#[test]
+#[serial]
+fn delete_me_in_next_pr_successful_request_to_open_ai_chat_completions() {
+    let ner_response = Entity {
+        score: 0.7,
+        text: String::from("hello"),
+        label: String::from("hello"),
+    };
+    let ner_response_buffer = serde_json::to_string(&ner_response).unwrap();
+    println!("{} is my length", ner_response_buffer.len());
+
+    let args = tester::MockSettings {
+        wasm_path: wasm_module(),
+        quiet: false,
+        allow_unexpected: false,
+    };
+    let mut module = tester::mock(args).unwrap();
+
+    module
+        .call_start()
+        .execute_and_expect(ReturnType::None)
+        .unwrap();
+
+    // Setup Filter
+    let root_context = 1;
+
+    module
+        .call_proxy_on_context_create(root_context, 0)
+        .expect_metric_creation(MetricType::Gauge, "active_http_calls")
+        .execute_and_expect(ReturnType::None)
+        .unwrap();
+
+    // Setup HTTP Stream
+    let http_context = 2;
+
+    module
+        .call_proxy_on_context_create(http_context, root_context)
+        .execute_and_expect(ReturnType::None)
+        .unwrap();
+
+    // Request Headers
+    module
+        .call_proxy_on_request_headers(http_context, 0, false)
+        .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some(":host"))
+        .returning(Some("api.openai.com"))
+        .expect_add_header_map_value(
+            Some(MapType::HttpRequestHeaders),
+            Some("content-length"),
+            Some(""),
+        )
+        .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some(":path"))
+        .returning(Some("/llmrouting"))
+        .expect_add_header_map_value(
+            Some(MapType::HttpRequestHeaders),
+            Some(":path"),
+            Some("/v1/chat/completions"),
+        )
+        .execute_and_expect(ReturnType::Action(Action::Continue))
+        .unwrap();
+
+    // Request Body
+    let chat_completions_request_body = "\
+    {\
+        \"messages\": [\
+        {\
+            \"role\": \"system\",\
+            \"content\": \"You are a poetic assistant, skilled in explaining complex programming concepts with creative flair.\"\
+        },\
+        {\
+            \"role\": \"user\",\
+            \"content\": \"Compose a poem that explains the concept of recursion in programming.\"\
+        }\
+        ]\
+    }";
+
+    module
+        .call_proxy_on_request_body(
+            http_context,
+            chat_completions_request_body.len() as i32,
+            true,
+        )
+        .expect_get_buffer_bytes(Some(BufferType::HttpRequestBody))
+        .returning(Some(chat_completions_request_body))
+        // TODO: assert that the model field was added.
+        .expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None)
+        .execute_and_expect(ReturnType::Action(Action::Pause))
+        .unwrap();
+}