llama-index

1
"""Program utils."""
2

3
from typing import Any, List, Type
4

5
from llama_index.legacy.bridge.pydantic import BaseModel, Field, create_model
6
from llama_index.legacy.llms.llm import LLM
7
from llama_index.legacy.output_parsers.pydantic import PydanticOutputParser
8
from llama_index.legacy.prompts.base import PromptTemplate
9
from llama_index.legacy.types import BasePydanticProgram, PydanticProgramMode
10

11

12
def create_list_model(base_cls: Type[BaseModel]) -> Type[BaseModel]:
13
    """Create a list version of an existing Pydantic object."""
14
    # NOTE: this is directly taken from
15
    # https://github.com/jxnl/openai_function_call/blob/main/examples/streaming_multitask/streaming_multitask.py
16
    # all credits go to the openai_function_call repo
17

18
    name = f"{base_cls.__name__}List"
19
    list_items = (
20
        List[base_cls],  # type: ignore
21
        Field(
22
            default_factory=list,
23
            repr=False,
24
            description=f"List of {base_cls.__name__} items",
25
        ),
26
    )
27

28
    new_cls = create_model(name, items=list_items)
29
    new_cls.__doc__ = f"A list of {base_cls.__name__} objects. "
30

31
    return new_cls
32

33

34
def get_program_for_llm(
35
    output_cls: BaseModel,
36
    prompt: PromptTemplate,
37
    llm: LLM,
38
    pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,
39
    **kwargs: Any,
40
) -> BasePydanticProgram:
41
    """Get a program based on the compatible LLM."""
42
    if pydantic_program_mode == PydanticProgramMode.DEFAULT:
43
        # in default mode, we try to use the OpenAI program if available else
44
        # we fall back to the LLM program
45
        try:
46
            from llama_index.legacy.program.openai_program import OpenAIPydanticProgram
47

48
            return OpenAIPydanticProgram.from_defaults(
49
                output_cls=output_cls,
50
                llm=llm,
51
                prompt=prompt,
52
                **kwargs,
53
            )
54
        except ValueError:
55
            from llama_index.legacy.program.llm_program import LLMTextCompletionProgram
56

57
            return LLMTextCompletionProgram.from_defaults(
58
                output_parser=PydanticOutputParser(output_cls=output_cls),
59
                llm=llm,
60
                prompt=prompt,
61
                **kwargs,
62
            )
63
    elif pydantic_program_mode == PydanticProgramMode.OPENAI:
64
        from llama_index.legacy.program.openai_program import OpenAIPydanticProgram
65

66
        return OpenAIPydanticProgram.from_defaults(
67
            output_cls=output_cls,
68
            llm=llm,
69
            prompt=prompt,
70
            **kwargs,
71
        )
72
    elif pydantic_program_mode == PydanticProgramMode.LLM:
73
        from llama_index.legacy.program.llm_program import LLMTextCompletionProgram
74

75
        return LLMTextCompletionProgram.from_defaults(
76
            output_parser=PydanticOutputParser(output_cls=output_cls),
77
            llm=llm,
78
            prompt=prompt,
79
            **kwargs,
80
        )
81
    elif pydantic_program_mode == PydanticProgramMode.LM_FORMAT_ENFORCER:
82
        from llama_index.legacy.program.lmformatenforcer_program import (
83
            LMFormatEnforcerPydanticProgram,
84
        )
85

86
        return LMFormatEnforcerPydanticProgram.from_defaults(
87
            output_cls=output_cls,
88
            llm=llm,
89
            prompt=prompt,
90
            **kwargs,
91
        )
92
    else:
93
        raise ValueError(f"Unsupported pydantic program mode: {pydantic_program_mode}")
94