mentra 0.6.0

An agent runtime for tool-using LLM applications
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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

use std::{path::PathBuf, sync::Arc};

use tokio::{
    runtime::{Builder as RuntimeBuilder, Handle as TokioHandle, Runtime as TokioRuntime},
    sync::{Mutex as AsyncMutex, mpsc},
    task::{AbortHandle, JoinHandle},
};

use crate::{agent::Agent, error::RuntimeError, runtime::CancellationToken};

use super::{TeamManager, teammate_actor_loop};

#[derive(Clone)]
pub(crate) struct TeammateHost {
    backend: Arc<TeammateRuntimeBackend>,
}

enum TeammateRuntimeBackend {
    Current(TokioHandle),
    Owned(Arc<TokioRuntime>),
}

pub(crate) struct TeammateActorHandle {
    pub(crate) wake_tx: mpsc::UnboundedSender<()>,
    pub(crate) cancellation: CancellationToken,
    pub(crate) abort: AbortHandle,
}

impl Drop for TeammateActorHandle {
    fn drop(&mut self) {
        self.cancellation.cancel();
        self.abort.abort();
    }
}

impl TeammateHost {
    pub(crate) fn new() -> Result<Self, RuntimeError> {
        let backend = match TokioHandle::try_current() {
            Ok(handle) => TeammateRuntimeBackend::Current(handle),
            Err(_) => TeammateRuntimeBackend::Owned(Arc::new(
                RuntimeBuilder::new_multi_thread()
                    .worker_threads(2)
                    .enable_all()
                    .build()
                    .map_err(|error| {
                        RuntimeError::Store(format!(
                            "Failed to create shared teammate runtime: {error}"
                        ))
                    })?,
            )),
        };
        Ok(Self {
            backend: Arc::new(backend),
        })
    }

    pub(crate) fn spawn_teammate(
        &self,
        manager: TeamManager,
        team_dir: PathBuf,
        teammate_name: String,
        agent: Arc<AsyncMutex<Agent>>,
    ) -> TeammateActorHandle {
        let (wake_tx, wake_rx) = mpsc::unbounded_channel();
        let cancellation = CancellationToken::default();
        let task = self.spawn_task(teammate_actor_loop(
            manager,
            team_dir,
            teammate_name,
            agent,
            wake_rx,
            cancellation.clone(),
        ));
        let abort = task.abort_handle();
        TeammateActorHandle {
            wake_tx,
            cancellation,
            abort,
        }
    }

    fn spawn_task(
        &self,
        future: impl std::future::Future<Output = ()> + Send + 'static,
    ) -> JoinHandle<()> {
        match self.backend.as_ref() {
            TeammateRuntimeBackend::Current(handle) => handle.spawn(future),
            TeammateRuntimeBackend::Owned(runtime) => runtime.handle().spawn(future),
        }
    }
}