apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-07-05 15:52:12 +02:00
Code Issues Projects Releases Packages Wiki Activity

Implement Client trait for StreamContext (#134)

Browse source 
Signed-off-by: José Ulises Niño Rivera <junr03@users.noreply.github.com>
This commit is contained in:
José Ulises Niño Rivera 2024-10-07 19:50:15 -04:00 • committed by GitHub
parent 5bfccd3959
commit c1cfbcd44d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 216 additions and 218 deletions
Download patch file Download diff file Expand all files Collapse all files

12
arch/Cargo.lock generated
View file

@ -410,6 +410,17 @@ dependencies = [
"uuid", "uuid",
] ]
[[package]]
name = "derivative"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fcc3dd5e9e9c0b295d6e1e4d811fb6f157d5ffd784b8d202fc62eac8035a770b"
dependencies = [
"proc-macro2",
"quote",
"syn 1.0.109",
]
[[package]] [[package]]
name = "digest" name = "digest"
version = "0.10.7" version = "0.10.7"
@ -747,6 +758,7 @@ name = "intelligent-prompt-gateway"
version = "0.1.0" version = "0.1.0"
dependencies = [ dependencies = [
"acap", "acap",
"derivative",
"governor", "governor",
"http", "http",
"log", "log",

1
arch/Cargo.toml
View file

@ -21,6 +21,7 @@ tiktoken-rs = "0.5.9"
acap = "0.3.0" acap = "0.3.0"
rand = "0.8.5" rand = "0.8.5"
thiserror = "1.0.64" thiserror = "1.0.64"
derivative = "2.2.0"
sha2 = "0.10.8" sha2 = "0.10.8"
[dev-dependencies] [dev-dependencies]

1
arch/src/filter_context.rs
View file

@ -97,6 +97,7 @@ impl FilterContext {
let call_args = CallArgs::new( let call_args = CallArgs::new(
MODEL_SERVER_NAME, MODEL_SERVER_NAME,
"/embeddings",
vec![ vec![
(":method", "POST"), (":method", "POST"),
(":path", "/embeddings"), (":path", "/embeddings"),

16
arch/src/http.rs
View file

@ -1,12 +1,16 @@
use crate::stats::{Gauge, IncrementingMetric}; use crate::stats::{Gauge, IncrementingMetric};
use derivative::Derivative;
use log::debug; use log::debug;
use proxy_wasm::{traits::Context, types::Status}; use proxy_wasm::{traits::Context, types::Status};
use std::{cell::RefCell, collections::HashMap, fmt::Debug, time::Duration}; use std::{cell::RefCell, collections::HashMap, fmt::Debug, time::Duration};
#[derive(Debug)] #[derive(Derivative)]
#[derivative(Debug)]
pub struct CallArgs<'a> { pub struct CallArgs<'a> {
upstream: &'a str, upstream: &'a str,
path: &'a str,
headers: Vec<(&'a str, &'a str)>, headers: Vec<(&'a str, &'a str)>,
#[derivative(Debug = "ignore")]
body: Option<&'a [u8]>, body: Option<&'a [u8]>,
trailers: Vec<(&'a str, &'a str)>, trailers: Vec<(&'a str, &'a str)>,
timeout: Duration, timeout: Duration,
@ -15,6 +19,7 @@ pub struct CallArgs<'a> {
impl<'a> CallArgs<'a> { impl<'a> CallArgs<'a> {
pub fn new( pub fn new(
upstream: &'a str, upstream: &'a str,
path: &'a str,
headers: Vec<(&'a str, &'a str)>, headers: Vec<(&'a str, &'a str)>,
body: Option<&'a [u8]>, body: Option<&'a [u8]>,
trailers: Vec<(&'a str, &'a str)>, trailers: Vec<(&'a str, &'a str)>,
@ -22,6 +27,7 @@ impl<'a> CallArgs<'a> {
) -> Self { ) -> Self {
CallArgs { CallArgs {
upstream, upstream,
path,
headers, headers,
body, body,
trailers, trailers,
@ -32,9 +38,10 @@ impl<'a> CallArgs<'a> {
#[derive(thiserror::Error, Debug)] #[derive(thiserror::Error, Debug)]
pub enum ClientError { pub enum ClientError {
#[error("Error dispatching HTTP call to `{upstream_name}`, error: {internal_status:?}")] #[error("Error dispatching HTTP call to `{upstream_name}/{path}`, error: {internal_status:?}")]
DispatchError { DispatchError {
upstream_name: String, upstream_name: String,
path: String,
internal_status: Status, internal_status: Status,
}, },
} }
@ -46,7 +53,7 @@ pub trait Client: Context {
&self, &self,
call_args: CallArgs, call_args: CallArgs,
call_context: Self::CallContext, call_context: Self::CallContext,
) -> Result<(), ClientError> { ) -> Result<u32, ClientError> {
debug!( debug!(
"dispatching http call with args={:?} context={:?}", "dispatching http call with args={:?} context={:?}",
call_args, call_context call_args, call_context
@ -61,10 +68,11 @@ pub trait Client: Context {
) { ) {
Ok(id) => { Ok(id) => {
self.add_call_context(id, call_context); self.add_call_context(id, call_context);
Ok(()) Ok(id)
} }
Err(status) => Err(ClientError::DispatchError { Err(status) => Err(ClientError::DispatchError {
upstream_name: String::from(call_args.upstream), upstream_name: String::from(call_args.upstream),
path: String::from(call_args.path),
internal_status: status, internal_status: status,
}), }),
} }

31
arch/src/ratelimit.rs
View file

@ -2,6 +2,7 @@ use governor::{DefaultKeyedRateLimiter, InsufficientCapacity, Quota};
use log::debug; use log::debug;
use public_types::configuration; use public_types::configuration;
use public_types::configuration::{Limit, Ratelimit, TimeUnit}; use public_types::configuration::{Limit, Ratelimit, TimeUnit};
use std::fmt::Display;
use std::num::{NonZero, NonZeroU32}; use std::num::{NonZero, NonZeroU32};
use std::sync::RwLock; use std::sync::RwLock;
use std::{collections::HashMap, sync::OnceLock}; use std::{collections::HashMap, sync::OnceLock};
@ -28,13 +29,18 @@ pub struct RatelimitMap {
} }
// This version of Header demands that the user passes a header value to match on. // This version of Header demands that the user passes a header value to match on.
#[allow(unused)] #[derive(Debug, Clone)]
#[derive(Debug)]
pub struct Header { pub struct Header {
pub key: String, pub key: String,
pub value: String, pub value: String,
} }
impl Display for Header {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{self:?}")
}
}
impl From<Header> for configuration::Header { impl From<Header> for configuration::Header {
fn from(header: Header) -> Self { fn from(header: Header) -> Self {
Self { Self {
@ -44,6 +50,16 @@ impl From<Header> for configuration::Header {
} }
} }
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("exceeded limit provider={provider}, selector={selector}, tokens_used={tokens_used}")]
ExceededLimit {
provider: String,
selector: Header,
tokens_used: NonZeroU32,
},
}
impl RatelimitMap { impl RatelimitMap {
// n.b new is private so that the only access to the Ratelimits can be done via the static // n.b new is private so that the only access to the Ratelimits can be done via the static
// reference inside a RwLock via ratelimit::ratelimits(). // reference inside a RwLock via ratelimit::ratelimits().
@ -82,7 +98,7 @@ impl RatelimitMap {
provider: String, provider: String,
selector: Header, selector: Header,
tokens_used: NonZeroU32, tokens_used: NonZeroU32,
) -> Result<(), String> { ) -> Result<(), Error> {
debug!( debug!(
"Checking limit for provider={}, with selector={:?}, consuming tokens={:?}", "Checking limit for provider={}, with selector={:?}, consuming tokens={:?}",
provider, selector, tokens_used provider, selector, tokens_used
@ -96,7 +112,7 @@ impl RatelimitMap {
Some(limit) => limit, Some(limit) => limit,
}; };
let mut config_selector = configuration::Header::from(selector); let mut config_selector = configuration::Header::from(selector.clone());
let (limit, limit_key) = match provider_limits.get(&config_selector) { let (limit, limit_key) = match provider_limits.get(&config_selector) {
// This is a specific limit, i.e one that was configured with both key, and value. // This is a specific limit, i.e one that was configured with both key, and value.
@ -119,8 +135,11 @@ impl RatelimitMap {
match limit.check_key_n(&limit_key, tokens_used) { match limit.check_key_n(&limit_key, tokens_used) {
Ok(Ok(())) => Ok(()), Ok(Ok(())) => Ok(()),
Ok(Err(_)) => Err(String::from("Not allowed")), Ok(Err(_)) | Err(InsufficientCapacity(_)) => Err(Error::ExceededLimit {
Err(InsufficientCapacity(_)) => Err(String::from("Not allowed")), provider,
selector,
tokens_used,
}),
} }
} }
} }

360
arch/src/stream_context.rs
View file

@ -5,6 +5,7 @@ use crate::consts::{
RATELIMIT_SELECTOR_HEADER_KEY, SYSTEM_ROLE, USER_ROLE, RATELIMIT_SELECTOR_HEADER_KEY, SYSTEM_ROLE, USER_ROLE,
}; };
use crate::filter_context::{EmbeddingsStore, WasmMetrics}; use crate::filter_context::{EmbeddingsStore, WasmMetrics};
use crate::http::{CallArgs, Client, ClientError};
use crate::llm_providers::LlmProviders; use crate::llm_providers::LlmProviders;
use crate::ratelimit::Header; use crate::ratelimit::Header;
use crate::stats::IncrementingMetric; use crate::stats::IncrementingMetric;
@ -30,11 +31,13 @@ use public_types::embeddings::{
}; };
use serde_json::Value; use serde_json::Value;
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
use std::cell::RefCell;
use std::collections::HashMap; use std::collections::HashMap;
use std::num::NonZero; use std::num::NonZero;
use std::rc::Rc; use std::rc::Rc;
use std::time::Duration; use std::time::Duration;
#[derive(Debug)]
enum ResponseHandlerType { enum ResponseHandlerType {
GetEmbeddings, GetEmbeddings,
FunctionResolver, FunctionResolver,
@ -44,7 +47,8 @@ enum ResponseHandlerType {
DefaultTarget, DefaultTarget,
} }
pub struct CallContext { #[derive(Debug)]
pub struct StreamCallContext {
response_handler_type: ResponseHandlerType, response_handler_type: ResponseHandlerType,
user_message: Option<String>, user_message: Option<String>,
prompt_target_name: Option<String>, prompt_target_name: Option<String>,
@ -54,13 +58,37 @@ pub struct CallContext {
upstream_cluster_path: Option<String>, upstream_cluster_path: Option<String>,
} }
#[derive(thiserror::Error, Debug)]
pub enum ServerError {
#[error(transparent)]
HttpDispatch(ClientError),
#[error(transparent)]
Deserialization(serde_json::Error),
#[error(transparent)]
Serialization(serde_json::Error),
#[error("{0}")]
LogicError(String),
#[error("upstream error response authority={authority}, path={path}, status={status}")]
Upstream {
authority: String,
path: String,
status: String,
},
#[error(transparent)]
ExceededRatelimit(ratelimit::Error),
#[error("jailbreak detected: {0}")]
Jailbreak(String),
#[error("{why}")]
BadRequest { why: String },
}
pub struct StreamContext { pub struct StreamContext {
context_id: u32, context_id: u32,
metrics: Rc<WasmMetrics>, metrics: Rc<WasmMetrics>,
prompt_targets: Rc<HashMap<String, PromptTarget>>, prompt_targets: Rc<HashMap<String, PromptTarget>>,
embeddings_store: Rc<EmbeddingsStore>, embeddings_store: Rc<EmbeddingsStore>,
overrides: Rc<Option<Overrides>>, overrides: Rc<Option<Overrides>>,
callouts: HashMap<u32, CallContext>, callouts: RefCell<HashMap<u32, StreamCallContext>>,
tool_calls: Option<Vec<ToolCall>>, tool_calls: Option<Vec<ToolCall>>,
tool_call_response: Option<String>, tool_call_response: Option<String>,
arch_state: Option<Vec<ArchState>>, arch_state: Option<Vec<ArchState>>,
@ -91,8 +119,8 @@ impl StreamContext {
metrics, metrics,
prompt_targets, prompt_targets,
embeddings_store, embeddings_store,
callouts: RefCell::new(HashMap::new()),
chat_completions_request: None, chat_completions_request: None,
callouts: HashMap::new(),
tool_calls: None, tool_calls: None,
tool_call_response: None, tool_call_response: None,
arch_state: None, arch_state: None,
@ -129,11 +157,17 @@ impl StreamContext {
self.add_http_request_header(ARCH_ROUTING_HEADER, &self.llm_provider().name); self.add_http_request_header(ARCH_ROUTING_HEADER, &self.llm_provider().name);
} }
fn modify_auth_headers(&mut self) -> Result<(), String> { fn modify_auth_headers(&mut self) -> Result<(), ServerError> {
let llm_provider_api_key_value = self.llm_provider().access_key.as_ref().ok_or(format!( let llm_provider_api_key_value =
"No access key configured for selected LLM Provider \"{}\"",
self.llm_provider() self.llm_provider()
))?; .access_key
.as_ref()
.ok_or(ServerError::BadRequest {
why: format!(
"No access key configured for selected LLM Provider \"{}\"",
self.llm_provider()
),
})?;
let authorization_header_value = format!("Bearer {}", llm_provider_api_key_value); let authorization_header_value = format!("Bearer {}", llm_provider_api_key_value);
@ -159,7 +193,7 @@ impl StreamContext {
}); });
} }
fn send_server_error(&self, error: String, override_status_code: Option<StatusCode>) { fn send_server_error(&self, error: ServerError, override_status_code: Option<StatusCode>) {
debug!("server error occurred: {}", error); debug!("server error occurred: {}", error);
self.send_http_response( self.send_http_response(
override_status_code override_status_code
@ -167,18 +201,15 @@ impl StreamContext {
.as_u16() .as_u16()
.into(), .into(),
vec![], vec![],
Some(error.as_bytes()), Some(format!("{error}").as_bytes()),
); );
} }
fn embeddings_handler(&mut self, body: Vec<u8>, mut callout_context: CallContext) { fn embeddings_handler(&mut self, body: Vec<u8>, mut callout_context: StreamCallContext) {
let embedding_response: CreateEmbeddingResponse = match serde_json::from_slice(&body) { let embedding_response: CreateEmbeddingResponse = match serde_json::from_slice(&body) {
Ok(embedding_response) => embedding_response, Ok(embedding_response) => embedding_response,
Err(e) => { Err(e) => {
return self.send_server_error( return self.send_server_error(ServerError::Deserialization(e), None);
format!("Error deserializing embedding response: {:?}", e),
None,
);
} }
}; };
@ -248,13 +279,13 @@ impl StreamContext {
let json_data: String = match serde_json::to_string(&zero_shot_classification_request) { let json_data: String = match serde_json::to_string(&zero_shot_classification_request) {
Ok(json_data) => json_data, Ok(json_data) => json_data,
Err(error) => { Err(error) => {
let error = format!("Error serializing zero shot request: {}", error); return self.send_server_error(ServerError::Serialization(error), None);
return self.send_server_error(error, None);
} }
}; };
let token_id = match self.dispatch_http_call( let call_args = CallArgs::new(
MODEL_SERVER_NAME, MODEL_SERVER_NAME,
"/zeroshot",
vec![ vec![
(":method", "POST"), (":method", "POST"),
(":path", "/zeroshot"), (":path", "/zeroshot"),
@ -266,49 +297,24 @@ impl StreamContext {
Some(json_data.as_bytes()), Some(json_data.as_bytes()),
vec![], vec![],
Duration::from_secs(5), Duration::from_secs(5),
) {
Ok(token_id) => token_id,
Err(e) => {
let error_msg = format!(
"Error dispatching embedding server HTTP call for zero-shot-intent-detection: {:?}",
e
);
return self.send_server_error(error_msg, None);
}
};
debug!(
"dispatched call to model_server/zeroshot token_id={}",
token_id
); );
self.metrics.active_http_calls.increment(1);
callout_context.response_handler_type = ResponseHandlerType::ZeroShotIntent; callout_context.response_handler_type = ResponseHandlerType::ZeroShotIntent;
if self.callouts.insert(token_id, callout_context).is_some() { if let Err(e) = self.http_call(call_args, callout_context) {
panic!( self.send_server_error(ServerError::HttpDispatch(e), None);
"duplicate token_id={} in embedding server requests",
token_id
)
} }
} }
fn zero_shot_intent_detection_resp_handler( fn zero_shot_intent_detection_resp_handler(
&mut self, &mut self,
body: Vec<u8>, body: Vec<u8>,
mut callout_context: CallContext, mut callout_context: StreamCallContext,
) { ) {
let zeroshot_intent_response: ZeroShotClassificationResponse = let zeroshot_intent_response: ZeroShotClassificationResponse =
match serde_json::from_slice(&body) { match serde_json::from_slice(&body) {
Ok(zeroshot_response) => zeroshot_response, Ok(zeroshot_response) => zeroshot_response,
Err(e) => { Err(e) => {
self.send_server_error( return self.send_server_error(ServerError::Deserialization(e), None);
format!(
"Error deserializing zeroshot intent detection response: {:?}",
e
),
None,
);
return;
} }
}; };
@ -390,40 +396,34 @@ impl StreamContext {
); );
let arch_messages_json = serde_json::to_string(&params).unwrap(); let arch_messages_json = serde_json::to_string(&params).unwrap();
debug!("no prompt target found with similarity score above threshold, using default prompt target"); debug!("no prompt target found with similarity score above threshold, using default prompt target");
let token_id = match self.dispatch_http_call(
let timeout_str = ARCH_FC_REQUEST_TIMEOUT_MS.to_string();
let call_args = CallArgs::new(
&upstream_endpoint, &upstream_endpoint,
&upstream_path,
vec![ vec![
(":method", "POST"), (":method", "POST"),
(":path", &upstream_path), (":path", &upstream_path),
(":authority", &upstream_endpoint), (":authority", &upstream_endpoint),
("content-type", "application/json"), ("content-type", "application/json"),
("x-envoy-max-retries", "3"), ("x-envoy-max-retries", "3"),
( ("x-envoy-upstream-rq-timeout-ms", timeout_str.as_str()),
"x-envoy-upstream-rq-timeout-ms",
ARCH_FC_REQUEST_TIMEOUT_MS.to_string().as_str(),
),
], ],
Some(arch_messages_json.as_bytes()), Some(arch_messages_json.as_bytes()),
vec![], vec![],
Duration::from_secs(5), Duration::from_secs(5),
) { );
Ok(token_id) => token_id,
Err(e) => {
let error_msg =
format!("Error dispatching HTTP call for default-target: {:?}", e);
return self
.send_server_error(error_msg, Some(StatusCode::BAD_REQUEST));
}
};
self.metrics.active_http_calls.increment(1);
callout_context.response_handler_type = ResponseHandlerType::DefaultTarget; callout_context.response_handler_type = ResponseHandlerType::DefaultTarget;
callout_context.prompt_target_name = Some(default_prompt_target.name.clone()); callout_context.prompt_target_name = Some(default_prompt_target.name.clone());
if self.callouts.insert(token_id, callout_context).is_some() {
panic!("duplicate token_id") if let Err(e) = self.http_call(call_args, callout_context) {
return self.send_server_error(
ServerError::HttpDispatch(e),
Some(StatusCode::BAD_REQUEST),
);
} }
return;
} }
self.resume_http_request(); self.resume_http_request();
return; return;
} }
@ -433,7 +433,9 @@ impl StreamContext {
Some(prompt_target) => prompt_target.clone(), Some(prompt_target) => prompt_target.clone(),
None => { None => {
return self.send_server_error( return self.send_server_error(
format!("Prompt target not found: {}", prompt_target_name), ServerError::LogicError(format!(
"Prompt target not found: {prompt_target_name}"
)),
None, None,
); );
} }
@ -499,62 +501,42 @@ impl StreamContext {
msg_body msg_body
} }
Err(e) => { Err(e) => {
return self return self.send_server_error(ServerError::Serialization(e), None);
.send_server_error(format!("Error serializing request_params: {:?}", e), None);
} }
}; };
let token_id = match self.dispatch_http_call( let timeout_str = ARCH_FC_REQUEST_TIMEOUT_MS.to_string();
let call_args = CallArgs::new(
ARC_FC_CLUSTER, ARC_FC_CLUSTER,
"/v1/chat/completions",
vec![ vec![
(":method", "POST"), (":method", "POST"),
(":path", "/v1/chat/completions"), (":path", "/v1/chat/completions"),
(":authority", ARC_FC_CLUSTER), (":authority", ARC_FC_CLUSTER),
("content-type", "application/json"), ("content-type", "application/json"),
("x-envoy-max-retries", "3"), ("x-envoy-max-retries", "3"),
( ("x-envoy-upstream-rq-timeout-ms", timeout_str.as_str()),
"x-envoy-upstream-rq-timeout-ms",
ARCH_FC_REQUEST_TIMEOUT_MS.to_string().as_str(),
),
], ],
Some(msg_body.as_bytes()), Some(msg_body.as_bytes()),
vec![], vec![],
Duration::from_secs(5), Duration::from_secs(5),
) {
Ok(token_id) => token_id,
Err(e) => {
let error_msg = format!("Error dispatching HTTP call for function-call: {:?}", e);
return self.send_server_error(error_msg, Some(StatusCode::BAD_REQUEST));
}
};
debug!(
"dispatched call to function {} token_id={}",
ARC_FC_CLUSTER, token_id
); );
self.metrics.active_http_calls.increment(1);
callout_context.response_handler_type = ResponseHandlerType::FunctionResolver; callout_context.response_handler_type = ResponseHandlerType::FunctionResolver;
callout_context.prompt_target_name = Some(prompt_target.name); callout_context.prompt_target_name = Some(prompt_target.name);
if self.callouts.insert(token_id, callout_context).is_some() {
panic!("duplicate token_id") if let Err(e) = self.http_call(call_args, callout_context) {
self.send_server_error(ServerError::HttpDispatch(e), Some(StatusCode::BAD_REQUEST));
} }
} }
fn function_resolver_handler(&mut self, body: Vec<u8>, mut callout_context: CallContext) { fn function_resolver_handler(&mut self, body: Vec<u8>, mut callout_context: StreamCallContext) {
let body_str = String::from_utf8(body).unwrap(); let body_str = String::from_utf8(body).unwrap();
debug!("arch <= app response body: {}", body_str); debug!("arch <= app response body: {}", body_str);
let arch_fc_response: ChatCompletionsResponse = match serde_json::from_str(&body_str) { let arch_fc_response: ChatCompletionsResponse = match serde_json::from_str(&body_str) {
Ok(arch_fc_response) => arch_fc_response, Ok(arch_fc_response) => arch_fc_response,
Err(e) => { Err(e) => {
return self.send_server_error( return self.send_server_error(ServerError::Deserialization(e), None);
format!(
"Error deserializing function resolver response into ChatCompletion: {:?}",
e
),
None,
);
} }
}; };
@ -607,11 +589,12 @@ impl StreamContext {
let endpoint = prompt_target.endpoint.unwrap(); let endpoint = prompt_target.endpoint.unwrap();
let path: String = endpoint.path.unwrap_or(String::from("/")); let path: String = endpoint.path.unwrap_or(String::from("/"));
let token_id = match self.dispatch_http_call( let call_args = CallArgs::new(
&endpoint.name, &endpoint.name,
&path,
vec![ vec![
(":method", "POST"), (":method", "POST"),
(":path", path.as_ref()), (":path", &path),
(":authority", endpoint.name.as_str()), (":authority", endpoint.name.as_str()),
("content-type", "application/json"), ("content-type", "application/json"),
("x-envoy-max-retries", "3"), ("x-envoy-max-retries", "3"),
@ -619,39 +602,33 @@ impl StreamContext {
Some(tool_params_json_str.as_bytes()), Some(tool_params_json_str.as_bytes()),
vec![], vec![],
Duration::from_secs(5), Duration::from_secs(5),
) { );
Ok(token_id) => token_id, callout_context.upstream_cluster = Some(endpoint.name.clone());
Err(e) => { callout_context.upstream_cluster_path = Some(path.clone());
let error_msg = format!( callout_context.response_handler_type = ResponseHandlerType::FunctionCall;
"Error dispatching call to cluster: {}, path: {}, err: {:?}",
&endpoint.name, path, e
);
debug!("{}", error_msg);
return self.send_server_error(error_msg, Some(StatusCode::BAD_REQUEST));
}
};
self.tool_calls = Some(tool_calls.clone()); self.tool_calls = Some(tool_calls.clone());
callout_context.upstream_cluster = Some(endpoint.name);
callout_context.upstream_cluster_path = Some(path); if let Err(e) = self.http_call(call_args, callout_context) {
callout_context.response_handler_type = ResponseHandlerType::FunctionCall; self.send_server_error(ServerError::HttpDispatch(e), Some(StatusCode::BAD_REQUEST));
if self.callouts.insert(token_id, callout_context).is_some() {
panic!("duplicate token_id")
} }
self.metrics.active_http_calls.increment(1);
} }
fn function_call_response_handler(&mut self, body: Vec<u8>, callout_context: CallContext) { fn function_call_response_handler(
let headers = self.get_http_call_response_headers(); &mut self,
if let Some(http_status) = headers.iter().find(|(key, _)| key == ":status") { body: Vec<u8>,
if http_status.1 != StatusCode::OK.as_str() { callout_context: StreamCallContext,
let error_msg = format!( ) {
"Error in function call response: cluster: {}, path: {}, status code: {}", if let Some(http_status) = self.get_http_call_response_header(":status") {
callout_context.upstream_cluster.unwrap(), if http_status != StatusCode::OK.as_str() {
callout_context.upstream_cluster_path.unwrap(), return self.send_server_error(
http_status.1 ServerError::Upstream {
authority: callout_context.upstream_cluster.unwrap(),
path: callout_context.upstream_cluster_path.unwrap(),
status: http_status,
},
None,
); );
return self.send_server_error(error_msg, Some(StatusCode::BAD_REQUEST));
} }
} else { } else {
warn!("http status code not found in api response"); warn!("http status code not found in api response");
@ -714,8 +691,7 @@ impl StreamContext {
let json_string = match serde_json::to_string(&chat_completions_request) { let json_string = match serde_json::to_string(&chat_completions_request) {
Ok(json_string) => json_string, Ok(json_string) => json_string,
Err(e) => { Err(e) => {
return self return self.send_server_error(ServerError::Serialization(e), None);
.send_server_error(format!("Error serializing request_body: {:?}", e), None);
} }
}; };
debug!("arch => openai request body: {}", json_string); debug!("arch => openai request body: {}", json_string);
@ -733,7 +709,7 @@ impl StreamContext {
Ok(_) => (), Ok(_) => (),
Err(err) => { Err(err) => {
self.send_server_error( self.send_server_error(
format!("Exceeded Ratelimit: {}", err), ServerError::ExceededRatelimit(err),
Some(StatusCode::TOO_MANY_REQUESTS), Some(StatusCode::TOO_MANY_REQUESTS),
); );
self.metrics.ratelimited_rq.increment(1); self.metrics.ratelimited_rq.increment(1);
@ -747,7 +723,7 @@ impl StreamContext {
self.resume_http_request(); self.resume_http_request();
} }
fn arch_guard_handler(&mut self, body: Vec<u8>, callout_context: CallContext) { fn arch_guard_handler(&mut self, body: Vec<u8>, callout_context: StreamCallContext) {
debug!("response received for arch guard"); debug!("response received for arch guard");
let prompt_guard_resp: PromptGuardResponse = serde_json::from_slice(&body).unwrap(); let prompt_guard_resp: PromptGuardResponse = serde_json::from_slice(&body).unwrap();
debug!("prompt_guard_resp: {:?}", prompt_guard_resp); debug!("prompt_guard_resp: {:?}", prompt_guard_resp);
@ -757,14 +733,17 @@ impl StreamContext {
let msg = self let msg = self
.prompt_guards .prompt_guards
.jailbreak_on_exception_message() .jailbreak_on_exception_message()
.unwrap_or("Jailbreak detected. Please refrain from discussing jailbreaking."); .unwrap_or("refrain from discussing jailbreaking.");
return self.send_server_error(msg.to_string(), Some(StatusCode::BAD_REQUEST)); return self.send_server_error(
ServerError::Jailbreak(String::from(msg)),
Some(StatusCode::BAD_REQUEST),
);
} }
self.get_embeddings(callout_context); self.get_embeddings(callout_context);
} }
fn get_embeddings(&mut self, callout_context: CallContext) { fn get_embeddings(&mut self, callout_context: StreamCallContext) {
let user_message = callout_context.user_message.unwrap(); let user_message = callout_context.user_message.unwrap();
let get_embeddings_input = CreateEmbeddingRequest { let get_embeddings_input = CreateEmbeddingRequest {
// Need to clone into input because user_message is used below. // Need to clone into input because user_message is used below.
@ -778,13 +757,13 @@ impl StreamContext {
let json_data: String = match serde_json::to_string(&get_embeddings_input) { let json_data: String = match serde_json::to_string(&get_embeddings_input) {
Ok(json_data) => json_data, Ok(json_data) => json_data,
Err(error) => { Err(error) => {
let error_msg = format!("Error serializing embeddings input: {}", error); return self.send_server_error(ServerError::Deserialization(error), None);
return self.send_server_error(error_msg, None);
} }
}; };
let token_id = match self.dispatch_http_call( let call_args = CallArgs::new(
MODEL_SERVER_NAME, MODEL_SERVER_NAME,
"/embeddings",
vec![ vec![
(":method", "POST"), (":method", "POST"),
(":path", "/embeddings"), (":path", "/embeddings"),
@ -796,19 +775,8 @@ impl StreamContext {
Some(json_data.as_bytes()), Some(json_data.as_bytes()),
vec![], vec![],
Duration::from_secs(5), Duration::from_secs(5),
) {
Ok(token_id) => token_id,
Err(e) => {
let error_msg = format!("dispatched call to model_server/embeddings: {:?}", e);
return self.send_server_error(error_msg, None);
}
};
debug!(
"dispatched call to model_server/embeddings token_id={}",
token_id
); );
let call_context = StreamCallContext {
let call_context = CallContext {
response_handler_type: ResponseHandlerType::GetEmbeddings, response_handler_type: ResponseHandlerType::GetEmbeddings,
user_message: Some(user_message), user_message: Some(user_message),
prompt_target_name: None, prompt_target_name: None,
@ -817,17 +785,13 @@ impl StreamContext {
upstream_cluster: None, upstream_cluster: None,
upstream_cluster_path: None, upstream_cluster_path: None,
}; };
if self.callouts.insert(token_id, call_context).is_some() {
panic!(
"duplicate token_id={} in embedding server requests",
token_id
)
}
self.metrics.active_http_calls.increment(1); if let Err(e) = self.http_call(call_args, call_context) {
self.send_server_error(ServerError::HttpDispatch(e), None);
}
} }
fn default_target_handler(&self, body: Vec<u8>, callout_context: CallContext) { fn default_target_handler(&self, body: Vec<u8>, callout_context: StreamCallContext) {
let prompt_target = self let prompt_target = self
.prompt_targets .prompt_targets
.get(callout_context.prompt_target_name.as_ref().unwrap()) .get(callout_context.prompt_target_name.as_ref().unwrap())
@ -856,10 +820,7 @@ impl StreamContext {
let chat_completions_resp: ChatCompletionsResponse = match serde_json::from_slice(&body) { let chat_completions_resp: ChatCompletionsResponse = match serde_json::from_slice(&body) {
Ok(chat_completions_resp) => chat_completions_resp, Ok(chat_completions_resp) => chat_completions_resp,
Err(e) => { Err(e) => {
return self.send_server_error( return self.send_server_error(ServerError::Deserialization(e), None);
format!("Error deserializing default target response: {:?}", e),
None,
);
} }
}; };
let api_resp = chat_completions_resp.choices[0] let api_resp = chat_completions_resp.choices[0]
@ -948,9 +909,9 @@ impl HttpContext for StreamContext {
match self.get_http_request_body(0, body_size) { match self.get_http_request_body(0, body_size) {
Some(body_bytes) => match serde_json::from_slice(&body_bytes) { Some(body_bytes) => match serde_json::from_slice(&body_bytes) {
Ok(deserialized) => deserialized, Ok(deserialized) => deserialized,
Err(msg) => { Err(e) => {
self.send_server_error( self.send_server_error(
format!("Failed to deserialize: {}", msg), ServerError::Deserialization(e),
Some(StatusCode::BAD_REQUEST), Some(StatusCode::BAD_REQUEST),
); );
return Action::Pause; return Action::Pause;
@ -958,10 +919,10 @@ impl HttpContext for StreamContext {
}, },
None => { None => {
self.send_server_error( self.send_server_error(
format!( ServerError::LogicError(format!(
"Failed to obtain body bytes even though body_size is {}", "Failed to obtain body bytes even though body_size is {}",
body_size body_size
), )),
None, None,
); );
return Action::Pause; return Action::Pause;
@ -1018,7 +979,7 @@ impl HttpContext for StreamContext {
if !prompt_guard_jailbreak_task { if !prompt_guard_jailbreak_task {
debug!("Missing input guard. Making inline call to retrieve"); debug!("Missing input guard. Making inline call to retrieve");
let callout_context = CallContext { let callout_context = StreamCallContext {
response_handler_type: ResponseHandlerType::ArchGuard, response_handler_type: ResponseHandlerType::ArchGuard,
user_message: user_message_str.clone(), user_message: user_message_str.clone(),
prompt_target_name: None, prompt_target_name: None,
@ -1046,14 +1007,14 @@ impl HttpContext for StreamContext {
let json_data: String = match serde_json::to_string(&get_prompt_guards_request) { let json_data: String = match serde_json::to_string(&get_prompt_guards_request) {
Ok(json_data) => json_data, Ok(json_data) => json_data,
Err(error) => { Err(error) => {
let error_msg = format!("Error serializing prompt guard request: {}", error); self.send_server_error(ServerError::Serialization(error), None);
self.send_server_error(error_msg, None);
return Action::Pause; return Action::Pause;
} }
}; };
let token_id = match self.dispatch_http_call( let call_args = CallArgs::new(
MODEL_SERVER_NAME, MODEL_SERVER_NAME,
"/guard",
vec![ vec![
(":method", "POST"), (":method", "POST"),
(":path", "/guard"), (":path", "/guard"),
@ -1065,21 +1026,8 @@ impl HttpContext for StreamContext {
Some(json_data.as_bytes()), Some(json_data.as_bytes()),
vec![], vec![],
Duration::from_secs(5), Duration::from_secs(5),
) { );
Ok(token_id) => token_id, let call_context = StreamCallContext {
Err(e) => {
let error_msg = format!(
"Error dispatching embedding server HTTP call for prompt-guard: {:?}",
e
);
self.send_server_error(error_msg, None);
return Action::Pause;
}
};
debug!("dispatched HTTP call to arch_guard token_id={}", token_id);
let call_context = CallContext {
response_handler_type: ResponseHandlerType::ArchGuard, response_handler_type: ResponseHandlerType::ArchGuard,
user_message: self.user_prompt.as_ref().unwrap().content.clone(), user_message: self.user_prompt.as_ref().unwrap().content.clone(),
prompt_target_name: None, prompt_target_name: None,
@ -1088,14 +1036,10 @@ impl HttpContext for StreamContext {
upstream_cluster: None, upstream_cluster: None,
upstream_cluster_path: None, upstream_cluster_path: None,
}; };
if self.callouts.insert(token_id, call_context).is_some() {
panic!(
"duplicate token_id={} in embedding server requests",
token_id
)
}
self.metrics.active_http_calls.increment(1); if let Err(e) = self.http_call(call_args, call_context) {
self.send_server_error(ServerError::HttpDispatch(e), None);
}
Action::Pause Action::Pause
} }
@ -1130,7 +1074,10 @@ impl HttpContext for StreamContext {
let chat_completions_data = match body_str.split_once("data: ") { let chat_completions_data = match body_str.split_once("data: ") {
Some((_, chat_completions_data)) => chat_completions_data, Some((_, chat_completions_data)) => chat_completions_data,
None => { None => {
self.send_server_error(String::from("parsing error in streaming data"), None); self.send_server_error(
ServerError::LogicError(String::from("parsing error in streaming data")),
None,
);
return Action::Pause; return Action::Pause;
} }
}; };
@ -1141,7 +1088,9 @@ impl HttpContext for StreamContext {
Err(_) => { Err(_) => {
if chat_completions_data != "[NONE]" { if chat_completions_data != "[NONE]" {
self.send_server_error( self.send_server_error(
String::from("error in streaming response"), ServerError::LogicError(String::from(
"error in streaming response",
)),
None, None,
); );
return Action::Continue; return Action::Continue;
@ -1168,14 +1117,7 @@ impl HttpContext for StreamContext {
match serde_json::from_slice(&body) { match serde_json::from_slice(&body) {
Ok(de) => de, Ok(de) => de,
Err(e) => { Err(e) => {
self.send_server_error( self.send_server_error(ServerError::Deserialization(e), None);
format!(
"error in non-streaming response: {}\n response was={}",
e,
String::from_utf8(body).unwrap()
),
None,
);
return Action::Pause; return Action::Pause;
} }
}; };
@ -1260,7 +1202,11 @@ impl Context for StreamContext {
body_size: usize, body_size: usize,
_num_trailers: usize, _num_trailers: usize,
) { ) {
let callout_context = self.callouts.remove(&token_id).expect("invalid token_id"); let callout_context = self
.callouts
.get_mut()
.remove(&token_id)
.expect("invalid token_id");
self.metrics.active_http_calls.increment(-1); self.metrics.active_http_calls.increment(-1);
if let Some(body) = self.get_http_call_response_body(0, body_size) { if let Some(body) = self.get_http_call_response_body(0, body_size) {
@ -1284,9 +1230,21 @@ impl Context for StreamContext {
} }
} else { } else {
self.send_server_error( self.send_server_error(
String::from("No response body in inline HTTP request"), ServerError::LogicError(String::from("No response body in inline HTTP request")),
None, None,
); );
} }
} }
} }
impl Client for StreamContext {
type CallContext = StreamCallContext;
fn callouts(&self) -> &RefCell<HashMap<u32, Self::CallContext>> {
&self.callouts
}
fn active_http_calls(&self) -> &crate::stats::Gauge {
&self.metrics.active_http_calls
}
}

13
arch/tests/integration.rs
View file

@ -571,6 +571,7 @@ fn request_ratelimited() {
.expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_http_call( .expect_http_call(
Some("api_server"), Some("api_server"),
Some(vec![ Some(vec![
@ -589,15 +590,14 @@ fn request_ratelimited() {
.execute_and_expect(ReturnType::None) .execute_and_expect(ReturnType::None)
.unwrap(); .unwrap();
let response_headers_with_200 = vec![(":status", "200"), ("content-type", "application/json")];
let body_text = String::from("test body"); let body_text = String::from("test body");
module module
.call_proxy_on_http_call_response(http_context, 5, 0, body_text.len() as i32, 0) .call_proxy_on_http_call_response(http_context, 5, 0, body_text.len() as i32, 0)
.expect_metric_increment("active_http_calls", -1) .expect_metric_increment("active_http_calls", -1)
.expect_get_buffer_bytes(Some(BufferType::HttpCallResponseBody)) .expect_get_buffer_bytes(Some(BufferType::HttpCallResponseBody))
.returning(Some(&body_text)) .returning(Some(&body_text))
.expect_get_header_map_pairs(Some(MapType::HttpCallResponseHeaders)) .expect_get_header_map_value(Some(MapType::HttpCallResponseHeaders), Some(":status"))
.returning(Some(response_headers_with_200)) .returning(Some("200"))
.expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None)
@ -679,6 +679,7 @@ fn request_not_ratelimited() {
.expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_http_call( .expect_http_call(
Some("api_server"), Some("api_server"),
Some(vec![ Some(vec![
@ -697,16 +698,14 @@ fn request_not_ratelimited() {
.execute_and_expect(ReturnType::None) .execute_and_expect(ReturnType::None)
.unwrap(); .unwrap();
let response_headers_with_200 = vec![(":status", "200"), ("content-type", "application/json")];
let body_text = String::from("test body"); let body_text = String::from("test body");
module module
.call_proxy_on_http_call_response(http_context, 5, 0, body_text.len() as i32, 0) .call_proxy_on_http_call_response(http_context, 5, 0, body_text.len() as i32, 0)
.expect_metric_increment("active_http_calls", -1) .expect_metric_increment("active_http_calls", -1)
.expect_get_buffer_bytes(Some(BufferType::HttpCallResponseBody)) .expect_get_buffer_bytes(Some(BufferType::HttpCallResponseBody))
.returning(Some(&body_text)) .returning(Some(&body_text))
.expect_get_header_map_pairs(Some(MapType::HttpCallResponseHeaders)) .expect_get_header_map_value(Some(MapType::HttpCallResponseHeaders), Some(":status"))
.returning(Some(response_headers_with_200)) .returning(Some("200"))
.expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None)