use std::{
future::Future,
pin::Pin,
sync::Arc,
task::{Context, Poll},
};
use bob_core::{
error::{LlmError, ToolError},
ports::{LlmPort, ToolPort},
types::{LlmRequest, LlmResponse, ToolCall, ToolDescriptor, ToolResult},
};
#[derive(Debug, Clone)]
pub struct ToolRequest {
pub call: ToolCall,
}
impl ToolRequest {
#[must_use]
pub fn new(call: ToolCall) -> Self {
Self { call }
}
#[must_use]
pub fn from_parts(name: impl Into<String>, arguments: serde_json::Value) -> Self {
Self { call: ToolCall::new(name, arguments) }
}
}
#[derive(Debug, Clone)]
pub struct ToolResponse {
pub result: ToolResult,
}
#[derive(Debug, Clone)]
pub struct LlmRequestWrapper {
pub request: LlmRequest,
}
#[derive(Debug, Clone)]
pub struct LlmResponseWrapper {
pub response: LlmResponse,
}
#[derive(Debug, Clone, Copy, Default)]
pub struct ToolListRequest;
pub struct ToolService {
inner: Arc<dyn ToolPort>,
}
impl ToolService {
#[must_use]
pub fn new(port: Arc<dyn ToolPort>) -> Self {
Self { inner: port }
}
}
impl std::fmt::Debug for ToolService {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ToolService").finish_non_exhaustive()
}
}
impl tower::Service<ToolRequest> for ToolService {
type Response = ToolResponse;
type Error = ToolError;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, req: ToolRequest) -> Self::Future {
let inner = self.inner.clone();
Box::pin(async move {
let result = inner.call_tool(req.call).await?;
Ok(ToolResponse { result })
})
}
}
pub struct LlmService {
inner: Arc<dyn LlmPort>,
}
impl LlmService {
#[must_use]
pub fn new(port: Arc<dyn LlmPort>) -> Self {
Self { inner: port }
}
}
impl std::fmt::Debug for LlmService {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LlmService").finish_non_exhaustive()
}
}
impl tower::Service<LlmRequestWrapper> for LlmService {
type Response = LlmResponseWrapper;
type Error = LlmError;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, req: LlmRequestWrapper) -> Self::Future {
let inner = self.inner.clone();
Box::pin(async move {
let response = inner.complete(req.request).await?;
Ok(LlmResponseWrapper { response })
})
}
}
pub struct ToolListService {
inner: Arc<dyn ToolPort>,
}
impl ToolListService {
#[must_use]
pub fn new(port: Arc<dyn ToolPort>) -> Self {
Self { inner: port }
}
}
impl std::fmt::Debug for ToolListService {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ToolListService").finish_non_exhaustive()
}
}
impl tower::Service<ToolListRequest> for ToolListService {
type Response = Vec<ToolDescriptor>;
type Error = ToolError;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, _req: ToolListRequest) -> Self::Future {
let inner = self.inner.clone();
Box::pin(async move { inner.list_tools().await })
}
}
pub trait ServiceExt<Request>: tower::Service<Request> + Sized {
fn with_timeout(self, timeout: std::time::Duration) -> tower::timeout::Timeout<Self> {
tower::timeout::Timeout::new(self, timeout)
}
fn with_rate_limit(
self,
max: u64,
interval: std::time::Duration,
) -> tower::limit::RateLimit<Self> {
tower::limit::RateLimit::new(self, tower::limit::rate::Rate::new(max, interval))
}
fn with_concurrency_limit(self, max: usize) -> tower::limit::ConcurrencyLimit<Self> {
tower::limit::ConcurrencyLimit::new(self, max)
}
fn map_err<F, E2>(self, f: F) -> tower::util::MapErr<Self, F>
where
F: FnOnce(Self::Error) -> E2,
{
tower::util::MapErr::new(self, f)
}
fn map_response<F, Response2>(self, f: F) -> tower::util::MapResponse<Self, F>
where
F: FnOnce(Self::Response) -> Response2,
{
tower::util::MapResponse::new(self, f)
}
fn boxed(self) -> tower::util::BoxService<Request, Self::Response, Self::Error>
where
Self: Send + 'static,
Request: Send + 'static,
Self::Future: Send + 'static,
{
tower::util::BoxService::new(self)
}
}
impl<T, Request> ServiceExt<Request> for T where T: tower::Service<Request> + Sized {}
pub trait ToolPortServiceExt: ToolPort {
fn into_tool_service(self: Arc<Self>) -> ToolService;
fn into_tool_list_service(self: Arc<Self>) -> ToolListService;
}
impl<T: ToolPort + 'static> ToolPortServiceExt for T {
fn into_tool_service(self: Arc<Self>) -> ToolService {
ToolService::new(self)
}
fn into_tool_list_service(self: Arc<Self>) -> ToolListService {
ToolListService::new(self)
}
}
impl ToolPortServiceExt for dyn ToolPort {
fn into_tool_service(self: Arc<Self>) -> ToolService {
ToolService::new(self)
}
fn into_tool_list_service(self: Arc<Self>) -> ToolListService {
ToolListService::new(self)
}
}
pub trait LlmPortServiceExt: LlmPort {
fn into_llm_service(self: Arc<Self>) -> LlmService;
}
impl<T: LlmPort + 'static> LlmPortServiceExt for T {
fn into_llm_service(self: Arc<Self>) -> LlmService {
LlmService::new(self)
}
}
impl LlmPortServiceExt for dyn LlmPort {
fn into_llm_service(self: Arc<Self>) -> LlmService {
LlmService::new(self)
}
}
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use bob_core::types::ToolDescriptor;
use tower::Service;
use super::*;
struct MockToolPort {
calls: Mutex<Vec<String>>,
}
impl MockToolPort {
fn new() -> Self {
Self { calls: Mutex::new(Vec::new()) }
}
}
#[async_trait::async_trait]
impl ToolPort for MockToolPort {
async fn list_tools(&self) -> Result<Vec<ToolDescriptor>, ToolError> {
Ok(vec![ToolDescriptor::new("mock/echo", "Echo tool")])
}
async fn call_tool(&self, call: ToolCall) -> Result<ToolResult, ToolError> {
let mut calls = self.calls.lock().unwrap_or_else(|p| p.into_inner());
calls.push(call.name.clone());
Ok(ToolResult {
name: call.name,
output: serde_json::json!({"ok": true}),
is_error: false,
})
}
}
#[tokio::test]
async fn tool_service_basic() {
let port: Arc<dyn ToolPort> = Arc::new(MockToolPort::new());
let mut svc = ToolService::new(port);
let resp = svc.call(ToolRequest::from_parts("mock/echo", serde_json::json!({}))).await;
assert!(resp.is_ok());
assert_eq!(resp.unwrap().result.name, "mock/echo");
}
#[tokio::test]
async fn tool_list_service() {
let port: Arc<dyn ToolPort> = Arc::new(MockToolPort::new());
let mut svc = ToolListService::new(port);
let tools = svc.call(ToolListRequest).await.expect("should list tools");
assert_eq!(tools.len(), 1);
assert_eq!(tools[0].id, "mock/echo");
}
#[tokio::test]
async fn service_ext_timeout() {
let port: Arc<dyn ToolPort> = Arc::new(MockToolPort::new());
let svc = ToolService::new(port);
let mut timeout_svc = svc.with_timeout(std::time::Duration::from_secs(1));
let resp =
timeout_svc.call(ToolRequest::from_parts("mock/echo", serde_json::json!({}))).await;
assert!(resp.is_ok());
}
#[tokio::test]
async fn port_service_ext() {
let port: Arc<dyn ToolPort> = Arc::new(MockToolPort::new());
let mut svc = port.into_tool_service();
let resp = svc.call(ToolRequest::from_parts("mock/echo", serde_json::json!({}))).await;
assert!(resp.is_ok());
}
}