1use super::config::BehaviorModelConfig;
8use super::context::StatefulAiContext;
9use super::llm_client::LlmClient;
10use super::rule_generator::PaginatedResponse;
11use super::types::LlmGenerationRequest;
12use mockforge_foundation::Result;
13use serde::{Deserialize, Serialize};
14use serde_json::Value;
15use std::collections::HashMap;
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct PaginationRequest {
20 pub path: String,
22 pub query_params: HashMap<String, String>,
24 pub request_body: Option<Value>,
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct PaginationMetadata {
31 pub page: Option<usize>,
33 pub page_size: usize,
35 pub total: usize,
37 pub total_pages: usize,
39 pub has_next: bool,
41 pub has_prev: bool,
43 pub offset: Option<usize>,
45 pub next_cursor: Option<String>,
47 pub prev_cursor: Option<String>,
49}
50
51pub struct PaginationIntelligence {
53 llm_client: Option<LlmClient>,
55 #[allow(dead_code)]
57 config: BehaviorModelConfig,
58 examples: Vec<PaginatedResponse>,
60 default_rule: PaginationRule,
62}
63
64#[derive(Debug, Clone, Serialize, Deserialize)]
66pub struct PaginationRule {
67 pub default_page_size: usize,
69 pub max_page_size: usize,
71 pub min_page_size: usize,
73 pub format: PaginationFormat,
75 pub parameter_names: HashMap<String, String>,
77}
78
79#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
81#[serde(rename_all = "lowercase")]
82pub enum PaginationFormat {
83 PageBased,
85 OffsetBased,
87 CursorBased,
89}
90
91impl PaginationIntelligence {
92 pub fn new(config: BehaviorModelConfig) -> Self {
94 let llm_client = if config.llm_provider != "disabled" {
95 Some(LlmClient::new(config.clone()))
96 } else {
97 None
98 };
99
100 Self {
101 llm_client,
102 config,
103 examples: Vec::new(),
104 default_rule: PaginationRule {
105 default_page_size: 20,
106 max_page_size: 100,
107 min_page_size: 1,
108 format: PaginationFormat::PageBased,
109 parameter_names: HashMap::new(),
110 },
111 }
112 }
113
114 pub fn learn_from_example(&mut self, example: PaginatedResponse) {
116 self.examples.push(example);
117 self.update_rule_from_examples();
119 }
120
121 pub async fn generate_pagination_metadata(
126 &self,
127 request: &PaginationRequest,
128 context: &StatefulAiContext,
129 ) -> Result<PaginationMetadata> {
130 let (page, page_size, offset, _cursor) = self.extract_pagination_params(request);
132
133 let page_size = page_size.unwrap_or_else(|| self.infer_page_size(request, &self.examples));
135
136 let total = self.generate_realistic_total(context, request).await?;
138
139 let total_pages = total.div_ceil(page_size); let current_page = page.unwrap_or(1);
142 let has_next = current_page < total_pages;
143 let has_prev = current_page > 1;
144
145 let (next_cursor, prev_cursor) =
147 if self.default_rule.format == PaginationFormat::CursorBased {
148 (
149 if has_next {
150 Some(self.generate_cursor(current_page + 1))
151 } else {
152 None
153 },
154 if has_prev {
155 Some(self.generate_cursor(current_page - 1))
156 } else {
157 None
158 },
159 )
160 } else {
161 (None, None)
162 };
163
164 let calculated_offset = if self.default_rule.format == PaginationFormat::OffsetBased {
166 Some(offset.unwrap_or_else(|| (current_page - 1) * page_size))
167 } else {
168 offset
169 };
170
171 Ok(PaginationMetadata {
172 page: Some(current_page),
173 page_size,
174 total,
175 total_pages,
176 has_next,
177 has_prev,
178 offset: calculated_offset,
179 next_cursor,
180 prev_cursor,
181 })
182 }
183
184 pub fn infer_page_size(
186 &self,
187 request: &PaginationRequest,
188 examples: &[PaginatedResponse],
189 ) -> usize {
190 for (key, value) in &request.query_params {
192 if matches!(key.to_lowercase().as_str(), "limit" | "per_page" | "page_size" | "size") {
193 if let Ok(size) = value.parse::<usize>() {
194 return size
195 .clamp(self.default_rule.min_page_size, self.default_rule.max_page_size);
196 }
197 }
198 }
199
200 if let Some(most_common) = self.find_most_common_page_size(examples) {
202 return most_common;
203 }
204
205 self.default_rule.default_page_size
207 }
208
209 pub async fn generate_realistic_total(
211 &self,
212 context: &StatefulAiContext,
213 request: &PaginationRequest,
214 ) -> Result<usize> {
215 if let Some(ref _llm_client) = self.llm_client {
217 return self.generate_total_with_llm(context, request).await;
218 }
219
220 Ok(self.generate_total_heuristic(context, request))
222 }
223
224 fn extract_pagination_params(
228 &self,
229 request: &PaginationRequest,
230 ) -> (Option<usize>, Option<usize>, Option<usize>, Option<String>) {
231 let mut page = None;
232 let mut page_size = None;
233 let mut offset = None;
234 let mut cursor = None;
235
236 for (key, value) in &request.query_params {
237 match key.to_lowercase().as_str() {
238 "page" | "p" => {
239 if let Ok(p) = value.parse::<usize>() {
240 page = Some(p);
241 }
242 }
243 "limit" | "per_page" | "page_size" | "size" => {
244 if let Ok(size) = value.parse::<usize>() {
245 page_size = Some(size);
246 }
247 }
248 "offset" => {
249 if let Ok(o) = value.parse::<usize>() {
250 offset = Some(o);
251 }
252 }
253 "cursor" => {
254 cursor = Some(value.clone());
255 }
256 _ => {}
257 }
258 }
259
260 (page, page_size, offset, cursor)
261 }
262
263 fn find_most_common_page_size(&self, examples: &[PaginatedResponse]) -> Option<usize> {
265 let mut size_counts: HashMap<usize, usize> = HashMap::new();
266
267 for example in examples {
268 if let Some(size) = example.page_size {
269 *size_counts.entry(size).or_insert(0) += 1;
270 }
271 }
272
273 size_counts.into_iter().max_by_key(|(_, count)| *count).map(|(size, _)| size)
274 }
275
276 fn update_rule_from_examples(&mut self) {
278 if self.examples.is_empty() {
279 return;
280 }
281
282 let page_sizes: Vec<usize> = self.examples.iter().filter_map(|e| e.page_size).collect();
284
285 if !page_sizes.is_empty() {
286 self.default_rule.default_page_size = *page_sizes.iter().min().unwrap();
287 self.default_rule.max_page_size = *page_sizes.iter().max().unwrap();
288 }
289
290 let mut has_offset = false;
292 let mut has_cursor = false;
293
294 for example in &self.examples {
295 for key in example.query_params.keys() {
296 match key.to_lowercase().as_str() {
297 "offset" => has_offset = true,
298 "cursor" => has_cursor = true,
299 "page" | "p" => {}
300 _ => {}
301 }
302 }
303 }
304
305 self.default_rule.format = if has_cursor {
306 PaginationFormat::CursorBased
307 } else if has_offset {
308 PaginationFormat::OffsetBased
309 } else {
310 PaginationFormat::PageBased
311 };
312 }
313
314 async fn generate_total_with_llm(
316 &self,
317 context: &StatefulAiContext,
318 request: &PaginationRequest,
319 ) -> Result<usize> {
320 let llm_client = self
321 .llm_client
322 .as_ref()
323 .ok_or_else(|| mockforge_foundation::Error::internal("LLM client not available"))?;
324
325 let context_summary = context.build_context_summary().await;
327 let request_summary = format!("Path: {}, Query: {:?}", request.path, request.query_params);
328
329 let system_prompt = "You are a pagination metadata generator. Generate realistic total item counts for API responses.";
330 let user_prompt = format!(
331 "Based on this API request context, generate a realistic total item count:\n\n{}\n\n{}\n\nReturn only a number between 0 and 10000. Consider the context and make it realistic.",
332 context_summary,
333 request_summary
334 );
335
336 let request_llm = LlmGenerationRequest {
337 system_prompt: system_prompt.to_string(),
338 user_prompt,
339 temperature: 0.5, max_tokens: 50,
341 schema: None,
342 };
343
344 let response = llm_client.generate(&request_llm).await?;
345
346 if let Some(num_str) = response.as_str() {
348 if let Some(num) =
350 num_str.split_whitespace().find_map(|word| word.parse::<usize>().ok())
351 {
352 return Ok(num.clamp(0, 10000));
353 }
354 }
355
356 Ok(self.generate_total_heuristic(context, request))
358 }
359
360 fn generate_total_heuristic(
362 &self,
363 _context: &StatefulAiContext,
364 _request: &PaginationRequest,
365 ) -> usize {
366 use std::collections::hash_map::DefaultHasher;
374 use std::hash::{Hash, Hasher};
375
376 let mut hasher = DefaultHasher::new();
377 _request.path.hash(&mut hasher);
378 let hash = hasher.finish();
379
380 let base = 50;
382 let range = 450;
383
384 base + (hash % range as u64) as usize
385 }
386
387 fn generate_cursor(&self, page: usize) -> String {
389 format!("cursor_{}", page)
392 }
393}
394
395impl Default for PaginationIntelligence {
396 fn default() -> Self {
397 Self::new(BehaviorModelConfig::default())
398 }
399}
400
401#[cfg(test)]
402mod tests {
403 use super::*;
404 use serde_json::json;
405
406 #[tokio::test]
407 async fn test_extract_pagination_params() {
408 let config = BehaviorModelConfig::default();
409 let intelligence = PaginationIntelligence::new(config);
410
411 let mut query_params = HashMap::new();
412 query_params.insert("page".to_string(), "2".to_string());
413 query_params.insert("limit".to_string(), "25".to_string());
414
415 let request = PaginationRequest {
416 path: "/api/users".to_string(),
417 query_params,
418 request_body: None,
419 };
420
421 let (page, page_size, offset, cursor) = intelligence.extract_pagination_params(&request);
422
423 assert_eq!(page, Some(2));
424 assert_eq!(page_size, Some(25));
425 assert_eq!(offset, None);
426 assert_eq!(cursor, None);
427 }
428
429 #[tokio::test]
430 async fn test_infer_page_size() {
431 let config = BehaviorModelConfig::default();
432 let intelligence = PaginationIntelligence::new(config);
433
434 let mut query_params = HashMap::new();
435 query_params.insert("limit".to_string(), "30".to_string());
436
437 let request = PaginationRequest {
438 path: "/api/users".to_string(),
439 query_params,
440 request_body: None,
441 };
442
443 let examples = vec![PaginatedResponse {
444 path: "/api/users".to_string(),
445 query_params: HashMap::new(),
446 response: json!({}),
447 page: Some(1),
448 page_size: Some(20),
449 total: Some(100),
450 }];
451
452 let page_size = intelligence.infer_page_size(&request, &examples);
453 assert_eq!(page_size, 30); }
455
456 #[test]
457 fn test_find_most_common_page_size() {
458 let config = BehaviorModelConfig::default();
459 let intelligence = PaginationIntelligence::new(config);
460
461 let examples = vec![
462 PaginatedResponse {
463 path: "/api/users".to_string(),
464 query_params: HashMap::new(),
465 response: json!({}),
466 page: Some(1),
467 page_size: Some(20),
468 total: Some(100),
469 },
470 PaginatedResponse {
471 path: "/api/users".to_string(),
472 query_params: HashMap::new(),
473 response: json!({}),
474 page: Some(2),
475 page_size: Some(20),
476 total: Some(100),
477 },
478 PaginatedResponse {
479 path: "/api/users".to_string(),
480 query_params: HashMap::new(),
481 response: json!({}),
482 page: Some(1),
483 page_size: Some(50),
484 total: Some(200),
485 },
486 ];
487
488 let most_common = intelligence.find_most_common_page_size(&examples);
489 assert_eq!(most_common, Some(20)); }
491}