tea_context/providers/
workspace.rs1use std::str::FromStr;
2
3use crate::{
4 BudgetBehavior, CacheScope, ContextError, ContextErrorCode, ContextProvider,
5 ContextProviderFuture, ContextProviderId, ContextRequest, PromptAuthority, PromptModule,
6 PromptModuleId, PromptPriority, PromptProvenance, PromptSegment, PromptSegmentId, TrustLevel,
7};
8
9pub const MAX_WORKSPACE_INSTRUCTIONS: usize = 128;
11
12#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct WorkspaceInstruction {
15 id: PromptSegmentId,
16 content: String,
17 locator: String,
18 trust: TrustLevel,
19}
20
21impl WorkspaceInstruction {
22 pub fn new(
28 id: PromptSegmentId,
29 content: impl Into<String>,
30 locator: impl Into<String>,
31 trust: TrustLevel,
32 ) -> Result<Self, ContextError> {
33 let content = content.into();
34 let locator = locator.into();
35 if locator.is_empty() || locator.len() > 2048 || locator.chars().any(char::is_control) {
36 return Err(ContextError::new(
37 ContextErrorCode::InvalidValue,
38 "workspace instruction locator is invalid",
39 ));
40 }
41 PromptSegment::new(
43 id.clone(),
44 content.clone(),
45 PromptProvenance::new(
46 ContextProviderId::from_str("builtin.workspace_instructions")
47 .map_err(value_error)?,
48 "workspace_file",
49 Some(locator.clone()),
50 )
51 .map_err(value_error)?,
52 trust,
53 CacheScope::Session,
54 BudgetBehavior::Omit,
55 )
56 .map_err(value_error)?;
57 Ok(Self {
58 id,
59 content,
60 locator,
61 trust,
62 })
63 }
64}
65
66#[derive(Debug, Clone)]
68pub struct WorkspaceInstructionProvider {
69 id: ContextProviderId,
70 instructions: Vec<WorkspaceInstruction>,
71}
72
73impl WorkspaceInstructionProvider {
74 pub fn new(mut instructions: Vec<WorkspaceInstruction>) -> Result<Self, ContextError> {
80 if instructions.len() > MAX_WORKSPACE_INSTRUCTIONS {
81 return Err(ContextError::new(
82 ContextErrorCode::BoundsExceeded,
83 "workspace instruction collection is too large",
84 ));
85 }
86 instructions.sort_by(|left, right| left.id.cmp(&right.id));
87 if instructions
88 .windows(2)
89 .any(|items| items[0].id == items[1].id)
90 {
91 return Err(ContextError::new(
92 ContextErrorCode::DuplicateIdentity,
93 "workspace instruction ID is duplicated",
94 ));
95 }
96 Ok(Self {
97 id: ContextProviderId::from_str("builtin.workspace_instructions")
98 .map_err(value_error)?,
99 instructions,
100 })
101 }
102}
103
104impl ContextProvider for WorkspaceInstructionProvider {
105 fn id(&self) -> &ContextProviderId {
106 &self.id
107 }
108
109 fn provide(&self, _request: ContextRequest) -> ContextProviderFuture<'_> {
110 let id = self.id.clone();
111 let instructions = self.instructions.clone();
112 Box::pin(async move {
113 if instructions.is_empty() {
114 return Ok(Vec::new());
115 }
116 let segments = instructions
117 .into_iter()
118 .map(|instruction| {
119 PromptSegment::new(
120 instruction.id,
121 instruction.content,
122 PromptProvenance::new(
123 id.clone(),
124 "workspace_file",
125 Some(instruction.locator),
126 )
127 .map_err(value_error)?,
128 instruction.trust,
129 CacheScope::Session,
130 BudgetBehavior::Omit,
131 )
132 .map_err(value_error)
133 })
134 .collect::<Result<Vec<_>, ContextError>>()?;
135 Ok(vec![
136 PromptModule::new(
137 PromptModuleId::from_str("workspace.instructions").map_err(value_error)?,
138 PromptAuthority::Workspace,
139 PromptPriority::new(0),
140 segments,
141 )
142 .map_err(value_error)?,
143 ])
144 })
145 }
146}
147
148fn value_error(error: impl std::fmt::Display) -> ContextError {
149 ContextError::new(ContextErrorCode::InvalidValue, error.to_string())
150}