myst-client 0.1.4

use async_openai::types::{CreateChatCompletionRequest, CreateChatCompletionStreamResponse};
use futures::Stream;
use serde::{Deserialize, Serialize};
use std::{net::SocketAddr, pin::Pin};

use crate::transport::{NodeId, RpcClient};

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct TaskContext {
    pub node_id: NodeId,
    pub proxy_addr: Option<SocketAddr>,
    pub model_ctx: ModelContext,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ModelContext {
    /// The base URL for the API associated with this model.
    pub api_base_url: Option<String>,

    /// The name of the environment variable containing the credentials for this model.
    pub credentials_env_name: Option<String>,
}

#[derive(Debug, Deserialize)]
pub enum Task {
    Text {
        node_id: Option<NodeId>,
        request: CreateChatCompletionRequest,
    },
    Image {
        node_id: Option<NodeId>,
        prompt: String,
    },
}

#[derive(Debug, Serialize)]
pub enum TaskResponse {
    Text {
        node_id: Option<NodeId>,
    },
    Image {
        node_id: Option<NodeId>,
        data: String,
    },
}

#[derive(thiserror::Error, Debug)]
pub enum TaskError {
    #[error("text task failed: {0}")]
    Text(String),
    #[error("image task failed: {0}")]
    Image(String),
    #[error("Network error: {0}")]
    Network(String),
}

#[derive(Default, Clone)]
pub struct Runner {
    // add configurations for caching, workflows, etc.
}

impl Runner {
    pub async fn start_text_stream(
        &self,
        ctx: TaskContext,
        request: CreateChatCompletionRequest,
    ) -> Result<
        (
            TaskResponse,
            Pin<
                Box<
                    dyn Stream<Item = Result<CreateChatCompletionStreamResponse, anyhow::Error>>
                        + Send,
                >,
            >,
        ),
        TaskError,
    > {
        eprintln!("starting text stream for ctx: {:?}", ctx);

        let mut client = RpcClient::new(ctx.node_id.clone(), ctx.proxy_addr)
            .await
            .map_err(|e| TaskError::Network(e.to_string()))?;

        let stream = client
            .compute_text(request, ctx.clone())
            .await
            .map_err(|e| TaskError::Text(e.to_string()))?;

        Ok((
            TaskResponse::Text {
                node_id: Some(ctx.node_id),
            },
            stream,
        ))
    }

    pub async fn run(&self, ctx: TaskContext, task: Task) -> Result<TaskResponse, TaskError> {
        eprintln!("running task with context: {:?}", ctx);
        let mut client = RpcClient::new(ctx.node_id.clone(), ctx.proxy_addr)
            .await
            .map_err(|e| TaskError::Network(e.to_string()))?;

        match task {
            Task::Image { node_id, prompt } => {
                let node_id = node_id.clone();
                let data = client
                    .generate_image(prompt, ctx)
                    .await
                    .map_err(|e| TaskError::Image(e.to_string()))?;

                Ok(TaskResponse::Image { node_id, data })
            }
            Task::Text { .. } => Err(TaskError::Text(
                "Use start_text_stream for text completion tasks".into(),
            )),
        }
    }
}