yopo 7.0.0

//! YOPO (You Only Prompt Once) - A simple library for testing SACP agents
//!
//! Provides a convenient API for running one-shot prompts against SACP components.

use sacp::ClientToAgent;
use sacp::schema::{
    AudioContent, ContentBlock, EmbeddedResourceResource, ImageContent, InitializeRequest,
    RequestPermissionOutcome, RequestPermissionRequest, RequestPermissionResponse,
    SessionNotification, TextContent, VERSION as PROTOCOL_VERSION,
};
use sacp::util::MatchMessage;
use sacp::{Component, Handled, MessageCx, UntypedMessage};
use std::path::PathBuf;

/// Converts a `ContentBlock` to its string representation.
///
/// This function provides standard string conversions for different content types:
/// - `Text`: Returns the text content
/// - `Image`: Returns a placeholder like `[Image: image/png]`
/// - `Audio`: Returns a placeholder like `[Audio: audio/wav]`
/// - `ResourceLink`: Returns the URI
/// - `Resource`: Returns the URI
///
/// # Example
///
/// ```no_run
/// use yopo::content_block_to_string;
/// use sacp::schema::{ContentBlock, TextContent};
///
/// let block = ContentBlock::Text(TextContent {
///     annotations: None,
///     text: "Hello".to_string(),
///     meta: None,
/// });
/// assert_eq!(content_block_to_string(&block), "Hello");
/// ```
pub fn content_block_to_string(block: &ContentBlock) -> String {
    match block {
        ContentBlock::Text(TextContent { text, .. }) => text.clone(),
        ContentBlock::Image(ImageContent { mime_type, .. }) => {
            format!("[Image: {}]", mime_type)
        }
        ContentBlock::Audio(AudioContent { mime_type, .. }) => {
            format!("[Audio: {}]", mime_type)
        }
        ContentBlock::ResourceLink(link) => link.uri.clone(),
        ContentBlock::Resource(resource) => match &resource.resource {
            EmbeddedResourceResource::TextResourceContents(text) => text.uri.clone(),
            EmbeddedResourceResource::BlobResourceContents(blob) => blob.uri.clone(),
        },
    }
}

/// Runs a single prompt against a component with a callback for each content block.
///
/// This function:
/// - Spawns the component
/// - Initializes the agent
/// - Creates a new session
/// - Sends the prompt
/// - Auto-approves all permission requests
/// - Calls the callback with each `ContentBlock` from agent messages
/// - Returns when the prompt completes
///
/// The callback receives each `ContentBlock` as it arrives and can process it
/// asynchronously (e.g., print it, accumulate it, etc.).
///
/// # Example
///
/// ```ignore
/// use yopo::{prompt_with_callback, content_block_to_string};
/// use sacp_tokio::AcpAgent;
/// use std::str::FromStr;
///
/// # async fn example() -> Result<(), sacp::Error> {
/// let agent = AcpAgent::from_str("python agent.py")?;
/// prompt_with_callback(agent, "What is 2+2?", async |block| {
///     print!("{}", content_block_to_string(&block));
/// }).await?;
/// # Ok(())
/// # }
/// ```
pub async fn prompt_with_callback(
    component: impl Component,
    prompt_text: impl ToString,
    mut callback: impl AsyncFnMut(ContentBlock) + Send,
) -> Result<(), sacp::Error> {
    // Convert prompt to String
    let prompt_text = prompt_text.to_string();

    // Run the client
    ClientToAgent::builder()
        .on_receive_message(
            async |message: MessageCx<UntypedMessage, UntypedMessage>, _cx| {
                tracing::trace!("received: {:?}", message.message());
                Ok(Handled::No {
                    message,
                    retry: false,
                })
            },
        )
        .connect_to(component)?
        .with_client(|cx: sacp::JrConnectionCx<ClientToAgent>| async move {
            // Initialize the agent
            let _init_response = cx
                .send_request(InitializeRequest {
                    protocol_version: PROTOCOL_VERSION,
                    client_capabilities: Default::default(),
                    client_info: None,
                    meta: None,
                })
                .block_task()
                .await?;

            let mut session = cx.build_session(PathBuf::from(".")).send_request().await?;

            session.send_prompt(prompt_text)?;

            loop {
                let update = session.read_update().await?;
                match update {
                    sacp::SessionMessage::SessionMessage(message) => {
                        MatchMessage::new(message)
                            .if_notification(async |notification: SessionNotification| {
                                tracing::debug!(
                                    ?notification,
                                    "yopo: received SessionNotification"
                                );
                                // Call the callback for each agent message chunk
                                if let sacp::schema::SessionUpdate::AgentMessageChunk(
                                    content_chunk,
                                ) = notification.update
                                {
                                    callback(content_chunk.content).await;
                                }
                                Ok(())
                            })
                            .await
                            .if_request(async |request: RequestPermissionRequest, request_cx| {
                                // Auto-approve all permission requests by selecting the first option
                                // that looks "allow-ish"
                                let outcome = request
                                    .options
                                    .iter()
                                    .find(|option| match option.kind {
                                        sacp::schema::PermissionOptionKind::AllowOnce
                                        | sacp::schema::PermissionOptionKind::AllowAlways => true,
                                        sacp::schema::PermissionOptionKind::RejectOnce
                                        | sacp::schema::PermissionOptionKind::RejectAlways => false,
                                    })
                                    .map(|option| RequestPermissionOutcome::Selected {
                                        option_id: option.id.clone(),
                                    })
                                    .unwrap_or(RequestPermissionOutcome::Cancelled);

                                request_cx.respond(RequestPermissionResponse {
                                    outcome,
                                    meta: None,
                                })?;

                                Ok(())
                            })
                            .await
                            .otherwise(async |_msg| Ok(()))
                            .await?;
                    }
                    sacp::SessionMessage::StopReason(stop_reason) => match stop_reason {
                        sacp::schema::StopReason::EndTurn => break,
                        sacp::schema::StopReason::MaxTokens => todo!(),
                        sacp::schema::StopReason::MaxTurnRequests => todo!(),
                        sacp::schema::StopReason::Refusal => todo!(),
                        sacp::schema::StopReason::Cancelled => todo!(),
                    },
                    _ => todo!(),
                }
            }

            Ok(())
        })
        .await?;

    Ok(())
}

/// Runs a single prompt against a component and returns the accumulated text response.
///
/// This function:
/// - Spawns the component
/// - Initializes the agent
/// - Creates a new session
/// - Sends the prompt
/// - Auto-approves all permission requests
/// - Accumulates all content from agent messages using [`content_block_to_string`]
/// - Returns the complete response as a String
///
/// This is a convenience wrapper around [`prompt_with_callback`] that accumulates
/// all content blocks into a single string.
///
/// # Example
///
/// ```ignore
/// use yopo::prompt;
/// use sacp_tokio::AcpAgent;
/// use std::str::FromStr;
///
/// # async fn example() -> Result<(), sacp::Error> {
/// let agent = AcpAgent::from_str("python agent.py")?;
/// let response = prompt(agent, "What is 2+2?").await?;
/// assert!(response.contains("4"));
/// # Ok(())
/// # }
/// ```
pub async fn prompt(
    component: impl Component,
    prompt_text: impl ToString,
) -> Result<String, sacp::Error> {
    let mut accumulated_text = String::new();
    prompt_with_callback(component, prompt_text, async |block| {
        let text = content_block_to_string(&block);
        accumulated_text.push_str(&text);
    })
    .await?;
    Ok(accumulated_text)
}