batch_mode_batch_workflow/
language_model_batch_workflow.rs

1// ---------------- [ File: batch-mode-batch-workflow/src/language_model_batch_workflow.rs ]
2crate::ix!();
3
4pub type LanguageModelClientArc = Arc<dyn LanguageModelClientInterface<LanguageModelBatchWorkflowError>>;
5
6/// Two new traits that users must implement:
7/// 1) `ComputeSystemMessage` to provide a static or dynamic system message.
8/// 2) `ComputeLanguageModelCoreQuery` to build requests for each seed item.
9///
10/// These traits are now required components of the overall workflow.
11/// We define them in the same `batch_mode_batch_workflow` (or relevant) crate
12/// so that the derive macro can reference them.
13pub trait ComputeSystemMessage {
14    /// Returns the system message to be applied to all requests.
15    fn system_message() -> String;
16}
17
18pub trait ComputeLanguageModelCoreQuery {
19    /// The seed item type (e.g., AiTomlWriterRequest).
20    type Seed: HasAssociatedOutputName + Named;
21
22    /// Builds a single language model API request for a given seed item.
23    /// The macro will call this once per seed item inside `compute_language_model_requests()`.
24    fn compute_language_model_core_query(
25        &self,
26        input: &Self::Seed
27    ) -> String;
28}
29
30#[async_trait]
31pub trait FinishProcessingUncompletedBatches {
32    type Error;
33
34    /// Possibly complete or discard partial data from prior
35    /// runs.
36    ///
37    async fn finish_processing_uncompleted_batches(
38        &self,
39        expected_content_type: &ExpectedContentType
40    ) -> Result<(), Self::Error>;
41}
42
43/// This is the trait we will typically need to implement manually
44pub trait ComputeLanguageModelRequests {
45
46    type Seed: HasAssociatedOutputName + Send + Sync;
47
48    /// Identify which new items need to be processed and
49    /// build the requests.
50    ///
51    fn compute_language_model_requests(
52        &self,
53        model:            &LanguageModelType,
54        input_tokens:     &[Self::Seed]
55    ) -> Vec<LanguageModelBatchAPIRequest>;
56}
57
58#[async_trait]
59pub trait ProcessBatchRequests {
60
61    type Error;
62
63    /// Process each batch, writing it to disk or sending it
64    /// to a remote server.
65    ///
66    async fn process_batch_requests(
67        &self,
68        batch_requests:        &[LanguageModelBatchAPIRequest],
69        expected_content_type: &ExpectedContentType,
70    ) -> Result<(), Self::Error>;
71}
72
73/// Trait describing a more general “batch workflow”
74/// specialized to GPT expansions.
75///
76/// This approach can unify:
77/// - Reconciling partial/incomplete state,
78/// - Computing new requests,
79/// - Chunking them,
80/// - Sending them to a remote server,
81/// - Handling the results.
82#[async_trait]
83pub trait LanguageModelBatchWorkflow<E: From<LanguageModelBatchCreationError>>: 
84    FinishProcessingUncompletedBatches<Error = E>
85    + ComputeLanguageModelRequests
86    + ProcessBatchRequests<Error = E>
87{
88    const REQUESTS_PER_BATCH: usize = 80;
89
90    async fn plant_seed_and_wait(
91        &mut self,
92        input_tokens: &[<Self as ComputeLanguageModelRequests>::Seed]
93    ) -> Result<(), E>;
94
95    /// High-level method that ties it all together.
96    /// Fixes the mismatch by enumerating chunk-slices properly and passing
97    /// `&[LanguageModelBatchAPIRequest]` to `process_batch_requests`.
98    async fn execute_language_model_batch_workflow(
99        &mut self,
100        model:                 LanguageModelType,
101        expected_content_type: ExpectedContentType,
102        input_tokens:          &[<Self as ComputeLanguageModelRequests>::Seed]
103    ) -> Result<(), E>
104    {
105        info!("Beginning full batch workflow execution");
106
107        self.finish_processing_uncompleted_batches(&expected_content_type).await?;
108
109        let requests: Vec<_> = self.compute_language_model_requests(&model, input_tokens);
110
111        // `construct_batches` presumably returns something like an iterator of chunks.
112        let enumerated_batches = construct_batches(&requests, Self::REQUESTS_PER_BATCH, false)?;
113
114        // Enumerate so we have (batch_idx, chunk_of_requests).
115        for (batch_idx, batch_requests) in enumerated_batches {
116            info!("Processing batch #{}", batch_idx);
117            // Here, `batch_requests` is a `&[LanguageModelBatchAPIRequest]`,
118            // matching the expected parameter type in `process_batch_requests`.
119            self.process_batch_requests(batch_requests, &expected_content_type).await?;
120        }
121
122        Ok(())
123    }
124}
125
126/// This new trait is used to gather the final AI expansions from disk, 
127/// matching each seed item to its parsed output JSON.
128#[async_trait]
129pub trait LanguageModelBatchWorkflowGatherResults {
130    type Error;
131    type Seed: HasAssociatedOutputName + Clone + Named;
132    type Output: LoadFromFile<Error = SaveLoadError>;
133
134    /// Gathers AI-generated JSON outputs for the given seeds in the same order, 
135    /// returning `(Seed, Output)` pairs.
136    async fn gather_results(
137        &self,
138        seeds: &[Self::Seed]
139    ) -> Result<Vec<(Self::Seed, Self::Output)>, Self::Error>;
140}
batch_mode_batch_workflow/language_model_batch_workflow.rs

batch_mode_batch_workflow/
language_model_batch_workflow.rs
