omniference 0.1.1

A multi-protocol inference engine with provider adapters
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

use crate::stream::StreamEvent;
use crate::types::ChatRequestIR;
use async_trait::async_trait;
use futures_util::Stream;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;

pub mod logging;

pub type ChatStream = Box<dyn Stream<Item = StreamEvent> + Send + Unpin>;

/// Represents the next step in the middleware chain.
/// This could be another middleware or the final router.
#[async_trait]
pub trait RequestHandler: Send + Sync {
    async fn handle(
        &self,
        request: ChatRequestIR,
        cancel: CancellationToken,
    ) -> anyhow::Result<ChatStream>;
}

/// Middleware trait for intercepting requests and responses
#[async_trait]
pub trait Middleware: Send + Sync {
    async fn handle(
        &self,
        request: ChatRequestIR,
        cancel: CancellationToken,
        next: &dyn RequestHandler,
    ) -> anyhow::Result<ChatStream>;
}

/// A wrapper to chain a middleware with the next handler
pub struct MiddlewareChain {
    middleware: Arc<dyn Middleware>,
    next: Arc<dyn RequestHandler>,
}

impl MiddlewareChain {
    pub fn new(middleware: Arc<dyn Middleware>, next: Arc<dyn RequestHandler>) -> Self {
        Self { middleware, next }
    }
}

#[async_trait]
impl RequestHandler for MiddlewareChain {
    async fn handle(
        &self,
        request: ChatRequestIR,
        cancel: CancellationToken,
    ) -> anyhow::Result<ChatStream> {
        self.middleware
            .handle(request, cancel, self.next.as_ref())
            .await
    }
}