add support for claude

This commit is contained in:
Adil Hafeez 2025-05-19 13:50:51 -07:00
parent a3352254fe
commit adda6aaa57
No known key found for this signature in database
GPG key ID: 9B18EF7691369645
6 changed files with 128 additions and 52 deletions

View file

@ -330,6 +330,7 @@ impl TryFrom<&str> for ChatCompletionStreamResponseServerEvents {
let response_chunks: VecDeque<ChatCompletionStreamResponse> = value
.lines()
.filter(|line| line.starts_with("data: "))
.filter(|line| !line.starts_with(r#"data: {"type": "ping"}"#))
.map(|line| line.get(6..).unwrap())
.filter(|data_chunk| *data_chunk != "[DONE]")
.map(serde_json::from_str::<ChatCompletionStreamResponse>)
@ -677,4 +678,37 @@ data: [DONE]
"Hello! How can I assist you today?"
);
}
#[test]
fn stream_chunk_parse_claude() {
const CHUNK_RESPONSE: &str = r#"data: {"id":"msg_01DZDMxYSgq8aPQxMQoBv6Kb","choices":[{"index":0,"delta":{"role":"assistant"}}],"created":1747685264,"model":"claude-3-7-sonnet-latest","object":"chat.completion.chunk"}
data: {"type": "ping"}
data: {"id":"msg_01DZDMxYSgq8aPQxMQoBv6Kb","choices":[{"index":0,"delta":{"content":"Hello!"}}],"created":1747685264,"model":"claude-3-7-sonnet-latest","object":"chat.completion.chunk"}
data: {"id":"msg_01DZDMxYSgq8aPQxMQoBv6Kb","choices":[{"index":0,"delta":{"content":" How can I assist you today? Whether"}}],"created":1747685264,"model":"claude-3-7-sonnet-latest","object":"chat.completion.chunk"}
data: {"id":"msg_01DZDMxYSgq8aPQxMQoBv6Kb","choices":[{"index":0,"delta":{"content":" you have a question, need information"}}],"created":1747685264,"model":"claude-3-7-sonnet-latest","object":"chat.completion.chunk"}
data: {"id":"msg_01DZDMxYSgq8aPQxMQoBv6Kb","choices":[{"index":0,"delta":{"content":", or just want to chat about"}}],"created":1747685264,"model":"claude-3-7-sonnet-latest","object":"chat.completion.chunk"}
data: {"id":"msg_01DZDMxYSgq8aPQxMQoBv6Kb","choices":[{"index":0,"delta":{"content":" something, I'm here to help. What woul"}}],"created":1747685264,"model":"claude-3-7-sonnet-latest","object":"chat.completion.chunk"}
data: {"id":"msg_01DZDMxYSgq8aPQxMQoBv6Kb","choices":[{"index":0,"delta":{"content":"d you like to talk about?"}}],"created":1747685264,"model":"claude-3-7-sonnet-latest","object":"chat.completion.chunk"}
data: {"id":"msg_01DZDMxYSgq8aPQxMQoBv6Kb","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"created":1747685264,"model":"claude-3-7-sonnet-latest","object":"chat.completion.chunk"}
data: [DONE]
"#;
let sever_events: ChatCompletionStreamResponseServerEvents =
ChatCompletionStreamResponseServerEvents::try_from(CHUNK_RESPONSE).unwrap();
assert_eq!(sever_events.events.len(), 8);
assert_eq!(
sever_events.to_string(),
"Hello! How can I assist you today? Whether you have a question, need information, or just want to chat about something, I'm here to help. What would you like to talk about?"
);
}
}

View file

@ -143,19 +143,28 @@ pub struct EmbeddingProviver {
pub model: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum LlmProviderType {
#[serde(rename = "openai")]
OpenAI,
#[serde(rename = "claude")]
Claude,
#[serde(rename = "deepseek")]
Deepseek,
#[serde(rename = "groq")]
Groq,
#[serde(rename = "mistral")]
Mistral,
#[serde(rename = "openai")]
OpenAI,
}
impl Display for LlmProviderType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
LlmProviderType::OpenAI => write!(f, "openai"),
LlmProviderType::Claude => write!(f, "claude"),
LlmProviderType::Deepseek => write!(f, "deepseek"),
LlmProviderType::Groq => write!(f, "groq"),
LlmProviderType::Mistral => write!(f, "mistral"),
LlmProviderType::OpenAI => write!(f, "openai"),
}
}
}
@ -175,6 +184,23 @@ pub struct LlmProvider {
pub usage: Option<String>,
}
impl Default for LlmProvider {
fn default() -> Self {
Self {
name: "openai".to_string(),
provider_interface: LlmProviderType::OpenAI,
access_key: None,
model: None,
default: Some(true),
stream: Some(false),
endpoint: None,
port: None,
rate_limits: None,
usage: None,
}
}
}
impl Display for LlmProvider {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.name)

View file

@ -89,15 +89,7 @@ impl StreamContext {
provider_hint,
));
// Check if we need to modify the path based on the provider's base_url
let needs_openai_prefix = self
.llm_provider
.as_ref()
.and_then(|provider| provider.endpoint.as_ref())
.map(|url| url.contains("api.groq.com"))
.unwrap_or(false);
if needs_openai_prefix {
if self.llm_provider.as_ref().unwrap().provider_interface == LlmProviderType::Groq {
if let Some(path) = self.get_http_request_header(":path") {
if path.starts_with("/v1/") {
let new_path = format!("/openai{}", path);
@ -221,14 +213,7 @@ impl HttpContext for StreamContext {
self.llm_provider = Some(Rc::new(LlmProvider {
name: routing_header_value.to_string(),
provider_interface: LlmProviderType::OpenAI,
access_key: None,
endpoint: None,
model: None,
default: None,
stream: None,
port: None,
rate_limits: None,
usage: None,
..Default::default()
}));
} else {
self.select_llm_provider();
@ -539,6 +524,9 @@ impl HttpContext for StreamContext {
}
streaming_chunk
} else {
if body_size == 0 {
return Action::Continue;
}
debug!("non streaming response bytes read: 0:{}", body_size);
match self.get_http_response_body(0, body_size) {
Some(body) => body,