wuenlp_tools.utils.batch_prompting

  1"""
  2Warning: This is fully AI generated and not checked at all
  3"""
  4
  5from __future__ import annotations
  6
  7import json
  8import os
  9import pickle
 10import tempfile
 11import time
 12import uuid
 13from typing import TYPE_CHECKING
 14
 15import requests
 16from loguru import logger
 17from openai import NOT_GIVEN, OpenAI
 18
 19from wuenlp_tools.keys import OPENAI_API_KEY
 20
 21from . import prompting
 22from .model_cost import get_cost
 23from .prompting import YesNo, YesNoReason
 24
 25if TYPE_CHECKING:
 26    from .prompting import LLM, LLMArchitecture, YesNoLLM
 27
 28logger.warning("This file is fully AI generated and not checked at all!")
 29
 30
 31class BatchRequest:
 32    """Represents a single request in a batch."""
 33
 34    def __init__(
 35        self,
 36        custom_id: str,
 37        prompt: str,
 38        system_prompt: str = "",
 39        endpoint: str = "/v1/chat/completions",
 40    ):
 41        self.custom_id = custom_id
 42        self.prompt = prompt
 43        self.system_prompt = system_prompt
 44        self.endpoint = endpoint
 45
 46    def to_jsonl_line(self, model: str, seed=NOT_GIVEN) -> str:
 47        """Convert to JSONL format for batch processing."""
 48        system_role = "system" if not model.startswith(("o3", "o4")) else "developer"
 49        body = {
 50            "model": model,
 51            "messages": [
 52                {"role": system_role, "content": self.system_prompt},
 53                {"role": "user", "content": self.prompt},
 54            ],
 55        }
 56        if seed != NOT_GIVEN:
 57            body["seed"] = seed
 58
 59        request = {
 60            "custom_id": self.custom_id,
 61            "method": "POST",
 62            "url": self.endpoint,
 63            "body": body,
 64        }
 65        return json.dumps(request)
 66
 67
 68class BatchLLM:
 69    """LLM class that supports OpenAI's Batch API for asynchronous processing."""
 70
 71    def __init__(
 72        self,
 73        model: LLMArchitecture,
 74        system_prompt: str,
 75        seed=NOT_GIVEN,
 76        endpoint: str = "/v1/chat/completions",
 77    ):
 78        if not model.openai:
 79            raise ValueError("Batch API is only supported for OpenAI models")
 80
 81        self.openai = OpenAI(api_key=OPENAI_API_KEY)
 82        self.model = model
 83        self.system_prompt = system_prompt
 84        self.seed = seed
 85        self.endpoint = endpoint
 86        self.requests = []
 87
 88    def add_request(self, prompt: str, custom_id: str = None) -> str:
 89        """Add a request to the batch. Returns the custom_id."""
 90        if custom_id is None:
 91            custom_id = str(uuid.uuid4())
 92
 93        # Truncate prompt if too long
 94        model_max_chars = self.model.estimated_num_characters
 95        if model_max_chars > -1 and len(prompt) > model_max_chars:
 96            prompt = prompt[:model_max_chars]
 97            logger.warning(
 98                f"Truncated prompt to {model_max_chars} characters for request {custom_id}"
 99            )
100
101        request = BatchRequest(custom_id, prompt, self.system_prompt, self.endpoint)
102        self.requests.append(request)
103        return custom_id
104
105    def clear_requests(self):
106        """Clear all pending requests."""
107        self.requests = []
108
109    def create_batch_file(self) -> str:
110        """Create a temporary JSONL file with all requests."""
111        if not self.requests:
112            raise ValueError("No requests to process")
113
114        # Create temporary file
115        with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f:
116            for request in self.requests:
117                f.write(request.to_jsonl_line(self.model.model, self.seed) + "\n")
118            temp_file_path = f.name
119
120        logger.info(
121            f"Created batch file with {len(self.requests)} requests: {temp_file_path}"
122        )
123        return temp_file_path
124
125    def submit_batch(
126        self, completion_window: str = "24h", metadata: dict = None
127    ) -> str:
128        """Submit the batch for processing. Returns batch_id."""
129        if not self.requests:
130            raise ValueError("No requests to submit")
131
132        # Create and upload batch file
133        batch_file_path = self.create_batch_file()
134
135        try:
136            # Upload file
137            with open(batch_file_path, "rb") as f:
138                batch_input_file = self.openai.files.create(file=f, purpose="batch")
139
140            # Create batch
141            batch = self.openai.batches.create(
142                input_file_id=batch_input_file.id,
143                endpoint=self.endpoint,
144                completion_window=completion_window,
145                metadata=metadata or {},
146            )
147
148            logger.info(
149                f"Submitted batch {batch.id} with {len(self.requests)} requests"
150            )
151            return batch.id
152
153        finally:
154            # Clean up temporary file
155            try:
156                os.unlink(batch_file_path)
157            except OSError:
158                pass
159
160    def get_batch_status(self, batch_id: str) -> dict:
161        """Get the status of a batch."""
162        batch = self.openai.batches.retrieve(batch_id)
163        return {
164            "id": batch.id,
165            "status": batch.status,
166            "request_counts": batch.request_counts,
167            "created_at": batch.created_at,
168            "completed_at": batch.completed_at,
169            "failed_at": batch.failed_at,
170            "expires_at": batch.expires_at,
171            "output_file_id": batch.output_file_id,
172            "error_file_id": batch.error_file_id,
173        }
174
175    def wait_for_completion(self, batch_id: str, poll_interval: int = 60) -> dict:
176        """Wait for batch to complete, polling at specified interval."""
177        logger.info(f"Waiting for batch {batch_id} to complete...")
178
179        while True:
180            status = self.get_batch_status(batch_id)
181            logger.debug(f"Batch {batch_id} status: {status['status']}")
182
183            if status["status"] in ["completed", "failed", "expired", "cancelled"]:
184                logger.info(
185                    f"Batch {batch_id} finished with status: {status['status']}"
186                )
187                return status
188
189            time.sleep(poll_interval)
190
191    def get_results(self, batch_id: str) -> dict:
192        """Retrieve results from a completed batch."""
193        status = self.get_batch_status(batch_id)
194
195        if status["status"] != "completed":
196            raise ValueError(
197                f"Batch {batch_id} is not completed (status: {status['status']})"
198            )
199
200        if not status["output_file_id"]:
201            raise ValueError(f"No output file available for batch {batch_id}")
202
203        # Download results
204        file_response = self.openai.files.content(status["output_file_id"])
205        results_text = file_response.text
206
207        # Parse results and map by custom_id
208        results = {}
209        errors = {}
210
211        for line in results_text.strip().split("\n"):
212            if not line:
213                continue
214
215            result = json.loads(line)
216            custom_id = result["custom_id"]
217
218            if result.get("error"):
219                errors[custom_id] = result["error"]
220            else:
221                response_body = result["response"]["body"]
222                content = response_body["choices"][0]["message"]["content"]
223                results[custom_id] = {
224                    "content": content,
225                    "usage": response_body.get("usage", {}),
226                    "model": response_body.get("model"),
227                    "finish_reason": response_body["choices"][0].get("finish_reason"),
228                }
229
230        # Get error results if any
231        if status.get("error_file_id"):
232            error_response = self.openai.files.content(status["error_file_id"])
233            error_text = error_response.text
234
235            for line in error_text.strip().split("\n"):
236                if not line:
237                    continue
238
239                error_result = json.loads(line)
240                custom_id = error_result["custom_id"]
241                errors[custom_id] = error_result["error"]
242
243        # Calculate total cost
244        total_cost = 0
245        for custom_id, result in results.items():
246            if "usage" in result:
247                # Create a mock response object for cost calculation
248                class MockResponse:
249                    def __init__(self, usage_dict):
250                        self.usage = type("Usage", (), usage_dict)()
251
252                mock_response = MockResponse(result["usage"])
253                cost = get_cost(mock_response, self.model.model)
254                total_cost += cost
255
256        # Update global run cost
257        prompting.RUN_COST += total_cost
258
259        logger.info(
260            f"Retrieved {len(results)} successful results and {len(errors)} errors from batch {batch_id}"
261        )
262        logger.info(f"Batch cost: {total_cost:.4f} €")
263
264        return {
265            "results": results,
266            "errors": errors,
267            "total_cost": total_cost,
268            "status": status,
269        }
270
271    def cancel_batch(self, batch_id: str):
272        """Cancel a running batch."""
273        batch = self.openai.batches.cancel(batch_id)
274        logger.info(f"Cancelled batch {batch_id}")
275        return batch
276
277    def list_batches(self, limit: int = 20):
278        """List all batches."""
279        return self.openai.batches.list(limit=limit)
280
281    def process_batch(
282        self,
283        prompts: list,
284        custom_ids: list = None,
285        completion_window: str = "24h",
286        metadata: dict = None,
287        wait: bool = True,
288        poll_interval: int = 60,
289    ) -> dict:
290        """Complete batch workflow: add requests, submit, wait, and return results."""
291        # Clear any existing requests
292        self.clear_requests()
293
294        # Add all requests
295        if custom_ids is None:
296            custom_ids = [None] * len(prompts)
297
298        added_ids = []
299        for prompt, custom_id in zip(prompts, custom_ids):
300            added_id = self.add_request(prompt, custom_id)
301            added_ids.append(added_id)
302
303        # Submit batch
304        batch_id = self.submit_batch(completion_window, metadata)
305
306        if not wait:
307            return {"batch_id": batch_id, "custom_ids": added_ids}
308
309        # Wait for completion and return results
310        self.wait_for_completion(batch_id, poll_interval)
311        return self.get_results(batch_id)
312
313    def submit_batch_async(
314        self,
315        prompts: list,
316        custom_ids: list = None,
317        completion_window: str = "24h",
318        metadata: dict = None,
319        state_file: str = None,
320    ) -> dict:
321        """Submit batch without waiting, optionally save state to file."""
322        # Clear any existing requests
323        self.clear_requests()
324
325        # Add all requests
326        if custom_ids is None:
327            custom_ids = [None] * len(prompts)
328
329        added_ids = []
330        for prompt, custom_id in zip(prompts, custom_ids):
331            added_id = self.add_request(prompt, custom_id)
332            added_ids.append(added_id)
333
334        # Submit batch
335        batch_id = self.submit_batch(completion_window, metadata)
336
337        result = {"batch_id": batch_id, "custom_ids": added_ids}
338
339        # Optionally save state
340        if state_file:
341            self.save_batch_state(batch_id, state_file)
342
343        return result
344
345    def check_and_retrieve_results(self, batch_id: str) -> dict:
346        """Check batch status and retrieve results if completed."""
347        status = self.get_batch_status(batch_id)
348
349        if status["status"] in ["completed", "failed", "expired", "cancelled"]:
350            if status["status"] == "completed":
351                return self.get_results(batch_id)
352            else:
353                return {
354                    "status": status["status"],
355                    "error": f"Batch finished with status: {status['status']}",
356                    "results": {},
357                    "errors": {},
358                }
359        else:
360            return {
361                "status": status["status"],
362                "pending": True,
363                "results": {},
364                "errors": {},
365            }
366
367    def save_batch_state(self, batch_id: str, state_file: str):
368        """Save batch state to file for later retrieval."""
369        state = {
370            "batch_id": batch_id,
371            "timestamp": time.time(),
372            "client_config": {"api_key": self.client.api_key, "model": self.model},
373        }
374
375        os.makedirs(os.path.dirname(state_file), exist_ok=True)
376        with open(state_file, "wb") as f:
377            pickle.dump(state, f)
378
379        logger.info(f"Batch state saved to {state_file}")
380
381    def load_batch_state(self, state_file: str) -> dict:
382        """Load batch state from file."""
383        with open(state_file, "rb") as f:
384            state = pickle.load(f)
385
386        logger.info(f"Batch state loaded from {state_file}")
387        return state
388
389
390class BatchLLMAdapter:
391    """Adapter that converts any LLM instance to use batch processing."""
392
393    def __init__(self, llm_instance: LLM):
394        if not llm_instance.model.openai:
395            raise ValueError("Batch API is only supported for OpenAI models")
396
397        self.original_llm = llm_instance
398        self.batch_llm = BatchLLM(
399            model=llm_instance.model,
400            system_prompt=llm_instance.system_prompt,
401            seed=llm_instance.seed,
402            endpoint="/v1/chat/completions",
403        )
404
405    def add_request(self, prompt: str, custom_id: str = None) -> str:
406        """Add a request to the batch."""
407        return self.batch_llm.add_request(prompt, custom_id)
408
409    def clear_requests(self):
410        """Clear all pending requests."""
411        self.batch_llm.clear_requests()
412
413    def submit_batch(
414        self, completion_window: str = "24h", metadata: dict = None
415    ) -> str:
416        """Submit the batch for processing."""
417        return self.batch_llm.submit_batch(completion_window, metadata)
418
419    def get_batch_status(self, batch_id: str) -> dict:
420        """Get the status of a batch."""
421        return self.batch_llm.get_batch_status(batch_id)
422
423    def wait_for_completion(self, batch_id: str, poll_interval: int = 60) -> dict:
424        """Wait for batch to complete."""
425        return self.batch_llm.wait_for_completion(batch_id, poll_interval)
426
427    def cancel_batch(self, batch_id: str):
428        """Cancel a running batch."""
429        return self.batch_llm.cancel_batch(batch_id)
430
431    def list_batches(self, limit: int = 20):
432        """List all batches."""
433        return self.batch_llm.list_batches(limit)
434
435    def get_results(self, batch_id: str) -> dict:
436        """Retrieve results from a completed batch with post-processing."""
437        raw_results = self.batch_llm.get_results(batch_id)
438
439        # Apply post-processing based on original LLM type
440        processed_results = {}
441        processed_errors = raw_results["errors"]
442
443        for custom_id, result in raw_results["results"].items():
444            try:
445                processed_result = self._post_process_result(result["content"])
446                processed_results[custom_id] = {
447                    **result,
448                    "content": result["content"],
449                    "processed": processed_result,
450                }
451            except Exception as e:
452                logger.warning(f"Post-processing failed for {custom_id}: {e}")
453                processed_errors[custom_id] = {
454                    "message": f"Post-processing failed: {e}"
455                }
456
457        return {**raw_results, "results": processed_results, "errors": processed_errors}
458
459    def _post_process_result(self, content: str):
460        """Apply the same post-processing logic as the original LLM."""
461        # Import YesNoLLM at runtime to avoid circular imports
462        from .prompting import YesNoLLM
463
464        if isinstance(self.original_llm, YesNoLLM):
465            return self._post_process_yes_no(content)
466        else:
467            # For other LLM types, return content as-is
468            return content
469
470    def _post_process_yes_no(self, content: str):
471        """Post-process results for YesNoLLM."""
472
473        if self.original_llm.model.openai:
474            # For OpenAI models, structured output should already be parsed
475            try:
476                parsed = json.loads(content) if isinstance(content, str) else content
477                if self.original_llm.reason:
478                    return YesNoReason(**parsed)
479                else:
480                    return YesNo(**parsed)
481            except:
482                pass
483
484        # Fallback to original parsing logic
485        content_lower = content.lower().strip("*").strip()
486
487        if self.original_llm.is_positive(
488            content_lower
489        ) or self.original_llm.is_negative(content_lower):
490            answer = self.original_llm.is_positive(content_lower)
491            if self.original_llm.reason:
492                return YesNoReason(answer=answer, reason=content)
493            else:
494                return YesNo(answer=answer)
495        else:
496            raise ValueError(f"Could not parse yes/no answer from: {content}")
497
498    def process_batch(
499        self,
500        prompts: list,
501        custom_ids: list = None,
502        completion_window: str = "24h",
503        metadata: dict = None,
504        wait: bool = True,
505        poll_interval: int = 60,
506    ) -> dict:
507        """Complete batch workflow with post-processing."""
508        # Add post_prompt to each prompt if needed
509        processed_prompts = []
510        for prompt in prompts:
511            full_prompt = prompt + (self.original_llm.post_prompt or "")
512            processed_prompts.append(full_prompt)
513
514        # Use the underlying batch LLM to process
515        raw_results = self.batch_llm.process_batch(
516            processed_prompts,
517            custom_ids,
518            completion_window,
519            metadata,
520            wait,
521            poll_interval,
522        )
523
524        if not wait:
525            return raw_results
526
527        # Apply post-processing to results
528        processed_results = {}
529        processed_errors = raw_results["errors"]
530
531        for custom_id, result in raw_results["results"].items():
532            try:
533                processed_result = self._post_process_result(result["content"])
534                processed_results[custom_id] = {
535                    **result,
536                    "content": result["content"],
537                    "processed": processed_result,
538                }
539            except Exception as e:
540                logger.warning(f"Post-processing failed for {custom_id}: {e}")
541                processed_errors[custom_id] = {
542                    "message": f"Post-processing failed: {e}"
543                }
544
545        return {**raw_results, "results": processed_results, "errors": processed_errors}
546
547    def submit_batch_async(
548        self,
549        prompts: list,
550        custom_ids: list = None,
551        completion_window: str = "24h",
552        metadata: dict = None,
553        state_file: str = None,
554    ) -> dict:
555        """Submit batch without waiting, with post-processing support."""
556        # Add post_prompt to each prompt if needed
557        processed_prompts = []
558        for prompt in prompts:
559            full_prompt = prompt + (self.original_llm.post_prompt or "")
560            processed_prompts.append(full_prompt)
561
562        # Submit via underlying batch LLM
563        return self.batch_llm.submit_batch_async(
564            processed_prompts,
565            custom_ids,
566            completion_window,
567            metadata,
568            state_file,
569        )
570
571    def check_and_retrieve_results(self, batch_id: str) -> dict:
572        """Check batch status and retrieve results with post-processing."""
573        raw_results = self.batch_llm.check_and_retrieve_results(batch_id)
574
575        # If not completed or has error, return as-is
576        if raw_results.get("pending") or raw_results.get("error"):
577            return raw_results
578
579        # Apply post-processing to completed results
580        if raw_results.get("results"):
581            processed_results = {}
582            processed_errors = raw_results["errors"]
583
584            for custom_id, result in raw_results["results"].items():
585                try:
586                    processed_result = self._post_process_result(result["content"])
587                    processed_results[custom_id] = {
588                        **result,
589                        "content": result["content"],
590                        "processed": processed_result,
591                    }
592                except Exception as e:
593                    logger.warning(f"Post-processing failed for {custom_id}: {e}")
594                    processed_errors[custom_id] = {
595                        "message": f"Post-processing failed: {e}"
596                    }
597
598            return {
599                **raw_results,
600                "results": processed_results,
601                "errors": processed_errors,
602            }
603
604        return raw_results
605
606    def save_batch_state(self, batch_id: str, state_file: str):
607        """Save batch state to file."""
608        return self.batch_llm.save_batch_state(batch_id, state_file)
609
610    def load_batch_state(self, state_file: str) -> dict:
611        """Load batch state from file."""
612        return self.batch_llm.load_batch_state(state_file)
613
614
615def make_batch_llm(llm: LLM) -> BatchLLMAdapter:
616    """Factory function to convert any LLM instance to batch processing.
617
618    Args:
619        llm: Any LLM instance (LLM, SentenceInContextLLM, YesNoLLM, etc.)
620
621    Returns:
622        BatchLLMAdapter that preserves the original LLM's behavior
623
624    Usage:
625        # Convert existing LLM to batch version
626        yes_no_llm = YesNoLLM(gpt4mini, "Is this positive?", reason=True)
627        batch_yes_no = make_batch_llm(yes_no_llm)
628
629        # Process multiple prompts
630        results = batch_yes_no.process_batch([
631            "I love this!",
632            "This is terrible.",
633            "It's okay I guess."
634        ])
635
636        # Access processed results
637        for custom_id, result in results["results"].items():
638            print(f"{custom_id}: {result['processed']}")  # YesNoReason objects
639    """
640    return BatchLLMAdapter(llm)