File size: 3,125 Bytes
e5034c3 | 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 88 89 90 91 92 93 94 | use std::sync::Arc;
use derive_setters::Setters;
use forge_domain::{
ChatCompletionMessageFull, Context, ContextMessage, ConversationId, ModelId, ProviderId,
ReasoningConfig, ResponseFormat, ResultStreamExt, UserPrompt,
};
use schemars::JsonSchema;
use serde::Deserialize;
use crate::TemplateEngine;
use crate::agent::AgentService as AS;
/// Structured response for title generation using JSON format
#[derive(Debug, Clone, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
#[schemars(title = "title")]
pub struct TitleResponse {
/// The generated title for the conversation
pub title: String,
}
/// Service for generating contextually appropriate titles
#[derive(Setters)]
pub struct TitleGenerator<S> {
/// Shared reference to the agent services used for AI interactions
services: Arc<S>,
/// The user prompt to generate a title for
user_prompt: UserPrompt,
/// The model ID to use for title generation
model_id: ModelId,
/// Reasoning configuration for the generator.
reasoning: Option<ReasoningConfig>,
/// The provider ID to use for title generation
provider_id: Option<ProviderId>,
}
impl<S: AS> TitleGenerator<S> {
pub fn new(
services: Arc<S>,
user_prompt: UserPrompt,
model_id: ModelId,
provider_id: Option<ProviderId>,
) -> Self {
Self {
services,
user_prompt,
model_id,
reasoning: None,
provider_id,
}
}
pub async fn generate(&self) -> anyhow::Result<Option<String>> {
let template = TemplateEngine::default().render(
"forge-system-prompt-title-generation.md",
&Default::default(),
)?;
let prompt = format!("<user_prompt>{}</user_prompt>", self.user_prompt.as_str());
// Generate JSON schema from TitleResponse using schemars
let schema = schemars::schema_for!(TitleResponse);
let mut ctx = Context::default()
.temperature(1.0f32)
.conversation_id(ConversationId::generate())
.add_message(ContextMessage::system(template))
.add_message(ContextMessage::user(prompt, Some(self.model_id.clone())))
.response_format(ResponseFormat::JsonSchema(Box::new(schema)));
// Set the reasoning if configured.
if let Some(reasoning) = self.reasoning.as_ref() {
ctx = ctx.reasoning(reasoning.clone());
}
let stream = self
.services
.chat_agent(&self.model_id, ctx, self.provider_id.clone())
.await?;
let ChatCompletionMessageFull { content, .. } = stream.into_full(false).await?;
// Parse the response - try JSON first (structured output), fallback to plain
// text
match serde_json::from_str::<TitleResponse>(&content) {
Ok(response) => Ok(Some(response.title)),
Err(_) => {
// Fallback: Some providers don't support structured output, treat as plain text
Ok(Some(content.trim().to_string()))
}
}
}
}
|