railtracks.llm

 32class ModelBase(ABC, Generic[_TStream]):
 33    """
 34    A simple base that represents the behavior of a model that can be used for chat, structured interactions, and streaming.
 35
 36    The base class allows for the insertion of hooks that can modify the messages before they are sent to the model,
 37    response after they are received, and map exceptions that may occur during the interaction.
 38
 39    All the hooks are optional and can be added or removed as needed.
 40    """
 41
 42    def __init__(
 43        self,
 44        __pre_hooks: List[Callable[[MessageHistory], MessageHistory]] | None = None,
 45        __post_hooks: List[Callable[[MessageHistory, Response], Response]]
 46        | None = None,
 47        __exception_hooks: List[Callable[[MessageHistory, Exception], None]]
 48        | None = None,
 49        stream: _TStream = False,
 50        retry_approach: RetryApproach | None = None,
 51    ):
 52        if __pre_hooks is None:
 53            pre_hooks: List[Callable[[MessageHistory], MessageHistory]] = []
 54        else:
 55            pre_hooks = __pre_hooks
 56
 57        if __post_hooks is None:
 58            post_hooks: List[Callable[[MessageHistory, Response], Response]] = []
 59        else:
 60            post_hooks = __post_hooks
 61
 62        if __exception_hooks is None:
 63            exception_hooks: List[Callable[[MessageHistory, Exception], None]] = []
 64        else:
 65            exception_hooks = __exception_hooks
 66
 67        self._pre_hooks = pre_hooks
 68        self._post_hooks = post_hooks
 69        self._exception_hooks = exception_hooks
 70        self.stream = stream
 71        self.retry_approach = retry_approach
 72
 73    def add_pre_hook(self, hook: Callable[[MessageHistory], MessageHistory]) -> None:
 74        """Adds a pre-hook to modify messages before sending them to the model."""
 75        self._pre_hooks.append(hook)
 76
 77    def add_post_hook(
 78        self, hook: Callable[[MessageHistory, Response], Response]
 79    ) -> None:
 80        """Adds a post-hook to modify the response after receiving it from the model."""
 81        self._post_hooks.append(hook)
 82
 83    def add_exception_hook(
 84        self, hook: Callable[[MessageHistory, Exception], None]
 85    ) -> None:
 86        """Adds an exception hook to handle exceptions during model interactions."""
 87        self._exception_hooks.append(hook)
 88
 89    def remove_pre_hooks(self) -> None:
 90        """Removes all of the hooks that modify messages before sending them to the model."""
 91        self._pre_hooks = []
 92
 93    def remove_post_hooks(self) -> None:
 94        """Removes all of the hooks that modify the response after receiving it from the model."""
 95        self._post_hooks = []
 96
 97    def remove_exception_hooks(self) -> None:
 98        """Removes all of the hooks that handle exceptions during model interactions."""
 99        self._exception_hooks = []
100
101    @abstractmethod
102    def model_name(self) -> str:
103        """
104        Returns the name of the model being used.
105
106        It can be treated as unique identifier for the model when paired with the `model_type`.
107        """
108        pass
109
110    @abstractmethod
111    def model_provider(self) -> ModelProvider:
112        """The name of the provider of this model (The Company that owns the model)"""
113        pass
114
115    @classmethod
116    @abstractmethod
117    def model_gateway(cls) -> ModelProvider:
118        """
119        Gets the API distrubutor of the model. Note nessecarily the same as the model itself.
120
121        E.g. if you are calling openai LLM through Azure AI foundry
122        """
123        pass
124
125    def _run_pre_hooks(self, message_history: MessageHistory) -> MessageHistory:
126        """Runs all pre-hooks on the provided message history."""
127        for hook in self._pre_hooks:
128            message_history = hook(message_history)
129        return message_history
130
131    def _run_post_hooks(
132        self, message_history: MessageHistory, result: Response
133    ) -> Response:
134        """Runs all post-hooks on the provided message history and result."""
135        for hook in self._post_hooks:
136            result = hook(message_history, result)
137        return result
138
139    def _run_exception_hooks(
140        self, message_history: MessageHistory, exception: Exception
141    ) -> None:
142        """Runs all exception hooks on the provided message history and exception."""
143        for hook in self._exception_hooks:
144            hook(message_history, exception)
145
146    def generator_wrapper(
147        self,
148        generator: Generator[str | Response, None, Response],
149        message_history: MessageHistory,
150    ) -> Generator[str | Response, None, Response]:
151        new_response: Response | None = None
152        for g in generator:
153            if isinstance(g, Response):
154                g.message_info
155                new_response = self._run_post_hooks(message_history, g)
156                yield new_response
157
158            yield g
159
160        assert new_response is not None, (
161            "The generator did not yield a final Response object so nothing could be done."
162        )
163
164        return new_response
165
166    @overload
167    def chat(self: ModelBase[Literal[False]], messages: MessageHistory) -> Response:
168        pass
169
170    @overload
171    def chat(
172        self: ModelBase[Literal[True]], messages: MessageHistory
173    ) -> Generator[str | Response, None, Response]:
174        pass
175
176    def chat(
177        self, messages: MessageHistory
178    ) -> Response | Generator[str | Response, None, Response]:
179        """Chat with the model using the provided messages."""
180
181        messages = self._run_pre_hooks(messages)
182
183        try:
184            response = self._chat(messages)
185        except Exception as e:
186            self._run_exception_hooks(messages, e)
187            raise e
188
189        if isinstance(response, Generator):
190            return self.generator_wrapper(response, messages)
191
192        response = self._run_post_hooks(messages, response)
193        return response
194
195    @overload
196    async def achat(
197        self: ModelBase[Literal[False]], messages: MessageHistory
198    ) -> Response:
199        pass
200
201    @overload
202    async def achat(
203        self: ModelBase[Literal[True]], messages: MessageHistory
204    ) -> Generator[str | Response, None, Response]:
205        pass
206
207    async def achat(self, messages: MessageHistory):
208        """Asynchronous chat with the model using the provided messages."""
209        messages = self._run_pre_hooks(messages)
210
211        try:
212            response = await self._achat(messages)
213        except Exception as e:
214            self._run_exception_hooks(messages, e)
215            raise e
216
217        if isinstance(response, Generator):
218            return self.generator_wrapper(response, messages)
219
220        response = self._run_post_hooks(messages, response)
221
222        return response
223
224    @overload
225    def structured(
226        self: ModelBase[Literal[False]],
227        messages: MessageHistory,
228        schema: Type[BaseModel],
229    ) -> Response:
230        pass
231
232    @overload
233    def structured(
234        self: ModelBase[Literal[True]],
235        messages: MessageHistory,
236        schema: Type[BaseModel],
237    ) -> Generator[str | Response, None, Response]:
238        pass
239
240    def structured(self, messages: MessageHistory, schema: Type[BaseModel]):
241        """Structured interaction with the model using the provided messages and output_schema."""
242        messages = self._run_pre_hooks(messages)
243
244        try:
245            response = self._structured(messages, schema)
246        except Exception as e:
247            self._run_exception_hooks(messages, e)
248            raise e
249
250        if isinstance(response, Generator):
251            return self.generator_wrapper(response, messages)
252
253        response = self._run_post_hooks(messages, response)
254
255        return response
256
257    @overload
258    async def astructured(
259        self: ModelBase[Literal[False]],
260        messages: MessageHistory,
261        schema: Type[BaseModel],
262    ) -> Response:
263        pass
264
265    @overload
266    async def astructured(
267        self: ModelBase[Literal[True]],
268        messages: MessageHistory,
269        schema: Type[BaseModel],
270    ) -> Generator[str | Response, None, Response]:
271        pass
272
273    async def astructured(self, messages: MessageHistory, schema: Type[BaseModel]):
274        """Asynchronous structured interaction with the model using the provided messages and output_schema."""
275        messages = self._run_pre_hooks(messages)
276
277        try:
278            response = await self._astructured(messages, schema)
279        except Exception as e:
280            self._run_exception_hooks(messages, e)
281            raise e
282
283        if isinstance(response, Generator):
284            return self.generator_wrapper(response, messages)
285
286        response = self._run_post_hooks(messages, response)
287
288        return response
289
290    @overload
291    def chat_with_tools(
292        self: ModelBase[Literal[False]], messages: MessageHistory, tools: List[Tool]
293    ) -> Response:
294        pass
295
296    @overload
297    def chat_with_tools(
298        self: ModelBase[Literal[True]], messages: MessageHistory, tools: List[Tool]
299    ) -> Generator[str | Response, None, Response]:
300        pass
301
302    def chat_with_tools(self, messages: MessageHistory, tools: List[Tool]):
303        """Chat with the model using the provided messages and tools."""
304        messages = self._run_pre_hooks(messages)
305
306        try:
307            response = self._chat_with_tools(messages, tools)
308        except Exception as e:
309            self._run_exception_hooks(messages, e)
310            raise e
311
312        if isinstance(response, Generator):
313            return self.generator_wrapper(response, messages)
314
315        response = self._run_post_hooks(messages, response)
316        return response
317
318    @overload
319    async def achat_with_tools(
320        self: ModelBase[Literal[False]], messages: MessageHistory, tools: List[Tool]
321    ) -> Response:
322        pass
323
324    @overload
325    async def achat_with_tools(
326        self: ModelBase[Literal[True]], messages: MessageHistory, tools: List[Tool]
327    ) -> Generator[str | Response, None, Response]:
328        pass
329
330    async def achat_with_tools(self, messages: MessageHistory, tools: List[Tool]):
331        """Asynchronous chat with the model using the provided messages and tools."""
332        messages = self._run_pre_hooks(messages)
333
334        try:
335            response = await self._achat_with_tools(messages, tools)
336        except Exception as e:
337            self._run_exception_hooks(messages, e)
338            raise e
339
340        if isinstance(response, Generator):
341            return self.generator_wrapper(response, messages)
342
343        response = self._run_post_hooks(messages, response)
344
345        return response
346
347    @abstractmethod
348    def _chat(
349        self, messages: MessageHistory
350    ) -> Response | Generator[str | Response, None, Response]:
351        pass
352
353    @abstractmethod
354    def _structured(
355        self, messages: MessageHistory, schema: Type[BaseModel]
356    ) -> Response | Generator[str | Response, None, Response]:
357        pass
358
359    @abstractmethod
360    def _chat_with_tools(
361        self, messages: MessageHistory, tools: List[Tool]
362    ) -> Response | Generator[str | Response, None, Response]:
363        pass
364
365    @abstractmethod
366    async def _achat(
367        self, messages: MessageHistory
368    ) -> Response | AsyncGenerator[str | Response, None]:
369        pass
370
371    @abstractmethod
372    async def _astructured(
373        self,
374        messages: MessageHistory,
375        schema: Type[BaseModel],
376    ) -> Response | AsyncGenerator[str | Response, None]:
377        pass
378
379    @abstractmethod
380    async def _achat_with_tools(
381        self, messages: MessageHistory, tools: List[Tool]
382    ) -> Response | AsyncGenerator[str | Response, None]:
383        pass

 6class ModelError(RTLLMError):
 7    """
 8    Any Large Language Model (LLM) error.
 9    """
10
11    def __init__(
12        self,
13        reason: str,
14        message_history: MessageHistory = None,
15    ):
16        self.reason = reason
17        self.message_history = message_history
18
19        message = f"{self._color('Failure reason: ', self.BOLD_RED)}{self._color(reason, self.RED)}"
20        super().__init__(message)
21
22    def __str__(self):
23        base = super().__str__()
24        details = []
25        if self.message_history:
26            mh_str = str(self.message_history)
27            indented_mh = "\n".join(
28                "    " + line for line in mh_str.splitlines()
29            )  # 2 indents (2-spaces) per indent
30            details.append(
31                self._color("Message History:\n", self.BOLD_GREEN)
32                + self._color(indented_mh, self.GREEN)
33            )
34        if details:
35            notes_str = (
36                "\n"
37                + self._color("Details:\n", self.BOLD_GREEN)
38                + "\n".join(f"  {d}" for d in details)
39            )
40            return f"\n{self._color(base, self.RED)}{notes_str}"
41        return self._color(base, self.RED)

44class ModelNotFoundError(RTLLMError):
45    def __init__(self, reason: str, notes: list[str] = None):
46        self.reason = reason
47        self.notes = notes or []
48        super().__init__(reason)
49
50    def __str__(self):
51        base = super().__str__()
52        if self.notes:
53            notes_str = (
54                "\n"
55                + self._color("Tips to debug:\n", self.GREEN)
56                + "\n".join(self._color(f"- {note}", self.GREEN) for note in self.notes)
57            )
58            return f"\n{self._color(base, self.RED)}{notes_str}"
59        return self._color(base, self.RED)

62class FunctionCallingNotSupportedError(ModelError):
63    """Error raised when a model does not support function calling."""
64
65    def __init__(self, model_name: str):
66        super().__init__(
67            reason=f"Model {model_name} does not support function calling. Chat with tools is not supported."
68        )

71class UnsupportedHyperparameterError(ModelError):
72    """Error raised when a model does not support a given common LLM hyperparameter."""
73
74    def __init__(self, model_name: str, hyperparameter: str, value):
75        super().__init__(
76            reason=(
77                f"Model {model_name} does not support '{hyperparameter}' "
78                f"(got {hyperparameter}={value!r})."
79            )
80        )

83class MutuallyExclusiveHyperparametersError(ModelError):
84    """Error raised when two or more common hyperparameters cannot be combined for
85    this model."""
86
87    def __init__(self, model_name: str, hyperparameters: list[str], values: dict):
88        joined = " and ".join(f"'{p}'" for p in hyperparameters)
89        super().__init__(
90            reason=(
91                f"Model {model_name} does not support specifying {joined} together "
92                f"(got {values!r}). Use only one."
93            )
94        )

12class ToolCall(BaseModel):
13    """
14    A simple model object that represents a tool call.
15
16    This simple model represents a moment when a tool is called.
17    """
18
19    identifier: str = Field(description="The identifier attatched to this tool call.")
20    name: str = Field(description="The name of the tool being called.")
21    arguments: Dict[str, Any] = Field(
22        description="The arguments provided as input to the tool."
23    )
24
25    def __str__(self):
26        return f"{self.name}({self.arguments})"

29class ToolResponse(BaseModel):
30    """
31    A simple model object that represents a tool response.
32
33    This simple model should be used when adding a response to a tool.
34    """
35
36    identifier: str = Field(
37        description="The identifier attached to this tool response. This should match the identifier of the tool call."
38    )
39    name: str = Field(description="The name of the tool that generated this response.")
40    result: AnyStr = Field(description="The result of the tool call.")
41
42    def __str__(self):
43        return f"{self.name} -> {self.result}"

279class UserMessage(_StringOnlyContent[Role.user]):
280    """
281    Note that we only support string input
282
283    Args:
284        content: The content of the user message.
285        attachment: The file attachment(s) for the user message. Can be a single string or a list of strings,
286                    containing file paths, URLs, or data URIs. Defaults to None.
287        inject_prompt: Whether to inject prompt with context variables. Defaults to True.
288        trust_urls: Allow in-process fetch for URL attachments. Defaults to False.
289                    When False, `.pdf` URLs raise and unknown-extension URLs are
290                    handed to the provider unprobed (as `image_url`). When True,
291                    `.pdf` URLs are downloaded and embedded as base64, and
292                    unknown-extension URLs are HEAD-probed to detect PDFs.
293                    Set True only when every URL in `attachment` is
294                    developer-controlled — end-user-supplied URLs are an SSRF
295                    surface once their bytes are fetched in-process.
296        attachment_timeout: Per-request timeout in seconds for the HEAD probe
297                    and the in-process PDF download. Defaults to 10. Only
298                    applies when `trust_urls=True`. Raise this for large PDFs
299                    over slow links.
300    """
301
302    def __init__(
303        self,
304        content: str | None = None,
305        attachment: str | list[str] | None = None,
306        inject_prompt: bool = True,
307        trust_urls: bool = False,
308        attachment_timeout: float = 10.0,
309    ):
310        if attachment is not None:
311            if isinstance(attachment, list):
312                self.attachment = [
313                    Attachment(
314                        att,
315                        trust_urls=trust_urls,
316                        attachment_timeout=attachment_timeout,
317                    )
318                    for att in attachment
319                ]
320            else:
321                self.attachment = [
322                    Attachment(
323                        attachment,
324                        trust_urls=trust_urls,
325                        attachment_timeout=attachment_timeout,
326                    )
327                ]
328
329            if content is None:
330                logger.warning(
331                    "UserMessage initialized without content, setting to empty string."
332                )
333                content = ""
334        else:
335            self.attachment = None
336
337        if content is None:
338            raise ValueError(
339                "UserMessage must have content if no attachment is provided."
340            )
341        super().__init__(content=content, role=Role.user, inject_prompt=inject_prompt)

344class SystemMessage(_StringOnlyContent[Role.system]):
345    """
346    A simple class that represents a system message.
347
348    Args:
349        content (str): The content of the system message.
350        inject_prompt (bool, optional): Whether to inject prompt with context  variables. Defaults to True.
351    """
352
353    def __init__(self, content: str, inject_prompt: bool = True):
354        super().__init__(content=content, role=Role.system, inject_prompt=inject_prompt)

357class AssistantMessage(Message[_T, Role.assistant], Generic[_T]):
358    """
359    A simple class that represents a message from the assistant.
360
361    Args:
362        content (_T): The content of the assistant message.
363        inject_prompt (bool, optional): Whether to inject prompt with context  variables. Defaults to True.
364    """
365
366    def __init__(self, content: _T, inject_prompt: bool = True):
367        super().__init__(
368            content=content, role=Role.assistant, inject_prompt=inject_prompt
369        )
370
371        # Optionally stores the raw litellm message object so providers that
372        # attach extra metadata (e.g. Gemini thought_signature) can round-trip
373        # it back without any manual reconstruction.
374        self.raw_litellm_message = None

186class Message(Generic[_T, _TRole]):
187    """
188    A base class that represents a message that an LLM can read.
189
190    Note the content may take on a variety of allowable types.
191    """
192
193    def __init__(
194        self,
195        content: _T,
196        role: _TRole,
197        inject_prompt: bool = True,
198    ):
199        """
200        A simple class that represents a message that an LLM can read.
201
202        Args:
203            content: The content of the message. It can take on any of the following types:
204                - str: A simple string message.
205                - List[ToolCall]: A list of tool calls.
206                - ToolResponse: A tool response.
207                - BaseModel: A custom base model object.
208                - Stream: A stream object with a final_message and a generator.
209            role: The role of the message (assistant, user, system, tool, etc.).
210            inject_prompt (bool, optional): Whether to inject prompt with context variables. Defaults to True.
211        """
212        assert isinstance(role, Role)
213        self.validate_content(content)
214        self._content = content
215        self._role = role
216        self._inject_prompt = inject_prompt
217
218    @classmethod
219    def validate_content(cls, content: _T):
220        pass
221
222    @property
223    def content(self) -> _T:
224        """Collects the content of the message."""
225        return self._content
226
227    @property
228    def role(self) -> _TRole:
229        """Collects the role of the message."""
230        return self._role
231
232    @property
233    def inject_prompt(self) -> bool:
234        """
235        A boolean that indicates whether this message should be injected into from context.
236        """
237        return self._inject_prompt
238
239    @inject_prompt.setter
240    def inject_prompt(self, value: bool):
241        """
242        Sets the inject_prompt property.
243        """
244        self._inject_prompt = value
245
246    def __str__(self):
247        return f"{self.role.value}: {self.content}"
248
249    def __repr__(self):
250        return str(self)
251
252    @property
253    def tool_calls(self):
254        """Gets the tool calls attached to this message, if any. If there are none return and empty list."""
255        tools: list[ToolCall] = []
256        if isinstance(self.content, list):
257            tools.extend(deepcopy(self.content))
258
259        return tools