use crate::in_memory::{InMemorySessionStorageStore, InMemorySessionStore};
use async_trait::async_trait;
use everruns_core::agent::Agent;
use everruns_core::error::Result;
use everruns_core::events::Event;
use everruns_core::harness::Harness;
use everruns_core::memory::{
InMemoryAgentStore, InMemoryEventEmitter, InMemoryHarnessStore, InMemoryLlmProviderStore,
InMemoryMemoryStore, InMemoryMessageRetriever,
};
use everruns_core::memory_store::MemoryStoreBackend;
use everruns_core::message::Message;
use everruns_core::message_retriever::{InputMessage, MessageRetriever};
use everruns_core::session::Session;
use everruns_core::traits::{
AgentStore, EventEmitter, HarnessStore, LlmProviderStore, ModelWithProvider, SessionMutator,
SessionStorageStore, SessionStore,
};
use everruns_core::typed_id::SessionId;
use std::sync::Arc;
#[async_trait]
pub trait RuntimeAgentStore: AgentStore + Send + Sync {
async fn add_agent(&self, agent: Agent) -> Result<()>;
}
#[async_trait]
pub trait RuntimeHarnessStore: HarnessStore + Send + Sync {
async fn add_harness(&self, harness: Harness) -> Result<()>;
}
#[async_trait]
pub trait RuntimeSessionStore: SessionStore + SessionMutator + Send + Sync {
async fn add_session(&self, session: Session) -> Result<()>;
}
#[async_trait]
pub trait RuntimeMessageStore: MessageRetriever + Send + Sync {
async fn add_input_message(
&self,
session_id: SessionId,
input: InputMessage,
) -> Result<Message>;
async fn store_message(&self, session_id: SessionId, message: Message) -> Result<()>;
}
#[async_trait]
pub trait RuntimeProviderStore: LlmProviderStore + Send + Sync {
async fn set_default_model(&self, model: ModelWithProvider) -> Result<()>;
}
#[async_trait]
pub trait EventBus: EventEmitter {
async fn collected_events(&self) -> Vec<Event> {
Vec::new()
}
}
#[async_trait]
impl<T: EventBus + ?Sized> EventBus for Arc<T> {
async fn collected_events(&self) -> Vec<Event> {
(**self).collected_events().await
}
}
#[derive(Clone)]
pub struct RuntimeBackends {
pub harness_store: Arc<dyn RuntimeHarnessStore>,
pub agent_store: Arc<dyn RuntimeAgentStore>,
pub session_store: Arc<dyn RuntimeSessionStore>,
pub message_store: Arc<dyn RuntimeMessageStore>,
pub provider_store: Arc<dyn RuntimeProviderStore>,
pub event_bus: Arc<dyn EventBus>,
pub storage_store: Arc<dyn SessionStorageStore>,
pub memory_store: Arc<dyn MemoryStoreBackend>,
}
impl RuntimeBackends {
pub fn in_memory() -> Self {
let event_bus = Arc::new(InMemoryEventEmitter::new());
Self {
harness_store: Arc::new(InMemoryHarnessStore::new()),
agent_store: Arc::new(InMemoryAgentStore::new()),
session_store: Arc::new(InMemorySessionStore::new()),
message_store: Arc::new(InMemoryMessageRetriever::new()),
provider_store: Arc::new(InMemoryLlmProviderStore::new()),
event_bus,
storage_store: Arc::new(InMemorySessionStorageStore::new()),
memory_store: Arc::new(InMemoryMemoryStore::new()),
}
}
pub fn with_harness_store(mut self, store: Arc<dyn RuntimeHarnessStore>) -> Self {
self.harness_store = store;
self
}
pub fn with_agent_store(mut self, store: Arc<dyn RuntimeAgentStore>) -> Self {
self.agent_store = store;
self
}
pub fn with_session_store(mut self, store: Arc<dyn RuntimeSessionStore>) -> Self {
self.session_store = store;
self
}
pub fn with_message_store(mut self, store: Arc<dyn RuntimeMessageStore>) -> Self {
self.message_store = store;
self
}
pub fn with_provider_store(mut self, store: Arc<dyn RuntimeProviderStore>) -> Self {
self.provider_store = store;
self
}
pub fn with_event_bus(mut self, bus: Arc<dyn EventBus>) -> Self {
self.event_bus = bus;
self
}
pub fn with_storage_store(mut self, store: Arc<dyn SessionStorageStore>) -> Self {
self.storage_store = store;
self
}
pub fn with_memory_store(mut self, store: Arc<dyn MemoryStoreBackend>) -> Self {
self.memory_store = store;
self
}
}
#[async_trait]
impl RuntimeAgentStore for InMemoryAgentStore {
async fn add_agent(&self, agent: Agent) -> Result<()> {
InMemoryAgentStore::add_agent(self, agent).await;
Ok(())
}
}
#[async_trait]
impl RuntimeHarnessStore for InMemoryHarnessStore {
async fn add_harness(&self, harness: Harness) -> Result<()> {
InMemoryHarnessStore::add_harness(self, harness).await;
Ok(())
}
}
#[async_trait]
impl RuntimeSessionStore for InMemorySessionStore {
async fn add_session(&self, session: Session) -> Result<()> {
InMemorySessionStore::add_session(self, session).await;
Ok(())
}
}
#[async_trait]
impl RuntimeMessageStore for InMemoryMessageRetriever {
async fn add_input_message(
&self,
session_id: SessionId,
input: InputMessage,
) -> Result<Message> {
self.add(session_id, input).await
}
async fn store_message(&self, session_id: SessionId, message: Message) -> Result<()> {
self.store(session_id, message).await
}
}
#[async_trait]
impl RuntimeProviderStore for InMemoryLlmProviderStore {
async fn set_default_model(&self, model: ModelWithProvider) -> Result<()> {
InMemoryLlmProviderStore::set_default_model(self, model).await;
Ok(())
}
}
#[async_trait]
impl EventBus for InMemoryEventEmitter {
async fn collected_events(&self) -> Vec<Event> {
InMemoryEventEmitter::events(self).await
}
}