langchain_rust/agent/
agent.rs

1use std::{pin::Pin, sync::Arc};
2
3use async_trait::async_trait;
4use futures::Stream;
5
6use crate::{
7    language_models::TokenUsage,
8    prompt::PromptArgs,
9    schemas::agent::{AgentAction, AgentEvent, AgentFinish},
10    tools::Tool,
11};
12
13use super::AgentError;
14
15/// An event from [`Agent::plan_stream`].
16#[derive(Debug)]
17#[non_exhaustive]
18pub enum PlanEvent {
19    /// Text the model writes while planning, as it arrives.
20    Text(String),
21    /// Tokens used by the model call(s) behind this plan, when the LLM reports them.
22    /// Comes before `Done`.
23    Usage(TokenUsage),
24    /// The plan: run tools or finish. Always the last event.
25    Done(AgentEvent),
26}
27
28/// Stream returned by [`Agent::plan_stream`].
29pub type PlanEventStream<'a> =
30    Pin<Box<dyn Stream<Item = Result<PlanEvent, AgentError>> + Send + 'a>>;
31
32/// Boxes a plan event stream; also pins down the error type for `try_stream!` blocks.
33pub(crate) fn box_plan_events<'a, S>(stream: S) -> PlanEventStream<'a>
34where
35    S: Stream<Item = Result<PlanEvent, AgentError>> + Send + 'a,
36{
37    Box::pin(stream)
38}
39
40#[async_trait]
41pub trait Agent: Send + Sync {
42    async fn plan(
43        &self,
44        intermediate_steps: &[(AgentAction, String)],
45        inputs: PromptArgs,
46    ) -> Result<AgentEvent, AgentError>;
47
48    fn get_tools(&self) -> Vec<Arc<dyn Tool>>;
49
50    /// Like [`plan`](Agent::plan), and also returns the tokens the model call(s) used, so
51    /// [`AgentExecutor`](crate::agent::AgentExecutor) can report a run's total.
52    ///
53    /// The default returns no usage. Agents built on an LLM chain override it.
54    async fn plan_with_usage(
55        &self,
56        intermediate_steps: &[(AgentAction, String)],
57        inputs: PromptArgs,
58    ) -> Result<(AgentEvent, Option<TokenUsage>), AgentError> {
59        Ok((self.plan(intermediate_steps, inputs).await?, None))
60    }
61
62    /// Like [`plan`](Agent::plan), but streams the model's text while it is generated.
63    ///
64    /// The default implementation calls `plan_with_usage` and, when the agent finishes,
65    /// emits the final answer as one `Text` chunk, then the usage (if any), then `Done`.
66    /// Agents whose raw model output is readable text (like
67    /// [`OpenAiToolAgent`](crate::agent::OpenAiToolAgent)) override it to stream token by
68    /// token.
69    fn plan_stream<'a>(
70        &'a self,
71        intermediate_steps: &'a [(AgentAction, String)],
72        inputs: PromptArgs,
73    ) -> PlanEventStream<'a> {
74        box_plan_events(async_stream::try_stream! {
75            let (event, usage) = self.plan_with_usage(intermediate_steps, inputs).await?;
76            if let AgentEvent::Finish(AgentFinish { output }) = &event {
77                if !output.is_empty() {
78                    yield PlanEvent::Text(output.clone());
79                }
80            }
81            if let Some(usage) = usage {
82                yield PlanEvent::Usage(usage);
83            }
84            yield PlanEvent::Done(event);
85        })
86    }
87}
langchain_rust/agent/agent.rs

langchain_rust/agent/
agent.rs
