llm-agent 0.3.0

The agent library to build LLM applications that work with any LLM providers.
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

use crate::{
    run::{RunSession, RunSessionRequest},
    types::AgentStream,
    AgentError, AgentParams, AgentRequest, AgentResponse,
};
use futures::stream::StreamExt;
use llm_sdk::LanguageModel;
use std::sync::Arc;

pub struct Agent<TCtx> {
    /// A unique name for the agent.
    /// The name can only contain letters and underscores.
    pub name: String,
    params: Arc<AgentParams<TCtx>>,
}

impl<TCtx> Agent<TCtx>
where
    TCtx: Send + Sync + 'static,
{
    #[must_use]
    pub fn new(params: AgentParams<TCtx>) -> Self {
        Self {
            name: params.name.clone(),
            params: Arc::new(params),
        }
    }
    /// Create a one-time run of the agent and generate a response.
    /// A session is created for the run and cleaned up afterwards.
    pub async fn run(&self, request: AgentRequest<TCtx>) -> Result<AgentResponse, AgentError> {
        let AgentRequest { input, context } = request;
        let run_session = self.create_session(context).await?;
        let result = run_session.run(RunSessionRequest { input }).await;
        match result {
            Ok(response) => {
                if let Err(close_err) = run_session.close().await {
                    Err(close_err)
                } else {
                    Ok(response)
                }
            }
            Err(err) => {
                let _ = run_session.close().await;
                Err(err)
            }
        }
    }

    /// Create a one-time streaming run of the agent and generate a response.
    /// A session is created for the run and cleaned up afterwards.
    pub async fn run_stream(&self, request: AgentRequest<TCtx>) -> Result<AgentStream, AgentError> {
        let AgentRequest { input, context } = request;
        let run_session = Arc::new(self.create_session(context).await?);
        let mut stream = match run_session.clone().run_stream(RunSessionRequest { input }) {
            Ok(stream) => stream,
            Err(err) => {
                if let Ok(session) = Arc::try_unwrap(run_session) {
                    let _ = session.close().await;
                }
                return Err(err);
            }
        };

        let wrapped_stream = async_stream::stream! {
            let run_session = run_session;
            while let Some(item) = stream.next().await {
                yield item;
            }
            if let Ok(session) = Arc::try_unwrap(run_session) {
                if let Err(close_err) = session.close().await {
                    yield Err(close_err);
                }
            }
        };

        Ok(AgentStream::from_stream(wrapped_stream))
    }

    /// Create a session for stateful multiple runs of the agent.
    pub async fn create_session(&self, context: TCtx) -> Result<RunSession<TCtx>, AgentError> {
        RunSession::new(self.params.clone(), context).await
    }

    pub fn builder(name: &str, model: Arc<dyn LanguageModel + Send + Sync>) -> AgentParams<TCtx> {
        AgentParams::new(name, model)
    }
}