Skip to main content

chronicle/agent/
mod.rs

1pub mod prompt;
2pub mod tools;
3
4use snafu::ResultExt;
5
6use crate::annotate::gather::AnnotationContext;
7use crate::error::agent_error::{MaxTurnsExceededSnafu, NoAnnotationsSnafu, ProviderSnafu};
8use crate::error::AgentError;
9use crate::git::GitOps;
10use crate::provider::{CompletionRequest, ContentBlock, LlmProvider, Message, Role, StopReason};
11use crate::schema::{CrossCuttingConcern, RegionAnnotation};
12
13const MAX_TURNS: u32 = 20;
14
15/// Run the annotation agent loop. Calls the LLM with tools until it finishes
16/// or hits the turn limit. Returns collected annotations.
17pub fn run_agent_loop(
18    provider: &dyn LlmProvider,
19    git_ops: &dyn GitOps,
20    context: &AnnotationContext,
21) -> Result<(Vec<RegionAnnotation>, Vec<CrossCuttingConcern>, String), AgentError> {
22    let system_prompt = prompt::build_system_prompt(context);
23    let tool_defs = tools::tool_definitions();
24
25    let mut messages = vec![Message {
26        role: Role::User,
27        content: vec![ContentBlock::Text {
28            text: "Please annotate this commit.".to_string(),
29        }],
30    }];
31
32    let mut collected_regions: Vec<RegionAnnotation> = Vec::new();
33    let mut collected_cross_cutting: Vec<CrossCuttingConcern> = Vec::new();
34    let mut summary = String::new();
35
36    for turn in 0..MAX_TURNS {
37        let request = CompletionRequest {
38            system: system_prompt.clone(),
39            messages: messages.clone(),
40            tools: tool_defs.clone(),
41            max_tokens: 4096,
42        };
43
44        let response = provider.complete(&request).context(ProviderSnafu)?;
45
46        // Collect any text from the response as potential summary
47        let mut assistant_text = String::new();
48        let mut tool_uses: Vec<(String, String, serde_json::Value)> = Vec::new();
49
50        for block in &response.content {
51            match block {
52                ContentBlock::Text { text } => {
53                    assistant_text.push_str(text);
54                }
55                ContentBlock::ToolUse { id, name, input } => {
56                    tool_uses.push((id.clone(), name.clone(), input.clone()));
57                }
58                _ => {}
59            }
60        }
61
62        // Add the assistant message to history
63        messages.push(Message {
64            role: Role::Assistant,
65            content: response.content.clone(),
66        });
67
68        // If stop reason is EndTurn or MaxTokens with no tool uses, we're done
69        if tool_uses.is_empty() {
70            summary = assistant_text;
71            break;
72        }
73
74        // Process tool uses
75        let mut tool_results: Vec<ContentBlock> = Vec::new();
76        for (id, name, input) in &tool_uses {
77            let result = tools::dispatch_tool(
78                name,
79                input,
80                git_ops,
81                context,
82                &mut collected_regions,
83                &mut collected_cross_cutting,
84            );
85            match result {
86                Ok(content) => {
87                    tool_results.push(ContentBlock::ToolResult {
88                        tool_use_id: id.clone(),
89                        content,
90                        is_error: None,
91                    });
92                }
93                Err(e) => {
94                    tool_results.push(ContentBlock::ToolResult {
95                        tool_use_id: id.clone(),
96                        content: format!("Error: {e}"),
97                        is_error: Some(true),
98                    });
99                }
100            }
101        }
102
103        // Add tool results as a user message
104        messages.push(Message {
105            role: Role::User,
106            content: tool_results,
107        });
108
109        // Check stop conditions
110        if response.stop_reason == StopReason::EndTurn {
111            summary = assistant_text;
112            break;
113        }
114        if response.stop_reason == StopReason::MaxTokens {
115            summary = assistant_text;
116            break;
117        }
118
119        if turn + 1 >= MAX_TURNS {
120            return MaxTurnsExceededSnafu { turns: MAX_TURNS }.fail();
121        }
122    }
123
124    if collected_regions.is_empty() {
125        return NoAnnotationsSnafu.fail();
126    }
127
128    // Use the last text block as summary, fall back to a default
129    if summary.is_empty() {
130        summary = "Annotation complete.".to_string();
131    }
132
133    Ok((collected_regions, collected_cross_cutting, summary))
134}