A wrapper class to interact with a language model.
| 63 | |
| 64 | |
| 65 | class LLM: |
| 66 | """ |
| 67 | A wrapper class to interact with a language model. |
| 68 | """ |
| 69 | |
| 70 | def __init__( |
| 71 | self, |
| 72 | model: str = "gpt-4o-2024-08-06", |
| 73 | api_base: str = None, |
| 74 | use_openai: bool = True, |
| 75 | use_batch: bool = False, |
| 76 | ) -> None: |
| 77 | """ |
| 78 | Initialize the LLM. |
| 79 | |
| 80 | Args: |
| 81 | model (str): The model name. |
| 82 | api_base (str): The base URL for the API. |
| 83 | use_openai (bool): Whether to use OpenAI. |
| 84 | use_batch (bool): Whether to use OpenAI's Batch API, which is single thread only. |
| 85 | """ |
| 86 | if use_openai and "OPENAI_API_KEY" in os.environ: |
| 87 | self.client = OpenAI(base_url=api_base) |
| 88 | if use_batch and "OPENAI_API_KEY" in os.environ: |
| 89 | assert use_openai, "use_batch must be used with use_openai" |
| 90 | self.oai_batch = Auto(loglevel=0) |
| 91 | if "OPENAI_API_KEY" not in os.environ: |
| 92 | print("Warning: no API key found") |
| 93 | self.model = model |
| 94 | self.api_base = api_base |
| 95 | self._use_openai = use_openai |
| 96 | self._use_batch = use_batch |
| 97 | |
| 98 | @tenacity |
| 99 | def __call__( |
| 100 | self, |
| 101 | content: str, |
| 102 | images: list[str] = None, |
| 103 | system_message: str = None, |
| 104 | history: list = None, |
| 105 | delay_batch: bool = False, |
| 106 | return_json: bool = False, |
| 107 | return_message: bool = False, |
| 108 | ) -> str | dict | list: |
| 109 | """ |
| 110 | Call the language model with a prompt and optional images. |
| 111 | |
| 112 | Args: |
| 113 | content (str): The prompt content. |
| 114 | images (list[str]): A list of image file paths. |
| 115 | system_message (str): The system message. |
| 116 | history (list): The conversation history. |
| 117 | delay_batch (bool): Whether to delay return of response. |
| 118 | return_json (bool): Whether to return the response as JSON. |
| 119 | return_message (bool): Whether to return the message. |
| 120 | |
| 121 | Returns: |
| 122 | str | dict | list: The response from the model. |