File size: 33,427 Bytes
5374a2d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
import io
import yaml
import inspect
import asyncio
import base64
from abc import ABC, abstractmethod
from pydantic import Field
from typing import Union, Optional, Type, Callable, List, Any, Dict

from ..core.parser import Parser
from .model_configs import LLMConfig
from ..core.module_utils import (
    parse_json_from_text, 
    get_type_name,
    parse_xml_from_text,
    parse_data_from_text
)


PARSER_VALID_MODE = ["str", "json", "xml", "title", "custom"]

class LLMOutputParser(Parser):
    """A basic parser for LLM-generated content.
    
    This parser stores the raw text generated by an LLM in the `.content` attribute
    and provides methods to extract structured data from this text using different
    parsing strategies.
    
    Attributes:
        content: The raw text generated by the LLM.
    """
    content: str = Field(default=None, exclude=True, description="the text generated by LLM")

    @classmethod
    def get_attrs(cls, return_type: bool = False) -> List[Union[str, tuple]]:
        """Returns the attributes of the LLMOutputParser class.
        
        Excludes ["class_name", "content"] by default.

        Args:
            return_type: Whether to return the type of the attributes along with their names.
        
        Returns:
            If `return_type` is True, returns a list of tuples where each tuple contains 
            the attribute name and its type. Otherwise, returns a list of attribute names.
        """
        attrs = [] 
        exclude_attrs = ["class_name", "content"]
        for field, field_info in cls.model_fields.items():
            if field not in exclude_attrs:
                if return_type:
                    field_type = get_type_name(field_info.annotation)
                    attrs.append((field, field_type))
                else:
                    attrs.append(field)
        return attrs
    
    @classmethod
    def get_attr_descriptions(cls) -> dict:
        """Returns the attributes and their descriptions.
        
        Returns:
            A dictionary mapping attribute names to their descriptions.
        """
        attrs = cls.get_attrs()
        results = {} 
        for field_name, field_info in cls.model_fields.items():
            if field_name not in attrs:
                continue
            field_desc = field_info.description if field_info.description is not None else "None"
            results[field_name] = field_desc
        return results

    @classmethod
    def get_content_data(cls, content: str, parse_mode: str = "json", parse_func: Optional[Callable] = None, **kwargs) -> dict:
        """Parses LLM-generated content into a dictionary.
        
        This method takes content from an LLM response and converts it to a structured
        dictionary based on the specified parsing mode.

        Args:
            content: The content to parse.
            parse_mode: The mode to parse the content. Must be one of:
                - 'str': Assigns the raw text content to all attributes of the parser. 
                - 'json': Extracts and parses JSON objects from LLM output. It will return a dictionary parsed from the first valid JSON string.
                - 'xml': Parses content using XML tags. It will return a dictionary parsed from the XML tags.
                - 'title': Parses content with Markdown-style headings.
                - 'custom': Uses custom parsing logic. Requires providing `parse_func` parameter as a custom parsing function.
            parse_func: The function to parse the content, only valid when parse_mode is 'custom'.
            **kwargs (Any): Additional arguments passed to the parsing function.
        
        Returns:
            The parsed content as a dictionary.
            
        Raises:
            ValueError: If parse_mode is invalid or if parse_func is not provided when parse_mode is 'custom'.
        """
        attrs = cls.get_attrs()
        if len(attrs) <= 0:
            return {} 
        
        if parse_mode == "str":
            parse_func = cls._parse_str_content
        elif parse_mode == "json":
            parse_func = cls._parse_json_content
        elif parse_mode == "xml":
            parse_func = cls._parse_xml_content
        elif parse_mode == "title":
            parse_func = cls._parse_title_content
        elif parse_mode == "custom":
            if parse_func is None:
                raise ValueError("`parse_func` must be provided when `parse_mode` is 'custom'.")
            # obtain the function inputs
            signature = inspect.signature(parse_func)
            if "content" not in signature.parameters:
                raise ValueError("`parse_func` must have an input argument `content`.")
            
            func_args = {}
            func_args["content"] = content
            for param_name, param in signature.parameters.items():
                if param_name == "content":
                    continue  # Already set
                if param_name in kwargs:
                    func_args[param_name] = kwargs[param_name]
            data = parse_func(**func_args)
            if not isinstance(data, dict):
                raise ValueError(f"The output of `parse_func` must be a dictionary, but found {type(data)}.")
            return data
        else:
            raise ValueError(f"Invalid value '{parse_mode}' detected for `parse_mode`. Available choices: {PARSER_VALID_MODE}")
        data = parse_func(content=content, **kwargs)
        return data
    
    @classmethod
    def _parse_str_content(cls, content: str, **kwargs) -> dict:
        """Parses content by setting all attributes to the raw content.
        
        Args:
            content: The content to parse.
            **kwargs: Additional arguments (not used).
        
        Returns:
            A dictionary mapping all attributes to the raw content.
        """
        # if `parse_mode=str` in llm.generate(), it will set all the defined parameters to `content`
        attrs = cls.get_attrs()
        return {attr: content for attr in attrs}
    
    @classmethod 
    def _parse_json_content(cls, content: str, **kwargs) -> dict:
        """Parses content by extracting and parsing a JSON object. 
        If the content contains multiple JSON objects, only the first one will be used. 
        
        Args:
            content: The content containing a JSON object.
            **kwargs: Additional arguments (not used).
        
        Returns:
            The parsed JSON as a dictionary.
            
        Raises:
            ValueError: If the content doesn't contain a valid JSON object.
        """
        extracted_json_list = parse_json_from_text(content)
        if len(extracted_json_list) > 0:
            json_str = extracted_json_list[0] # only use the first JSON
            try:
                data = yaml.safe_load(json_str)
                if not isinstance(data, dict):
                    if isinstance(data, list):
                        # LLM returns a list of JSON strings, without specifying the attribute name
                        attrs = cls.get_attrs()
                        if len(attrs) == 1:
                            # if there is only one attribute, use it as the attribute name 
                            return {attrs[0]: data}
                        else:
                            # if there are multiple attributes, raise an error
                            raise ValueError("The generated content is a list of JSON strings, but the attribute name for the list is not specified. You should instruct the LLM to specify the attribute name for the list.")
                    else:
                        raise ValueError(f"The generated content is not a valid JSON string:\n{json_str}")
            except Exception:
                raise ValueError(f"The generated content is not a valid JSON string:\n{json_str}")
        else:
            raise ValueError(f"The following generated content does not contain JSON string!\n{content}")
        
        return data
    
    @classmethod
    def _parse_xml_content(cls, content: str, **kwargs) -> dict:
        """Parses content by extracting values from XML tags.
        
        Each attribute of the parser is expected to be enclosed in XML tags
        with the attribute name as the tag name.
        
        Args:
            content: The content containing XML tags.
            **kwargs: Additional arguments (not used).
        
        Returns:
            A dictionary mapping attributes to their extracted values.
            
        Raises:
            ValueError: If the content is missing expected XML tags or if the
                        extracted values can't be converted to the expected types.
        """
        attrs_with_types: List[tuple] = cls.get_attrs(return_type=True)
        data = {} 
        for attr, attr_type in attrs_with_types:
            attr_raw_value_list = parse_xml_from_text(text=content, label=attr)
            if len(attr_raw_value_list) > 0:
                attr_raw_value = attr_raw_value_list[0]
                try:
                    attr_value = parse_data_from_text(text=attr_raw_value, datatype=attr_type)
                except Exception:
                    raise ValueError(f"Cannot parse text: {attr_raw_value} into {attr_type} data!")
            else:
                raise ValueError(f"The following generated content does not contain xml label <{attr}>xxx</{attr}>!\n{content}")
            data[attr] = attr_value
        return data
    
    @classmethod
    def _parse_title_content(cls, content: str, title_format: str = "## {title}", **kwargs) -> dict:
        """Parses content with markdown-style titles.
        
        Extracts sections from content that are divided by titles following
        the specified format described in `title_format`. The default format is "## {title}".
        For example:
        ```
        ## title1
        content1
        ## title2
        content2
        ```
        This content will be parsed into:
        ```
        {
            "title1": "content1",
            "title2": "content2"
        }
        ```
        Args:
            content: The content with title-divided sections.
            title_format: The format of the titles, default is "## {title}".
            **kwargs: Additional arguments (not used).

        Returns:
            A dictionary mapping title names to their section contents.
        """
        attrs: List[str] = cls.get_attrs()
        if not attrs:
            return {}
        
        output_titles = [title_format.format(title=attr) for attr in attrs]

        def is_output_title(text: str):
            for title in output_titles:
                if text.strip().lower().startswith(title.lower()):
                    return True, title
            return False, None

        data = {}
        current_output_name: str = None
        current_output_content: list = None
        for line in content.split("\n"):
            is_title, title = is_output_title(line)
            if is_title:
                if current_output_name is not None and current_output_content is not None:
                    data[current_output_name] = "\n".join(current_output_content)
                current_output_content = []
                current_output_name = title.replace("#", "").strip()
                output_titles.remove(title)
            else: 
                if current_output_content is not None:
                    current_output_content.append(line)
        if current_output_name is not None and current_output_content is not None:
            data[current_output_name] = "\n".join(current_output_content)
        return data
    
    @classmethod
    def parse(cls, content: str, parse_mode: str = "json", parse_func: Optional[Callable] = None, **kwargs) -> "LLMOutputParser":
        """Parses LLM-generated text into a structured parser instance.
        
        This is the main method for creating parser instances from LLM output.
        
        Args:
            content: The text generated by the LLM.
            parse_mode: The mode to parse the content, must be one of:
                - 'str': Assigns the raw text content to all attributes of the parser. 
                - 'json': Extracts and parses JSON objects from LLM output. Uses the first valid JSON string to create an instance of LLMOutputParser.
                - 'xml': Parses content using XML tags. Uses the XML tags to create an instance of LLMOutputParser.
                - 'title': Parses content with Markdown-style headings. Uses the Markdown-style headings to create an instance of LLMOutputParser. The default title format is "## {title}", you can change it by providing `title_format` parameter, which should be a string that contains `{title}` placeholder. 
                - 'custom': Uses custom parsing logic. Requires providing `parse_func` parameter as a custom parsing function. The `parse_func` must have a parameter named `content` and return a dictionary where the keys are the attribute names and the values are the parsed data. 
            parse_func: The function to parse the content, only valid when `parse_mode` is 'custom'.
            **kwargs (Any): Additional arguments passed to parsing functions, such as:
                - `title_format` for `parse_mode="title"`.
            
        Returns:
            An instance of LLMOutputParser containing the parsed data.
            
        Raises:
            ValueError: If parse_mode is invalid or if content is not a string.
        """
        if parse_mode not in PARSER_VALID_MODE:
            raise ValueError(f"'{parse_mode}' is an invalid value for `parse_mode`. Available choices: {PARSER_VALID_MODE}.")
        if not isinstance(content, str):
            raise ValueError(f"The input to {cls.__name__}.parse should be a str, but found {type(content)}.")
        data = cls.get_content_data(content=content, parse_mode=parse_mode, parse_func=parse_func, **kwargs)
        data.update({"content": content})
        parser = cls.from_dict(data, **kwargs)
        # parser.content = content
        return parser

    def __str__(self) -> str:
        """
        Returns a string representation of the parser.
        """
        return self.to_str()
    
    def to_str(self, **kwargs) -> str:
        """
        Converts the parser to a string.
        """
        return self.content
    
    def get_structured_data(self) -> dict:
        """Extracts structured data from the parser.
        
        Returns:
            A dictionary containing only the defined attributes and their values,
            excluding metadata like class_name.
        """
        attrs = type(self).get_attrs()
        data = self.to_dict(ignore=["class_name"])
        # structured_data = {attr: data[attr] for attr in attrs}
        structured_data = {key: value for key, value in data.items() if key in attrs}
        return structured_data


def _is_multimodal_content(content: Any) -> bool:
    """Check if content contains multimodal objects (TextChunk, ImageChunk, etc.)."""
    try:
        from ..rag.schema import TextChunk, ImageChunk
        
        # Handle different content formats
        if isinstance(content, list):
            return any(isinstance(item, (TextChunk, ImageChunk)) for item in content)
        elif isinstance(content, (TextChunk, ImageChunk)):
            return True
        
        return False
    except ImportError:
        return False


def _process_multimodal_content(content: List[Any], model_type: str = "openai") -> List[Dict[str, Any]]:
    """Convert multimodal content (TextChunk, ImageChunk) to model-specific message format."""
    try:
        from ..rag.schema import TextChunk, ImageChunk
    except ImportError:
        raise ImportError("Cannot import TextChunk/ImageChunk from rag.schema for multimodal processing")
    
    processed_content = []
    
    for item in content:
        if isinstance(item, TextChunk):
            processed_content.append({
                "type": "text",
                "text": item.text
            })
        elif isinstance(item, ImageChunk):
            if model_type.lower() in ["openai", "openrouter", "litellm"]:
                # OpenAI-style format
                image_data = _get_image_data_url(item)
                processed_content.append({
                    "type": "image_url",
                    "image_url": {"url": image_data}
                })
            else:
                # For other models, include image path/data as appropriate
                processed_content.append({
                    "type": "image",
                    "image_path": item.image_path,
                    "image_mimetype": item.image_mimetype
                })
        else:
            # Handle other types (strings, etc.)
            if isinstance(item, str):
                processed_content.append({
                    "type": "text", 
                    "text": item
                })
            else:
                # Convert to string as fallback
                processed_content.append({
                    "type": "text",
                    "text": str(item)
                })
    
    return processed_content


def _get_image_data_url(image_chunk) -> str:
    """Convert ImageChunk to data URL format for model consumption."""
    
    try:
        # Load image using the chunk's lazy loading
        image = image_chunk.get_image()
        if image is None:
            raise ValueError(f"Could not load image from path: {image_chunk.image_path}")
        
        # Convert to base64
        buffer = io.BytesIO()
        # Determine format from mimetype or default to PNG
        format_name = "PNG"
        if image_chunk.image_mimetype:
            format_name = image_chunk.image_mimetype.split('/')[-1].upper()
            if format_name not in ['PNG', 'JPEG', 'JPG', 'GIF', 'WEBP']:
                format_name = "PNG"
        
        image.save(buffer, format=format_name)
        image_data = base64.b64encode(buffer.getvalue()).decode('utf-8')
        
        # Create data URL
        mime_type = image_chunk.image_mimetype or f"image/{format_name.lower()}"
        return f"data:{mime_type};base64,{image_data}"
        
    except Exception as e:
        raise RuntimeError(f"Failed to convert image to data URL: {str(e)}")


class BaseLLM(ABC):
    """Abstract base class for Large Language Model implementations.
    
    This class defines the interface that all LLM implementations must follow,
    providing methods for generating text, formatting messages, and parsing output.
    
    Attributes:
        config: Configuration for the LLM.
        kwargs: Additional keyword arguments provided during initialization.
    """
    
    def __init__(self, config: LLMConfig, **kwargs):
        """Initializes the LLM with configuration.
        
        Args:
            config: Configuration object for the LLM.
            **kwargs (Any): Additional keyword arguments.
        """
        self.config = config
        self.kwargs = kwargs
        self.init_model()
    
    @abstractmethod
    def init_model(self):
        """Initializes the underlying model.
        
        This method should be implemented by subclasses to set up the actual LLM.
        """
        pass

    def __deepcopy__(self, memo) -> "BaseLLM":
        """Handles deep copying of the LLM instance.
        
        Returns the same instance when deepcopy is called, as LLM instances
        often cannot be meaningfully deep-copied.
        
        Args:
            memo (Dict[int, Any]): Memo dictionary used by the deepcopy process.
            
        Returns:
            The same LLM instance.
        """
        # return the same instance when deepcopy
        memo[id(self)] = self
        return self

    @abstractmethod
    def formulate_messages(self, prompts: List[str], system_messages: Optional[List[str]] = None) -> List[List[dict]]:
        """Converts input prompts into the chat format compatible with different LLMs.

        Args:
            prompts: A list of user prompts that need to be converted.
            system_messages: An optional list of system messages that provide instructions or context to the model.
        
        Returns:
            A list of message lists, where each inner list contains messages in the chat format required by LLMs. 
        """
        pass

    @abstractmethod
    def single_generate(self, messages: List[dict], **kwargs) -> str:
        """Generates LLM output for a single set of messages.

        Args:
            messages: The input messages to the LLM in chat format.
            **kwargs (Any): Additional keyword arguments for generation settings.
        
        Returns:
            The generated output text from the LLM.
        """
        pass

    @abstractmethod
    def batch_generate(self, batch_messages: List[List[dict]], **kwargs) -> List[str]:
        """Generates outputs for a batch of message sets.

        Args: 
            batch_messages: A list of message lists, where each inner list contains messages for a single generation.
            **kwargs (Any): Additional keyword arguments for generation settings.
        
        Returns:
            A list of generated outputs from the LLM, one for each input message set.
        """
        pass
    
    async def single_generate_async(self, messages: List[dict], **kwargs) -> str:
        """Asynchronously generates LLM output for a single set of messages.
        
        This default implementation wraps the synchronous method in an async executor.
        Subclasses should override this for true async implementation if supported.
        
        Args:
            messages: The input messages to the LLM in chat format.
            **kwargs (Any): Additional keyword arguments for generation settings.
        
        Returns:
            The generated output text from the LLM.
        """
        # Default implementation for backward compatibility
        loop = asyncio.get_event_loop()
        result = await loop.run_in_executor(None, self.single_generate, messages, **kwargs)
        return result
    
    async def batch_generate_async(self, batch_messages: List[List[dict]], **kwargs) -> List[str]:
        """Asynchronously generates outputs for a batch of message sets.
        
        This default implementation runs each generation as a separate async task.
        Subclasses should override this for more efficient async batching if supported.
        
        Args: 
            batch_messages: A list of message lists, where each inner list contains messages for a single generation.
            **kwargs (Any): Additional keyword arguments for generation settings.
        
        Returns:
            A list of generated outputs from the LLM, one for each input message set.
        """
        # Default implementation for backward compatibility
        tasks = [self.single_generate_async(messages, **kwargs) for messages in batch_messages]
        return await asyncio.gather(*tasks)
    
    def parse_generated_text(self, text: str, parser: Optional[Type[LLMOutputParser]]=None, parse_mode: Optional[str] = "json", parse_func: Optional[Callable] = None, **kwargs) -> LLMOutputParser:
        """Parses generated text into a structured output using a parser.

        Args: 
            text: The text generated by the LLM.
            parser: An LLMOutputParser class to use for parsing. If None, the default LLMOutputParser is used.
            parse_mode: The mode to use for parsing, must be the `parse_mode` supported by the `parser`. 
            **kwargs (Any): Additional arguments passed to the parser.
        
        Returns:
            An LLMOutputParser instance containing the parsed data.
        """
        if not parser:
            parser = LLMOutputParser
        return parser.parse(text, parse_mode=parse_mode, parse_func=parse_func)

    def parse_generated_texts(self, texts: List[str], parser: Optional[Type[LLMOutputParser]]=None, parse_mode: Optional[str] = "json", parse_func: Optional[Callable] = None, **kwargs) -> List[LLMOutputParser]:
        """Parses multiple generated texts into structured outputs.
        
        Args:
            texts: A list of texts generated by the LLM.
            parser: An LLMOutputParser class to use for parsing.
            parse_mode: The mode to use for parsing, must be the `parse_mode` supported by the `parser`. 
            **kwargs (Any): Additional arguments passed to the parser.
            
        Returns:
            A list of LLMOutputParser instances containing the parsed data.
        """
        parsed_results = [self.parse_generated_text(text=text, parser=parser, parse_mode=parse_mode, parse_func=parse_func, **kwargs) for text in texts]
        return parsed_results

    def _prepare_messages(
        self,
        prompt: Optional[Union[str, List[str]]] = None,
        system_message: Optional[Union[str, List[str]]] = None,
        messages: Optional[Union[List[dict],List[List[dict]]]] = None,
    ) -> tuple[List[List[dict]], bool]:
        """Prepares and validates input messages for generation.
        
        This internal method handles the various input formats (prompt strings, system messages,
        or pre-formatted message dictionaries) and converts them to a consistent format for generation.
        
        Args:
            prompt: Input prompt(s) to the LLM.
            system_message: System message(s) for the LLM.
            messages: Chat message(s) for the LLM, already in the required format.

        Returns:
            A tuple containing:
            - prepared_messages: List of message lists ready for generation
            - is_single_generate: Boolean indicating if this is a single generation (vs. batch)
            
        Raises:
            ValueError: If neither prompt nor messages is provided, or if both are provided.
            TypeError: If the inputs have inconsistent types or formats.
        """
        if not (prompt or messages):
            raise ValueError("Either 'prompt' or 'messages' must be provided.")
        if prompt and messages:
            raise ValueError("Both 'prompt' and 'messages' are provided. Please provide only one of them.")

        single_generate = False
        if messages is not None:
            if not messages:  # empty messages
                return [], False
            if isinstance(messages[0], dict):
                single_generate = True
                messages = [messages]
            
            # Process multimodal content in messages
            processed_messages = self._process_messages_for_multimodal(messages)
            return processed_messages, single_generate

        if isinstance(prompt, str):
            single_generate = True
            prompt = [prompt]
            if system_message:
                if not isinstance(system_message, str):
                    raise TypeError(f"'system_message' should be a string when passing a single prompt, but found {type(system_message)}.")
                system_message = [system_message]
        elif isinstance(prompt, list) and all(isinstance(p, str) for p in prompt):
            single_generate = False
            if not prompt:  # empty prompt
                return [], False
            if system_message:
                if not isinstance(system_message, list) or len(prompt) != len(system_message):
                    raise ValueError(f"'system_message' should be a list of string when passing multiple prompts and the number of prompts ({len(prompt)}) must match the number of system messages ({len(system_message)}).")
        else:
            raise ValueError(f"'prompt' must be a str or List[str], but found {type(prompt)}.")
        
        prepared_messages = self.formulate_messages(prompts=prompt, system_messages=system_message)
        return prepared_messages, single_generate

    def _process_messages_for_multimodal(self, messages: List[List[dict]]) -> List[List[dict]]:
        """Process messages to handle multimodal content (TextChunk, ImageChunk)."""
        processed_messages = []
        
        for message_list in messages:
            processed_message_list = []
            
            for message in message_list:
                processed_message = message.copy()
                content = message.get("content")
                
                # Check if content contains multimodal objects
                if _is_multimodal_content(content):
                    # Get model type for proper formatting
                    llm_type = getattr(self.config, 'llm_type', 'openai')
                    # Map LLM types to processing types
                    if llm_type.lower() in ["openaillm", "openai"]:
                        model_type = "openai"
                    elif llm_type.lower() in ["litellm"]:
                        model_type = "litellm" 
                    elif llm_type.lower() in ["openrouter"]:
                        model_type = "openrouter"
                    else:
                        model_type = "openai"  # Default to OpenAI format
                    
                    from ..core.logging import logger
                    logger.debug(f"Processing multimodal content: llm_type={llm_type}, model_type={model_type}")
                    
                    # Convert multimodal content to appropriate format
                    if isinstance(content, list):
                        processed_message["content"] = _process_multimodal_content(content, model_type)
                    else:
                        # Single chunk - wrap in list for processing
                        processed_message["content"] = _process_multimodal_content([content], model_type)
                
                processed_message_list.append(processed_message)
            
            processed_messages.append(processed_message_list)
        
        return processed_messages

    def generate(
        self,
        prompt: Optional[Union[str, List[str]]] = None,
        system_message: Optional[Union[str, List[str]]] = None,
        messages: Optional[Union[List[dict],List[List[dict]]]] = None,
        parser: Optional[Type[LLMOutputParser]] = None,
        parse_mode: Optional[str] = "json", 
        parse_func: Optional[Callable] = None,
        **kwargs
    ) -> Union[LLMOutputParser, List[LLMOutputParser]]:
        """Generates LLM output(s) and parses the result(s).
        
        This is the main method for generating text with the LLM. It handles both
        single and batch generation, and automatically parses the outputs.
        
        Args:
            prompt: Input prompt(s) to the LLM.
            system_message: System message(s) for the LLM.
            messages: Chat message(s) for the LLM, already in the required format (either `prompt` or `messages` must be provided).
            parser: Parser class to use for processing the output.
            parse_mode: The mode to use for parsing, must be the `parse_mode` supported by the `parser`. 
            **kwargs (Any): Additional generation configuration parameters.
        
        Returns:
            For single generation: An LLMOutputParser instance.
            For batch generation: A list of LLMOutputParser instances.
            
        Raises:
            ValueError: If the input format is invalid.
            
        Note:
            Either prompt or messages must be provided. If both or neither is provided,
            an error will be raised.
        """
        prepared_messages, single_generate = self._prepare_messages(prompt, system_message, messages)
        if not prepared_messages:  # Handle empty messages case
            return []

        generated_texts = self.batch_generate(batch_messages=prepared_messages, **kwargs)
        parsed_outputs = self.parse_generated_texts(texts=generated_texts, parser=parser, parse_mode=parse_mode, parse_func=parse_func, **kwargs)
        return parsed_outputs[0] if single_generate else parsed_outputs
    
    async def async_generate(
        self,
        prompt: Optional[Union[str, List[str]]] = None,
        system_message: Optional[Union[str, List[str]]] = None,
        messages: Optional[Union[List[dict],List[List[dict]]]] = None,
        parser: Optional[Type[LLMOutputParser]] = None,
        parse_mode: Optional[str] = "json", 
        parse_func: Optional[Callable] = None,
        **kwargs
    ) -> Union[LLMOutputParser, List[LLMOutputParser]]:
        """Asynchronously generates LLM output(s) and parses the result(s).
        
        This is the async version of the generate method. It works identically but
        performs the generation asynchronously.
        """
        prepared_messages, single_generate = self._prepare_messages(prompt, system_message, messages)
        if not prepared_messages:  # Handle empty messages case
            return []

        generated_texts = await self.batch_generate_async(batch_messages=prepared_messages, **kwargs)
        parsed_outputs = self.parse_generated_texts(texts=generated_texts, parser=parser, parse_mode=parse_mode, parse_func=parse_func, **kwargs)
        return parsed_outputs[0] if single_generate else parsed_outputs


__all__ = ["LLMConfig", "BaseLLM", "LLMOutputParser"]