P0 deep-dive: core, providers, agent, memory, TUI, setup wizard
DEEP DIVE — Agent Loop (gyandoot-agent): - Real conversation loop with interactive (run()) and non-interactive (run_conversation()) modes - Retry logic with exponential backoff and error classification (RateLimit, Auth, Timeout, ServerError, ContextOverflow) - Context window management — token estimation, 75% threshold compression - 7 built-in tools: Read, Write, Edit, Bash, Grep, Find, Ls - BashTool with destructive-command security blocks - ContextCompressor — prune tool results, summarize middle messages DEEP DIVE — Memory (gyandoot-memory): - Real MemPalace with ~/.gyandoot/palace/wings/<wing>/ directory tree - Daily transcript files (YYYY-MM-DD_transcript.md) with append-only format - L0/L1 identity persistence (< 157 tokens for L0) - Session digest to drawer.md - File persistence support for transcripts and identity store DEEP DIVE — Core (gyandoot-core): - GyandootConfig with TOML deserialization for all sections - ConfigLoader with env > profile > file > default chain - Rich types: Message, ToolCall, ToolResult, SessionId - Comprehensive error types with thiserror derive DEEP DIVE — Providers (gyandoot-providers): - Provider trait with chat/chat_stream/models - OpenAI provider with streaming SSE and function calling - Anthropic provider with Messages API + tool use - Ollama provider with merge_system/no_tools auto-detect - SseParser, StreamEvent, TokenStream for SSE parsing - ProviderFactory for creation from config NEW — TUI crate (gyandoot-tui): - Banner: Kapil-approved robot head ASCII art with ANSI colors - Setup wizard: 7 providers, interactive config, connection test - Terminal helpers: spinner, confirm, text input, selection menu - Ratatui TUI: keyboard input, commands, status display All 12 crates + CLI binary compile with 0 errors.
This commit is contained in:
@@ -10,7 +10,7 @@ gyandoot-core = { path = "../gyandoot-core" }
|
||||
tokio.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
reqwest.workspace = true
|
||||
reqwest = { workspace = true, features = ["json", "stream"] }
|
||||
thiserror.workspace = true
|
||||
anyhow.workspace = true
|
||||
tracing.workspace = true
|
||||
|
||||
@@ -1,17 +1,28 @@
|
||||
//! Anthropic API provider implementation.
|
||||
//!
|
||||
//! Full Anthropic Messages API (`/v1/messages`) with streaming SSE support,
|
||||
//! tool use via the `tools` parameter, proper `x-api-key` + `anthropic-version`
|
||||
//! header auth, and structured error handling.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use gyandoot_core::{ConversationContext, MessageRole, ProviderResult};
|
||||
use futures::StreamExt;
|
||||
use gyandoot_core::{ConversationContext, ProviderError, ProviderResult, ToolCall};
|
||||
use reqwest::Client;
|
||||
use tracing::debug;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::{debug, trace, warn};
|
||||
|
||||
use crate::streaming::StreamEvent;
|
||||
use crate::{Provider, ProviderResponse};
|
||||
|
||||
/// Provider for the Anthropic (Claude) API.
|
||||
const ANTHROPIC_VERSION: &str = "2023-06-01";
|
||||
|
||||
/// Provider for the Anthropic (Claude) Messages API.
|
||||
pub struct AnthropicProvider {
|
||||
api_key: String,
|
||||
model: String,
|
||||
client: Client,
|
||||
request_timeout: Duration,
|
||||
}
|
||||
|
||||
impl AnthropicProvider {
|
||||
@@ -20,39 +31,158 @@ impl AnthropicProvider {
|
||||
Self {
|
||||
api_key,
|
||||
model,
|
||||
client: Client::new(),
|
||||
client: Client::builder()
|
||||
.timeout(Duration::from_secs(120))
|
||||
.build()
|
||||
.expect("Failed to create HTTP client"),
|
||||
request_timeout: Duration::from_secs(60),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set a custom request timeout.
|
||||
pub fn with_timeout(mut self, timeout: Duration) -> Self {
|
||||
self.request_timeout = timeout;
|
||||
self
|
||||
}
|
||||
|
||||
/// Extract the system prompt from context (Anthropic uses a separate field).
|
||||
fn extract_system(&self, context: &ConversationContext) -> Option<String> {
|
||||
let system_msg = context.messages.iter().find(|m| m.is_system());
|
||||
system_msg.map(|m| m.content.clone())
|
||||
}
|
||||
|
||||
/// Build the messages array in Anthropic's format.
|
||||
///
|
||||
/// Anthropic expects alternating `user` and `assistant` messages.
|
||||
/// System messages are extracted separately. Tool results are sent
|
||||
/// as `user` messages with `content` blocks containing `tool_result`.
|
||||
fn build_messages(&self, context: &ConversationContext) -> Vec<serde_json::Value> {
|
||||
let mut result = Vec::new();
|
||||
|
||||
for m in &context.messages {
|
||||
if m.is_system() {
|
||||
continue; // handled separately
|
||||
}
|
||||
|
||||
let serialized = match m.role.to_string().as_str() {
|
||||
"tool" => {
|
||||
// Tool results: send as user with content blocks
|
||||
serde_json::json!({
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": m.tool_result.as_ref().map(|_| "tool_result"),
|
||||
"content": m.content
|
||||
}
|
||||
]
|
||||
})
|
||||
}
|
||||
role_str => {
|
||||
let mut msg = serde_json::json!({
|
||||
"role": role_str,
|
||||
"content": m.content,
|
||||
});
|
||||
|
||||
// Attach tool_use blocks for assistant tool calls
|
||||
if let Some(calls) = &m.tool_calls {
|
||||
let blocks: Vec<serde_json::Value> = calls
|
||||
.iter()
|
||||
.map(|tc| {
|
||||
serde_json::json!({
|
||||
"type": "tool_use",
|
||||
"id": tc.id,
|
||||
"name": tc.name,
|
||||
"input": tc.arguments,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
msg["content"] = serde_json::to_value(blocks).unwrap_or_default();
|
||||
}
|
||||
|
||||
msg
|
||||
}
|
||||
};
|
||||
|
||||
result.push(serialized);
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Build the tool definitions in Anthropic's format.
|
||||
fn build_tools(&self, context: &ConversationContext) -> Vec<serde_json::Value> {
|
||||
context
|
||||
.messages
|
||||
.tools
|
||||
.iter()
|
||||
.filter(|m| m.role != MessageRole::System)
|
||||
.map(|m| {
|
||||
let role = match m.role {
|
||||
MessageRole::User => "user",
|
||||
MessageRole::Assistant => "assistant",
|
||||
MessageRole::Tool => "user",
|
||||
MessageRole::System => "user", // filtered above
|
||||
};
|
||||
.map(|td| {
|
||||
serde_json::json!({
|
||||
"role": role,
|
||||
"content": m.content,
|
||||
"name": td.name,
|
||||
"description": td.description,
|
||||
"input_schema": td.parameters,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Extract the system prompt from context (Anthropic uses a separate field).
|
||||
fn extract_system(&self, context: &ConversationContext) -> String {
|
||||
context
|
||||
.messages
|
||||
.iter()
|
||||
.find(|m| m.role == MessageRole::System)
|
||||
.map(|m| m.content.clone())
|
||||
.unwrap_or_default()
|
||||
/// Send the non-streaming request.
|
||||
async fn send_request(&self, body: &serde_json::Value) -> ProviderResult<reqwest::Response> {
|
||||
let response = self
|
||||
.client
|
||||
.post("https://api.anthropic.com/v1/messages")
|
||||
.header("x-api-key", &self.api_key)
|
||||
.header("anthropic-version", ANTHROPIC_VERSION)
|
||||
.header("content-type", "application/json")
|
||||
.json(body)
|
||||
.timeout(self.request_timeout)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let error_text = response.text().await.unwrap_or_default();
|
||||
return Err(ProviderError::from_response(status.as_u16(), &error_text));
|
||||
}
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// Parse a non-streaming response.
|
||||
fn parse_response(&self, data: &serde_json::Value) -> ProviderResult<ProviderResponse> {
|
||||
// Extract content blocks
|
||||
let content_blocks = data["content"].as_array().cloned().unwrap_or_default();
|
||||
let mut text_parts = Vec::new();
|
||||
let mut tool_calls = Vec::new();
|
||||
|
||||
for block in &content_blocks {
|
||||
match block["type"].as_str() {
|
||||
Some("text") => {
|
||||
if let Some(text) = block["text"].as_str() {
|
||||
text_parts.push(text.to_string());
|
||||
}
|
||||
}
|
||||
Some("tool_use") => {
|
||||
let id = block["id"].as_str().unwrap_or_default().to_string();
|
||||
let name = block["name"].as_str().unwrap_or_default().to_string();
|
||||
let arguments = block["input"].clone();
|
||||
tool_calls.push(ToolCall::new(id, name, arguments));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
let content = text_parts.join("");
|
||||
let finish_reason = data["stop_reason"].as_str().unwrap_or("end_turn").to_string();
|
||||
|
||||
Ok(ProviderResponse {
|
||||
content,
|
||||
tool_calls: if tool_calls.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(tool_calls)
|
||||
},
|
||||
finish_reason,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -63,6 +193,7 @@ impl Provider for AnthropicProvider {
|
||||
|
||||
let system = self.extract_system(&context);
|
||||
let messages = self.build_messages(&context);
|
||||
let tools = self.build_tools(&context);
|
||||
|
||||
let mut body = serde_json::json!({
|
||||
"model": self.model,
|
||||
@@ -70,40 +201,18 @@ impl Provider for AnthropicProvider {
|
||||
"max_tokens": context.max_tokens,
|
||||
});
|
||||
|
||||
if !system.is_empty() {
|
||||
body["system"] = serde_json::Value::String(system);
|
||||
if let Some(sys) = &system {
|
||||
body["system"] = serde_json::Value::String(sys.clone());
|
||||
}
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post("https://api.anthropic.com/v1/messages")
|
||||
.header("x-api-key", &self.api_key)
|
||||
.header("anthropic-version", "2023-06-01")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let error_text = response.text().await.unwrap_or_default();
|
||||
return Err(gyandoot_core::ProviderError::Api(format!(
|
||||
"Anthropic API error ({}): {}",
|
||||
status, error_text
|
||||
)));
|
||||
if !tools.is_empty() {
|
||||
body["tools"] = serde_json::Value::Array(tools);
|
||||
}
|
||||
|
||||
let response = self.send_request(&body).await?;
|
||||
let data: serde_json::Value = response.json().await?;
|
||||
let content = data["content"][0]["text"]
|
||||
.as_str()
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let finish_reason = data["stop_reason"].as_str().unwrap_or("end_turn").to_string();
|
||||
|
||||
Ok(ProviderResponse {
|
||||
content,
|
||||
tool_calls: None,
|
||||
finish_reason,
|
||||
})
|
||||
Ok(self.parse_response(&data)?)
|
||||
}
|
||||
|
||||
async fn chat_stream(
|
||||
@@ -111,12 +220,238 @@ impl Provider for AnthropicProvider {
|
||||
context: ConversationContext,
|
||||
) -> ProviderResult<Box<dyn tokio_stream::Stream<Item = ProviderResult<String>> + Send + Unpin>>
|
||||
{
|
||||
// Fall back to non-streaming for now
|
||||
let result = self.chat(context).await?;
|
||||
let stream = tokio_stream::StreamExt::chain(
|
||||
tokio_stream::once(Ok(result.content)),
|
||||
tokio_stream::empty(),
|
||||
);
|
||||
debug!("Anthropic streaming chat (model: {})", self.model);
|
||||
|
||||
let system = self.extract_system(&context);
|
||||
let messages = self.build_messages(&context);
|
||||
let tools = self.build_tools(&context);
|
||||
|
||||
let mut body = serde_json::json!({
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"max_tokens": context.max_tokens,
|
||||
"stream": true,
|
||||
});
|
||||
|
||||
if let Some(sys) = &system {
|
||||
body["system"] = serde_json::Value::String(sys.clone());
|
||||
}
|
||||
|
||||
if !tools.is_empty() {
|
||||
body["tools"] = serde_json::Value::Array(tools);
|
||||
}
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post("https://api.anthropic.com/v1/messages")
|
||||
.header("x-api-key", &self.api_key)
|
||||
.header("anthropic-version", ANTHROPIC_VERSION)
|
||||
.header("content-type", "application/json")
|
||||
.json(&body)
|
||||
.timeout(self.request_timeout)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let error_text = response.text().await.unwrap_or_default();
|
||||
return Err(ProviderError::from_response(status.as_u16(), &error_text));
|
||||
}
|
||||
|
||||
let (tx, rx) = mpsc::channel::<ProviderResult<String>>(256);
|
||||
let mut byte_stream = response.bytes_stream();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut parser = crate::streaming::SseParser::new(50);
|
||||
let mut buffer = Vec::new();
|
||||
|
||||
loop {
|
||||
match tokio::time::timeout(Duration::from_secs(30), byte_stream.next()).await {
|
||||
Ok(Some(Ok(chunk))) => {
|
||||
buffer.extend_from_slice(&chunk);
|
||||
let buffer_str = String::from_utf8_lossy(&buffer);
|
||||
|
||||
for event in parser.feed_chunk(buffer_str.as_ref()) {
|
||||
match event {
|
||||
StreamEvent::Chunk(text) => {
|
||||
// Anthropic SSE includes content_block_delta events
|
||||
// with "text" deltas. Parse them from the JSON.
|
||||
if let Ok(val) =
|
||||
serde_json::from_str::<serde_json::Value>(&text)
|
||||
{
|
||||
if let Some(delta) = val["delta"]["text"].as_str() {
|
||||
trace!("Anthropic delta: {delta}");
|
||||
if tx.send(Ok(delta.to_string())).await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
StreamEvent::Done => {
|
||||
debug!("Anthropic stream done");
|
||||
return;
|
||||
}
|
||||
StreamEvent::Error(msg) => {
|
||||
warn!("Anthropic SSE error: {msg}");
|
||||
let _ = tx.send(Err(ProviderError::Stream(msg))).await;
|
||||
return;
|
||||
}
|
||||
StreamEvent::ToolCall(_tc) => {
|
||||
// Anthropic tool calls appear in the final message event,
|
||||
// not as streaming deltas. Handled in non-streaming path.
|
||||
}
|
||||
}
|
||||
}
|
||||
buffer.drain(0..parser.consumed());
|
||||
}
|
||||
Ok(Some(Err(e))) => {
|
||||
let _ = tx.send(Err(ProviderError::Request(e))).await;
|
||||
return;
|
||||
}
|
||||
Ok(None) => return,
|
||||
Err(_) => {
|
||||
let _ = tx
|
||||
.send(Err(ProviderError::Timeout { duration: 30 }))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let stream = tokio_stream::wrappers::ReceiverStream::new(rx);
|
||||
Ok(Box::new(stream))
|
||||
}
|
||||
|
||||
async fn models(&self) -> ProviderResult<Vec<String>> {
|
||||
// Anthropic doesn't have a public models endpoint.
|
||||
// Return known models.
|
||||
Ok(vec![
|
||||
"claude-3-5-sonnet-20241022".to_string(),
|
||||
"claude-3-5-haiku-20241022".to_string(),
|
||||
"claude-3-opus-20240229".to_string(),
|
||||
])
|
||||
}
|
||||
|
||||
fn supports_tools(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use gyandoot_core::{Message, Role, ToolDefinition};
|
||||
|
||||
fn test_provider() -> AnthropicProvider {
|
||||
AnthropicProvider::new("sk-ant-test".to_string(), "claude-3-5-sonnet-20241022".to_string())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_system() {
|
||||
let provider = test_provider();
|
||||
let context = ConversationContext {
|
||||
messages: vec![
|
||||
Message::new(Role::System, "You are Claude."),
|
||||
Message::new(Role::User, "Hello"),
|
||||
],
|
||||
system_prompt: "You are Claude.".to_string(),
|
||||
tools: vec![],
|
||||
max_tokens: 4096,
|
||||
};
|
||||
let system = provider.extract_system(&context);
|
||||
assert_eq!(system, Some("You are Claude.".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_messages_filters_system() {
|
||||
let provider = test_provider();
|
||||
let context = ConversationContext {
|
||||
messages: vec![
|
||||
Message::new(Role::System, "System prompt"),
|
||||
Message::new(Role::User, "Hello"),
|
||||
Message::new(Role::Assistant, "Hi there!"),
|
||||
],
|
||||
system_prompt: String::new(),
|
||||
tools: vec![],
|
||||
max_tokens: 4096,
|
||||
};
|
||||
let messages = provider.build_messages(&context);
|
||||
assert_eq!(messages.len(), 2);
|
||||
assert_eq!(messages[0]["role"], "user");
|
||||
assert_eq!(messages[1]["role"], "assistant");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_tools() {
|
||||
let provider = test_provider();
|
||||
let context = ConversationContext {
|
||||
messages: vec![Message::new(Role::User, "Weather?")],
|
||||
system_prompt: String::new(),
|
||||
tools: vec![ToolDefinition::new(
|
||||
"get_weather",
|
||||
"Get weather",
|
||||
serde_json::json!({"type": "object"}),
|
||||
)],
|
||||
max_tokens: 4096,
|
||||
};
|
||||
let tools = provider.build_tools(&context);
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0]["name"], "get_weather");
|
||||
assert!(tools[0].get("input_schema").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_response_text_only() {
|
||||
let provider = test_provider();
|
||||
let data = serde_json::json!({
|
||||
"id": "msg_123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Hello! I'm Claude."}
|
||||
],
|
||||
"stop_reason": "end_turn",
|
||||
"model": "claude-3-5-sonnet-20241022",
|
||||
"usage": {
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 5
|
||||
}
|
||||
});
|
||||
|
||||
let response = provider.parse_response(&data).unwrap();
|
||||
assert_eq!(response.content, "Hello! I'm Claude.");
|
||||
assert_eq!(response.finish_reason, "end_turn");
|
||||
assert!(response.tool_calls.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_response_with_tool_use() {
|
||||
let provider = test_provider();
|
||||
let data = serde_json::json!({
|
||||
"id": "msg_456",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Let me check the weather."},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_123",
|
||||
"name": "get_weather",
|
||||
"input": {"city": "NYC"}
|
||||
}
|
||||
],
|
||||
"stop_reason": "tool_use",
|
||||
"model": "claude-3-5-sonnet-20241022"
|
||||
});
|
||||
|
||||
let response = provider.parse_response(&data).unwrap();
|
||||
assert_eq!(response.content, "Let me check the weather.");
|
||||
assert_eq!(response.finish_reason, "tool_use");
|
||||
assert!(response.has_tool_calls());
|
||||
let calls = response.tool_calls.unwrap();
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].id, "toolu_123");
|
||||
assert_eq!(calls[0].name, "get_weather");
|
||||
}
|
||||
}
|
||||
@@ -1,34 +1,327 @@
|
||||
//! GyanDoot Providers — LLM provider abstraction layer.
|
||||
//!
|
||||
//! Defines the `Provider` trait and implementations for OpenAI, Anthropic,
|
||||
//! and Ollama. Each provider handles authentication, request formatting,
|
||||
//! response parsing, and streaming.
|
||||
//! Defines the [`Provider`] trait and implementations for OpenAI, Anthropic,
|
||||
//! Ollama, and other providers. Each implementation handles authentication,
|
||||
//! request formatting, response parsing, and streaming.
|
||||
//!
|
||||
//! The [`ProviderFactory`] creates provider instances from
|
||||
//! [`gyandoot_core::ProviderConfig`].
|
||||
|
||||
use async_trait::async_trait;
|
||||
use gyandoot_core::{ConversationContext, ProviderResult};
|
||||
use gyandoot_core::{ConversationContext, ProviderResult, ToolCall};
|
||||
|
||||
pub mod openai;
|
||||
pub mod anthropic;
|
||||
pub mod ollama;
|
||||
pub mod streaming;
|
||||
|
||||
// ─── Re-exports ───────────────────────────────────────────────────────────
|
||||
|
||||
pub use openai::OpenAIProvider;
|
||||
pub use anthropic::AnthropicProvider;
|
||||
pub use ollama::OllamaProvider;
|
||||
pub use streaming::{SseParser, StreamEvent, TokenStream, collect_stream};
|
||||
|
||||
// ─── ProviderResponse ─────────────────────────────────────────────────────
|
||||
|
||||
/// A provider response containing the generated text and any tool calls.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProviderResponse {
|
||||
pub content: String,
|
||||
pub tool_calls: Option<Vec<gyandoot_core::ToolCall>>,
|
||||
pub tool_calls: Option<Vec<ToolCall>>,
|
||||
pub finish_reason: String,
|
||||
}
|
||||
|
||||
impl ProviderResponse {
|
||||
/// Create a new provider response with text content only.
|
||||
pub fn new(content: impl Into<String>) -> Self {
|
||||
Self {
|
||||
content: content.into(),
|
||||
tool_calls: None,
|
||||
finish_reason: "stop".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a provider response with tool calls.
|
||||
pub fn with_tool_calls(
|
||||
content: impl Into<String>,
|
||||
calls: Vec<ToolCall>,
|
||||
finish_reason: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
content: content.into(),
|
||||
tool_calls: Some(calls),
|
||||
finish_reason: finish_reason.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Return true if this response contains tool calls.
|
||||
pub fn has_tool_calls(&self) -> bool {
|
||||
self.tool_calls.as_ref().map_or(false, |c| !c.is_empty())
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Provider trait ───────────────────────────────────────────────────────
|
||||
|
||||
/// The trait that all LLM providers must implement.
|
||||
#[async_trait]
|
||||
pub trait Provider: Send + Sync {
|
||||
/// Send a chat completion request and return the response.
|
||||
/// Send a chat completion request and return the full response.
|
||||
async fn chat(&self, context: ConversationContext) -> ProviderResult<ProviderResponse>;
|
||||
|
||||
/// Send a chat completion request with streaming response.
|
||||
///
|
||||
/// Each item in the stream represents a text delta. Use
|
||||
/// [`streaming::collect_stream`] to aggregate the full text.
|
||||
async fn chat_stream(
|
||||
&self,
|
||||
context: ConversationContext,
|
||||
) -> ProviderResult<Box<dyn tokio_stream::Stream<Item = ProviderResult<String>> + Send + Unpin>>;
|
||||
|
||||
/// Return the list of models available for this provider.
|
||||
async fn models(&self) -> ProviderResult<Vec<String>> {
|
||||
// Default: return an empty list; providers can override.
|
||||
let _ = self;
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
/// Return whether this provider supports tool/function calling.
|
||||
fn supports_tools(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
// ─── ProviderKind (runtime enum) ──────────────────────────────────────────
|
||||
|
||||
/// Runtime provider kind for factory dispatch.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum ProviderKind {
|
||||
OpenAI,
|
||||
Anthropic,
|
||||
Ollama,
|
||||
Gemini,
|
||||
DeepSeek,
|
||||
OpenRouter,
|
||||
}
|
||||
|
||||
impl ProviderKind {
|
||||
/// Parse from a string, case-insensitive.
|
||||
pub fn from_str(s: &str) -> Option<Self> {
|
||||
match s.to_lowercase().as_str() {
|
||||
"openai" => Some(ProviderKind::OpenAI),
|
||||
"anthropic" => Some(ProviderKind::Anthropic),
|
||||
"ollama" => Some(ProviderKind::Ollama),
|
||||
"gemini" => Some(ProviderKind::Gemini),
|
||||
"deepseek" => Some(ProviderKind::DeepSeek),
|
||||
"openrouter" => Some(ProviderKind::OpenRouter),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the canonical string name.
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
ProviderKind::OpenAI => "openai",
|
||||
ProviderKind::Anthropic => "anthropic",
|
||||
ProviderKind::Ollama => "ollama",
|
||||
ProviderKind::Gemini => "gemini",
|
||||
ProviderKind::DeepSeek => "deepseek",
|
||||
ProviderKind::OpenRouter => "openrouter",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ProviderKind {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}", self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
// ─── ProviderFactory ───────────────────────────────────────────────────────
|
||||
|
||||
/// Factory for creating provider instances from configuration.
|
||||
pub struct ProviderFactory;
|
||||
|
||||
impl ProviderFactory {
|
||||
/// Create a provider from a [`gyandoot_core::ProviderRootConfig`].
|
||||
pub fn create(config: &gyandoot_core::ProviderRootConfig) -> ProviderResult<Box<dyn Provider>> {
|
||||
let kind = ProviderKind::from_str(&config.default)
|
||||
.ok_or_else(|| gyandoot_core::ProviderError::Api {
|
||||
status: 0,
|
||||
message: format!("Unknown provider name: {}", config.default),
|
||||
})?;
|
||||
|
||||
Self::create_from_kind(kind, config)
|
||||
}
|
||||
|
||||
/// Create a provider from a [`ProviderKind`] and config.
|
||||
pub fn create_from_kind(
|
||||
kind: ProviderKind,
|
||||
config: &gyandoot_core::ProviderRootConfig,
|
||||
) -> ProviderResult<Box<dyn Provider>> {
|
||||
match kind {
|
||||
ProviderKind::OpenAI | ProviderKind::DeepSeek | ProviderKind::OpenRouter => {
|
||||
let cfg = config.openai.as_ref().ok_or_else(|| gyandoot_core::ProviderError::Api {
|
||||
status: 0,
|
||||
message: "OpenAI config not found".to_string(),
|
||||
})?;
|
||||
Ok(Box::new(openai::OpenAIProvider::new(
|
||||
cfg.api_key.clone(),
|
||||
cfg.model.clone(),
|
||||
cfg.base_url.clone(),
|
||||
)))
|
||||
}
|
||||
ProviderKind::Anthropic => {
|
||||
let cfg = config.anthropic.as_ref().ok_or_else(|| gyandoot_core::ProviderError::Api {
|
||||
status: 0,
|
||||
message: "Anthropic config not found".to_string(),
|
||||
})?;
|
||||
Ok(Box::new(anthropic::AnthropicProvider::new(
|
||||
cfg.api_key.clone(),
|
||||
cfg.model.clone(),
|
||||
)))
|
||||
}
|
||||
ProviderKind::Ollama => {
|
||||
let cfg = config.ollama.as_ref().ok_or_else(|| gyandoot_core::ProviderError::Api {
|
||||
status: 0,
|
||||
message: "Ollama config not found".to_string(),
|
||||
})?;
|
||||
Ok(Box::new(ollama::OllamaProvider::new(
|
||||
cfg.base_url.clone(),
|
||||
cfg.model.clone(),
|
||||
)))
|
||||
}
|
||||
ProviderKind::Gemini => {
|
||||
let cfg = config.openai.as_ref().ok_or_else(|| gyandoot_core::ProviderError::Api {
|
||||
status: 0,
|
||||
message: "Gemini config not found (use openai section)".to_string(),
|
||||
})?;
|
||||
// Gemini uses the OpenAI-compatible endpoint
|
||||
Ok(Box::new(openai::OpenAIProvider::new(
|
||||
cfg.api_key.clone(),
|
||||
cfg.model.clone(),
|
||||
cfg.base_url.clone(),
|
||||
)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a provider from just a kind name string and minimal config.
|
||||
pub fn from_name(
|
||||
name: &str,
|
||||
api_key: &str,
|
||||
model: &str,
|
||||
) -> ProviderResult<Box<dyn Provider>> {
|
||||
let kind = ProviderKind::from_str(name)
|
||||
.ok_or_else(|| gyandoot_core::ProviderError::Api {
|
||||
status: 0,
|
||||
message: format!("Unknown provider name: {name}"),
|
||||
})?;
|
||||
|
||||
let config = gyandoot_core::ProviderRootConfig {
|
||||
default: name.to_string(),
|
||||
openai: if matches!(kind, ProviderKind::OpenAI | ProviderKind::DeepSeek | ProviderKind::OpenRouter | ProviderKind::Gemini) {
|
||||
Some(gyandoot_core::OpenAIConfig {
|
||||
api_key: api_key.to_string(),
|
||||
model: model.to_string(),
|
||||
base_url: Some(default_base_url(&kind)),
|
||||
max_retries: 3,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
},
|
||||
anthropic: if matches!(kind, ProviderKind::Anthropic) {
|
||||
Some(gyandoot_core::AnthropicConfig {
|
||||
api_key: api_key.to_string(),
|
||||
model: model.to_string(),
|
||||
max_retries: 3,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
},
|
||||
ollama: if matches!(kind, ProviderKind::Ollama) {
|
||||
Some(gyandoot_core::OllamaConfig {
|
||||
base_url: default_base_url(&kind),
|
||||
model: model.to_string(),
|
||||
max_retries: 2,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
},
|
||||
};
|
||||
|
||||
Self::create_from_kind(kind, &config)
|
||||
}
|
||||
}
|
||||
|
||||
/// Default base URL for each provider kind.
|
||||
fn default_base_url(kind: &ProviderKind) -> String {
|
||||
match kind {
|
||||
ProviderKind::OpenAI => "https://api.openai.com/v1".to_string(),
|
||||
ProviderKind::Anthropic => "https://api.anthropic.com".to_string(),
|
||||
ProviderKind::Ollama => "http://localhost:11434".to_string(),
|
||||
ProviderKind::Gemini => "https://generativelanguage.googleapis.com/v1".to_string(),
|
||||
ProviderKind::DeepSeek => "https://api.deepseek.com/v1".to_string(),
|
||||
ProviderKind::OpenRouter => "https://openrouter.ai/api/v1".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Tests ────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_provider_kind_roundtrip() {
|
||||
for kind in &[
|
||||
ProviderKind::OpenAI,
|
||||
ProviderKind::Anthropic,
|
||||
ProviderKind::Ollama,
|
||||
ProviderKind::Gemini,
|
||||
ProviderKind::DeepSeek,
|
||||
ProviderKind::OpenRouter,
|
||||
] {
|
||||
let s = kind.to_string();
|
||||
let parsed = ProviderKind::from_str(&s).unwrap();
|
||||
assert_eq!(*kind, parsed);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_provider_response_has_tool_calls() {
|
||||
let r = ProviderResponse::new("hello");
|
||||
assert!(!r.has_tool_calls());
|
||||
|
||||
let r = ProviderResponse::with_tool_calls(
|
||||
"",
|
||||
vec![ToolCall::new("1", "test", serde_json::Value::Null)],
|
||||
"tool_calls",
|
||||
);
|
||||
assert!(r.has_tool_calls());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_base_urls() {
|
||||
assert!(default_base_url(&ProviderKind::OpenAI).contains("openai.com"));
|
||||
assert!(default_base_url(&ProviderKind::Anthropic).contains("anthropic.com"));
|
||||
assert!(default_base_url(&ProviderKind::Ollama).contains("localhost"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_factory_from_name() {
|
||||
let provider = ProviderFactory::from_name("openai", "sk-test", "gpt-4o").unwrap();
|
||||
assert!(provider.supports_tools());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_factory_invalid_name() {
|
||||
let result = ProviderFactory::from_name("nonexistent", "", "");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_provider_kind_unknown_string() {
|
||||
assert_eq!(ProviderKind::from_str("unknown"), None);
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,7 @@
|
||||
//! Ollama API provider implementation.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use gyandoot_core::{ConversationContext, MessageRole, ProviderResult};
|
||||
use gyandoot_core::{ConversationContext, ProviderResult};
|
||||
use reqwest::Client;
|
||||
use tracing::debug;
|
||||
|
||||
@@ -59,10 +59,13 @@ impl Provider for OllamaProvider {
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(gyandoot_core::ProviderError::Api(format!(
|
||||
"Ollama API error: {}",
|
||||
response.status()
|
||||
)));
|
||||
return Err(gyandoot_core::ProviderError::Api {
|
||||
status: response.status().as_u16(),
|
||||
message: format!(
|
||||
"Ollama API error: {}",
|
||||
response.status()
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
let data: serde_json::Value = response.json().await?;
|
||||
|
||||
@@ -1,31 +1,56 @@
|
||||
//! OpenAI API provider implementation.
|
||||
//!
|
||||
//! Full OpenAI-compatible chat completions with function calling, streaming
|
||||
//! SSE parsing, and error handling for auth, rate limits, and timeouts.
|
||||
//! Works with any OpenAI-compatible endpoint (OpenAI, DeepSeek, OpenRouter,
|
||||
//! Gemini, etc.).
|
||||
|
||||
use async_trait::async_trait;
|
||||
use gyandoot_core::{ConversationContext, MessageRole, ProviderResult};
|
||||
use futures::StreamExt;
|
||||
use gyandoot_core::{ConversationContext, ProviderError, ProviderResult, ToolCall};
|
||||
use reqwest::Client;
|
||||
use tracing::debug;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::{debug, trace, warn};
|
||||
|
||||
use crate::streaming::StreamEvent;
|
||||
use crate::{Provider, ProviderResponse};
|
||||
|
||||
/// Provider for the OpenAI chat completion API.
|
||||
/// Provider for OpenAI-compatible chat completion APIs.
|
||||
pub struct OpenAIProvider {
|
||||
api_key: String,
|
||||
model: String,
|
||||
base_url: String,
|
||||
client: Client,
|
||||
request_timeout: Duration,
|
||||
}
|
||||
|
||||
impl OpenAIProvider {
|
||||
/// Create a new OpenAI provider.
|
||||
///
|
||||
/// - `api_key`: Bearer token for authentication.
|
||||
/// - `model`: Model identifier (e.g. "gpt-4o").
|
||||
/// - `base_url`: Base URL of the API (e.g. "https://api.openai.com/v1").
|
||||
/// Defaults to OpenAI's endpoint if `None`.
|
||||
pub fn new(api_key: String, model: String, base_url: Option<String>) -> Self {
|
||||
Self {
|
||||
api_key,
|
||||
model,
|
||||
base_url: base_url.unwrap_or_else(|| "https://api.openai.com/v1".to_string()),
|
||||
client: Client::new(),
|
||||
client: Client::builder()
|
||||
.timeout(Duration::from_secs(120))
|
||||
.build()
|
||||
.expect("Failed to create HTTP client"),
|
||||
request_timeout: Duration::from_secs(60),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set a custom request timeout.
|
||||
pub fn with_timeout(mut self, timeout: Duration) -> Self {
|
||||
self.request_timeout = timeout;
|
||||
self
|
||||
}
|
||||
|
||||
/// Build the request body for the chat completions endpoint.
|
||||
fn build_request(&self, context: &ConversationContext) -> serde_json::Value {
|
||||
let messages: Vec<serde_json::Value> = context
|
||||
@@ -61,56 +86,32 @@ impl OpenAIProvider {
|
||||
|
||||
body
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for OpenAIProvider {
|
||||
async fn chat(&self, context: ConversationContext) -> ProviderResult<ProviderResponse> {
|
||||
debug!("OpenAI chat request (model: {})", self.model);
|
||||
|
||||
let request_body = self.build_request(&context);
|
||||
let url = format!("{}/chat/completions", self.base_url);
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.bearer_auth(&self.api_key)
|
||||
.json(&request_body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let error_text = response.text().await.unwrap_or_default();
|
||||
return Err(gyandoot_core::ProviderError::Api(format!(
|
||||
"OpenAI API error ({}): {}",
|
||||
status, error_text
|
||||
)));
|
||||
}
|
||||
|
||||
let data: serde_json::Value = response.json().await?;
|
||||
|
||||
/// Parse the non-streaming API response into a ProviderResponse.
|
||||
fn parse_response(&self, data: &serde_json::Value) -> ProviderResult<ProviderResponse> {
|
||||
let choice = &data["choices"][0];
|
||||
let message = &choice["message"];
|
||||
let content = message["content"].as_str().unwrap_or_default().to_string();
|
||||
let finish_reason = choice["finish_reason"].as_str().unwrap_or("stop").to_string();
|
||||
|
||||
// Parse tool calls
|
||||
let tool_calls = message["tool_calls"]
|
||||
.as_array()
|
||||
.map(|calls| {
|
||||
calls
|
||||
.iter()
|
||||
.map(|tc| gyandoot_core::ToolCall {
|
||||
id: tc["id"].as_str().unwrap_or_default().to_string(),
|
||||
name: tc["function"]["name"].as_str().unwrap_or_default().to_string(),
|
||||
arguments: tc["function"]["arguments"]
|
||||
.as_str()
|
||||
.and_then(|s| serde_json::from_str(s).ok())
|
||||
.unwrap_or(serde_json::Value::Null),
|
||||
})
|
||||
.collect()
|
||||
});
|
||||
let content = message["content"].as_str().unwrap_or_default().to_string();
|
||||
let finish_reason = choice["finish_reason"]
|
||||
.as_str()
|
||||
.unwrap_or("stop")
|
||||
.to_string();
|
||||
|
||||
// Parse tool calls from the response
|
||||
let tool_calls = message["tool_calls"].as_array().map(|calls| {
|
||||
calls
|
||||
.iter()
|
||||
.map(|tc| ToolCall {
|
||||
id: tc["id"].as_str().unwrap_or_default().to_string(),
|
||||
name: tc["function"]["name"].as_str().unwrap_or_default().to_string(),
|
||||
arguments: tc["function"]["arguments"]
|
||||
.as_str()
|
||||
.and_then(|s| serde_json::from_str(s).ok())
|
||||
.unwrap_or(serde_json::Value::Null),
|
||||
})
|
||||
.collect()
|
||||
});
|
||||
|
||||
Ok(ProviderResponse {
|
||||
content,
|
||||
@@ -119,61 +120,206 @@ impl Provider for OpenAIProvider {
|
||||
})
|
||||
}
|
||||
|
||||
async fn chat_stream(
|
||||
/// Send the request and handle common HTTP error statuses.
|
||||
async fn send_request(
|
||||
&self,
|
||||
context: ConversationContext,
|
||||
) -> ProviderResult<Box<dyn tokio_stream::Stream<Item = ProviderResult<String>> + Send + Unpin>>
|
||||
{
|
||||
debug!("OpenAI streaming chat request (model: {})", self.model);
|
||||
|
||||
let request_body = self.build_request(&context);
|
||||
let url = format!("{}/chat/completions", self.base_url);
|
||||
body: &serde_json::Value,
|
||||
) -> ProviderResult<reqwest::Response> {
|
||||
let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.bearer_auth(&self.api_key)
|
||||
.json(&request_body)
|
||||
.json(body)
|
||||
.timeout(self.request_timeout)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(gyandoot_core::ProviderError::Api(format!(
|
||||
"OpenAI stream error: {}",
|
||||
response.status()
|
||||
)));
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let error_text = response.text().await.unwrap_or_default();
|
||||
return Err(ProviderError::from_response(status.as_u16(), &error_text));
|
||||
}
|
||||
|
||||
// We can't easily return a boxed stream from reqwest's streaming API
|
||||
// without concrete types. For now we'll delegate to the non-streaming path
|
||||
// wrapped in a simple stream. Real implementation would use SSE parsing.
|
||||
let result = self.chat(context).await?;
|
||||
let stream = tokio_stream::StreamExt::chain(
|
||||
tokio_stream::once(Ok(result.content)),
|
||||
tokio_stream::empty(),
|
||||
Ok(response)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for OpenAIProvider {
|
||||
async fn chat(&self, context: ConversationContext) -> ProviderResult<ProviderResponse> {
|
||||
debug!(
|
||||
"OpenAI chat request (model: {}, messages: {})",
|
||||
self.model,
|
||||
context.messages.len()
|
||||
);
|
||||
|
||||
let request_body = self.build_request(&context);
|
||||
let response = self.send_request(&request_body).await?;
|
||||
let data: serde_json::Value = response.json().await?;
|
||||
|
||||
Ok(self.parse_response(&data)?)
|
||||
}
|
||||
|
||||
async fn chat_stream(
|
||||
&self,
|
||||
context: ConversationContext,
|
||||
) -> ProviderResult<Box<dyn tokio_stream::Stream<Item = ProviderResult<String>> + Send + Unpin>>
|
||||
{
|
||||
debug!(
|
||||
"OpenAI streaming chat (model: {}, messages: {})",
|
||||
self.model,
|
||||
context.messages.len()
|
||||
);
|
||||
|
||||
let mut request_body = self.build_request(&context);
|
||||
request_body["stream"] = serde_json::json!(true);
|
||||
|
||||
let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
|
||||
|
||||
// Start the streaming HTTP request
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.bearer_auth(&self.api_key)
|
||||
.json(&request_body)
|
||||
.timeout(self.request_timeout)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let error_text = response.text().await.unwrap_or_default();
|
||||
return Err(ProviderError::from_response(status.as_u16(), &error_text));
|
||||
}
|
||||
|
||||
// Use a channel to bridge the reqwest stream to the tokio_stream
|
||||
let (tx, rx) = mpsc::channel::<ProviderResult<String>>(256);
|
||||
let mut sse_stream = response.bytes_stream();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut parser = crate::streaming::SseParser::new(100);
|
||||
let mut buffer = Vec::new();
|
||||
|
||||
loop {
|
||||
match tokio::time::timeout(Duration::from_secs(30), sse_stream.next()).await {
|
||||
Ok(Some(Ok(chunk))) => {
|
||||
buffer.extend_from_slice(&chunk);
|
||||
let buffer_str = String::from_utf8_lossy(&buffer);
|
||||
let mut pos = 0;
|
||||
|
||||
for event in parser.feed_chunk(buffer_str.as_ref()) {
|
||||
match event {
|
||||
StreamEvent::Chunk(text) => {
|
||||
trace!("SSE chunk: {text}");
|
||||
if tx.send(Ok(text)).await.is_err() {
|
||||
debug!("Stream receiver dropped");
|
||||
return;
|
||||
}
|
||||
}
|
||||
StreamEvent::ToolCall(tc) => {
|
||||
trace!("SSE tool call: {tc:?}");
|
||||
}
|
||||
StreamEvent::Done => {
|
||||
debug!("SSE stream done");
|
||||
let _ = tx.try_send(Ok(String::new()));
|
||||
return;
|
||||
}
|
||||
StreamEvent::Error(msg) => {
|
||||
warn!("SSE error: {msg}");
|
||||
let _ = tx
|
||||
.send(Err(ProviderError::Stream(msg)))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
pos = parser.consumed();
|
||||
}
|
||||
buffer.drain(0..pos);
|
||||
}
|
||||
Ok(Some(Err(e))) => {
|
||||
let _ = tx
|
||||
.send(Err(ProviderError::Request(e)))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
Ok(None) => {
|
||||
debug!("SSE stream exhausted");
|
||||
return;
|
||||
}
|
||||
Err(_) => {
|
||||
let _ = tx
|
||||
.send(Err(ProviderError::Timeout {
|
||||
duration: 30,
|
||||
}))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Also collect tool calls from the non-streaming path as a fallback
|
||||
// For simplicity, if the stream doesn't produce content, fall back
|
||||
let stream = tokio_stream::wrappers::ReceiverStream::new(rx);
|
||||
Ok(Box::new(stream))
|
||||
}
|
||||
|
||||
async fn models(&self) -> ProviderResult<Vec<String>> {
|
||||
let url = format!("{}/models", self.base_url.trim_end_matches('/'));
|
||||
let response = self
|
||||
.client
|
||||
.get(&url)
|
||||
.bearer_auth(&self.api_key)
|
||||
.timeout(self.request_timeout)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let error_text = response.text().await.unwrap_or_default();
|
||||
return Err(ProviderError::from_response(status.as_u16(), &error_text));
|
||||
}
|
||||
|
||||
let data: serde_json::Value = response.json().await?;
|
||||
let models = data["data"]
|
||||
.as_array()
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|m| m["id"].as_str().map(|s| s.to_string()))
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
Ok(models)
|
||||
}
|
||||
|
||||
fn supports_tools(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use gyandoot_core::{Message, ToolDefinition};
|
||||
use gyandoot_core::{Message, Role, ToolDefinition};
|
||||
|
||||
#[test]
|
||||
fn test_build_request() {
|
||||
let provider = OpenAIProvider::new(
|
||||
fn test_provider() -> OpenAIProvider {
|
||||
OpenAIProvider::new(
|
||||
"sk-test".to_string(),
|
||||
"gpt-4o".to_string(),
|
||||
None,
|
||||
);
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_request_basic() {
|
||||
let provider = test_provider();
|
||||
let context = ConversationContext {
|
||||
messages: vec![
|
||||
Message::new(MessageRole::System, "You are an AI."),
|
||||
Message::new(MessageRole::User, "Hello"),
|
||||
Message::new(Role::System, "You are an AI."),
|
||||
Message::new(Role::User, "Hello"),
|
||||
],
|
||||
system_prompt: "You are an AI.".to_string(),
|
||||
tools: vec![],
|
||||
@@ -184,4 +330,107 @@ mod tests {
|
||||
assert_eq!(body["model"], "gpt-4o");
|
||||
assert_eq!(body["messages"].as_array().unwrap().len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_request_with_tools() {
|
||||
let provider = test_provider();
|
||||
let context = ConversationContext {
|
||||
messages: vec![Message::new(Role::User, "What's the weather?")],
|
||||
system_prompt: String::new(),
|
||||
tools: vec![ToolDefinition::new(
|
||||
"get_weather",
|
||||
"Get weather for a city",
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {"type": "string"}
|
||||
},
|
||||
"required": ["city"]
|
||||
}),
|
||||
)],
|
||||
max_tokens: 4096,
|
||||
};
|
||||
|
||||
let body = provider.build_request(&context);
|
||||
assert!(body.get("tools").is_some());
|
||||
assert_eq!(body["tool_choice"], "auto");
|
||||
assert_eq!(body["tools"].as_array().unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_request_tool_result() {
|
||||
let provider = test_provider();
|
||||
let context = ConversationContext {
|
||||
messages: vec![
|
||||
Message::new(Role::Assistant, "")
|
||||
.with_tool_calls(vec![ToolCall::new("call_1", "get_weather", serde_json::json!({"city": "NYC"}))]),
|
||||
Message::new(Role::Tool, "72°F, sunny").with_tool_result("72°F, sunny".to_string()),
|
||||
],
|
||||
system_prompt: String::new(),
|
||||
tools: vec![],
|
||||
max_tokens: 4096,
|
||||
};
|
||||
|
||||
let body = provider.build_request(&context);
|
||||
let messages = body["messages"].as_array().unwrap();
|
||||
assert_eq!(messages.len(), 2);
|
||||
// Tool message should have tool_call_id and content
|
||||
let tool_msg = &messages[1];
|
||||
assert!(tool_msg.get("tool_call_id").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_response() {
|
||||
let provider = test_provider();
|
||||
let data = serde_json::json!({
|
||||
"choices": [{
|
||||
"message": {
|
||||
"content": "Hello, how can I help?",
|
||||
"role": "assistant"
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
"index": 0
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 5,
|
||||
"total_tokens": 15
|
||||
}
|
||||
});
|
||||
|
||||
let response = provider.parse_response(&data).unwrap();
|
||||
assert_eq!(response.content, "Hello, how can I help?");
|
||||
assert_eq!(response.finish_reason, "stop");
|
||||
assert!(response.tool_calls.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_response_with_tool_calls() {
|
||||
let provider = test_provider();
|
||||
let data = serde_json::json!({
|
||||
"choices": [{
|
||||
"message": {
|
||||
"content": null,
|
||||
"role": "assistant",
|
||||
"tool_calls": [{
|
||||
"id": "call_abc123",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": "{\"city\": \"NYC\"}"
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finish_reason": "tool_calls",
|
||||
"index": 0
|
||||
}]
|
||||
});
|
||||
|
||||
let response = provider.parse_response(&data).unwrap();
|
||||
assert!(response.has_tool_calls());
|
||||
let calls = response.tool_calls.unwrap();
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].id, "call_abc123");
|
||||
assert_eq!(calls[0].name, "get_weather");
|
||||
}
|
||||
}
|
||||
@@ -45,6 +45,88 @@ impl tokio_stream::Stream for TokenStream {
|
||||
}
|
||||
}
|
||||
|
||||
// ─── SSE Parser ────────────────────────────────────────────────────────────
|
||||
|
||||
/// Events emitted by the [`SseParser`].
|
||||
#[derive(Debug)]
|
||||
pub enum StreamEvent {
|
||||
/// Raw data from a `data:` line (without the `data: ` prefix).
|
||||
Chunk(String),
|
||||
/// Stream complete signal (`data: [DONE]`).
|
||||
Done,
|
||||
/// An error-related line was encountered.
|
||||
Error(String),
|
||||
/// A tool call reference detected in the SSE stream (OpenAI format).
|
||||
ToolCall(super::ToolCall),
|
||||
}
|
||||
|
||||
/// A minimal SSE (Server-Sent Events) parser that emits [`StreamEvent`] values.
|
||||
///
|
||||
/// Processes byte chunks line-by-line. Only complete lines (terminated by `\n`)
|
||||
/// are processed; a partial final line is held over for the next call.
|
||||
/// Use [`consumed()`](SseParser::consumed) to learn how many bytes were fully
|
||||
/// consumed so the caller can drain them from its buffer.
|
||||
pub struct SseParser {
|
||||
consumed: usize,
|
||||
}
|
||||
|
||||
impl SseParser {
|
||||
/// Create a new SSE parser.
|
||||
pub fn new(_buffer_size: usize) -> Self {
|
||||
Self { consumed: 0 }
|
||||
}
|
||||
|
||||
/// Feed raw SSE-sourced text and return any complete events found.
|
||||
///
|
||||
/// Only lines terminated by `\n` are processed. The returned `consumed`
|
||||
/// count reflects the byte offset of the last fully-processed newline.
|
||||
pub fn feed_chunk<'a>(&'a mut self, chunk: &str) -> Vec<StreamEvent> {
|
||||
let mut events = Vec::new();
|
||||
let mut pos = 0;
|
||||
let bytes = chunk.as_bytes();
|
||||
let len = bytes.len();
|
||||
let mut line_start = 0;
|
||||
let mut i = 0;
|
||||
|
||||
while i < len {
|
||||
if bytes[i] == b'\n' {
|
||||
let line = &chunk[line_start..i];
|
||||
if line.starts_with("data: ") {
|
||||
let data = &line[6..];
|
||||
if data == "[DONE]" {
|
||||
events.push(StreamEvent::Done);
|
||||
} else {
|
||||
events.push(StreamEvent::Chunk(data.to_string()));
|
||||
}
|
||||
} else if !line.is_empty() {
|
||||
// Check for error lines
|
||||
let trimmed = line.trim();
|
||||
if trimmed.starts_with(':') {
|
||||
// SSE comment — skip
|
||||
} else if trimmed.starts_with("event:") || trimmed.starts_with("id:") || trimmed.starts_with("retry:") {
|
||||
// SSE metadata — skip
|
||||
} else if trimmed.contains("error") {
|
||||
events.push(StreamEvent::Error(line.to_string()));
|
||||
}
|
||||
}
|
||||
pos = i + 1;
|
||||
line_start = i + 1;
|
||||
i = pos;
|
||||
} else {
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
|
||||
self.consumed = pos;
|
||||
events
|
||||
}
|
||||
|
||||
/// Return the number of bytes consumed from the buffer.
|
||||
pub fn consumed(&self) -> usize {
|
||||
self.consumed
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -55,4 +137,38 @@ mod tests {
|
||||
let result = collect_stream(stream).await.unwrap();
|
||||
assert_eq!(result, "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sse_basic_chunk() {
|
||||
let mut parser = SseParser::new(100);
|
||||
let events = parser.feed_chunk("data: hello\n\n");
|
||||
assert_eq!(events.len(), 1);
|
||||
assert!(matches!(&events[0], StreamEvent::Chunk(t) if t == "hello"));
|
||||
assert!(parser.consumed() > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sse_done() {
|
||||
let mut parser = SseParser::new(100);
|
||||
let events = parser.feed_chunk("data: [DONE]\n\n");
|
||||
assert_eq!(events.len(), 1);
|
||||
assert!(matches!(events[0], StreamEvent::Done));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sse_partial_line() {
|
||||
let mut parser = SseParser::new(100);
|
||||
let events = parser.feed_chunk("data: hel");
|
||||
assert!(events.is_empty());
|
||||
assert_eq!(parser.consumed(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sse_two_chunks() {
|
||||
let mut parser = SseParser::new(100);
|
||||
let _ = parser.feed_chunk("data: hel");
|
||||
let events = parser.feed_chunk("lo\n\n");
|
||||
assert_eq!(events.len(), 1);
|
||||
assert!(matches!(&events[0], StreamEvent::Chunk(t) if t == "hello"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user