railtracks.llm

 25class ModelBase(ABC):
 26    """
 27    A simple base that represents the behavior of a model that can be used for chat, structured interactions, and streaming.
 28    """
 29
 30    def __init__(
 31        self,
 32        retry_approach: RetryApproach | None = None,
 33    ):
 34        self.retry_approach = retry_approach
 35        self.id = str(uuid4())
 36
 37    @abstractmethod
 38    def model_name(self) -> str:
 39        """
 40        Returns the name of the model being used.
 41
 42        It can be treated as unique identifier for the model when paired with the `model_type`.
 43        """
 44        pass
 45
 46    @abstractmethod
 47    def model_provider(self) -> ModelProvider:
 48        """The name of the provider of this model (The Company that owns the model)"""
 49        pass
 50
 51    @classmethod
 52    @abstractmethod
 53    def model_gateway(cls) -> ModelProvider:
 54        """
 55        Gets the API distrubutor of the model. Note nessecarily the same as the model itself.
 56
 57        E.g. if you are calling openai LLM through Azure AI foundry
 58        """
 59        pass
 60
 61    def chat(self, messages: MessageHistory) -> Response:
 62        """Chat with the model using the provided messages."""
 63        return self._chat(messages)
 64
 65    async def achat(self, messages: MessageHistory) -> Response:
 66        """Asynchronous chat with the model using the provided messages."""
 67        return await self._achat(messages)
 68
 69    def structured(self, messages: MessageHistory, schema: Type[BaseModel]) -> Response:
 70        """Structured interaction with the model using the provided messages and output_schema."""
 71        return self._structured(messages, schema)
 72
 73    async def astructured(
 74        self, messages: MessageHistory, schema: Type[BaseModel]
 75    ) -> Response:
 76        """Asynchronous structured interaction with the model using the provided messages and output_schema."""
 77        return await self._astructured(messages, schema)
 78
 79    def chat_with_tools(self, messages: MessageHistory, tools: List[Tool]) -> Response:
 80        """Chat with the model using the provided messages and tools."""
 81        return self._chat_with_tools(messages, tools)
 82
 83    async def achat_with_tools(
 84        self, messages: MessageHistory, tools: List[Tool]
 85    ) -> Response:
 86        """Asynchronous chat with the model using the provided messages and tools."""
 87        return await self._achat_with_tools(messages, tools)
 88
 89    # ================ START Streaming (per-call) LLM calls ===============
 90    # These methods request a streamed response for a single call. They are the model-level
 91    # building blocks of railtracks streaming (see `rt.astream` at the framework level).
 92
 93    async def astream_chat(
 94        self, messages: MessageHistory
 95    ) -> AsyncGenerator[str | Response, None]:
 96        """
 97        Chat with the model, streaming the response.
 98
 99        Returns an async generator that yields `str` token chunks as they arrive, followed by a
100        single final `Response` object containing the complete message (and usage info).
101
102        ```python
103        async for item in model.astream_chat(MessageHistory([UserMessage("hi")])):
104            if isinstance(item, str):
105                ...  # token chunk
106            else:
107                final = item  # the terminal Response
108        ```
109
110        Args:
111            messages: The conversation so far, as a `MessageHistory`.
112
113        Yields:
114            str | Response: `str` token chunks, then one final complete `Response`.
115        """
116        async for item in self._astream_chat(messages):
117            yield item
118
119    async def astream_chat_with_tools(
120        self, messages: MessageHistory, tools: List[Tool]
121    ) -> AsyncGenerator[str | Response, None]:
122        """
123        Chat with the model using tools, streaming the response.
124
125        Yields `str` content chunks as they arrive, followed by a single final `Response`. The
126        final `Response` contains either the complete assistant text or the requested tool
127        calls (tool-call deltas are accumulated internally and are not yielded as chunks).
128
129        Args:
130            messages: The conversation so far, as a `MessageHistory`.
131            tools: The tools to make available to the model.
132
133        Yields:
134            str | Response: `str` content chunks, then one final complete `Response`.
135        """
136        async for item in self._astream_chat_with_tools(messages, tools):
137            yield item
138
139    async def astream_structured(
140        self, messages: MessageHistory, schema: Type[BaseModel]
141    ) -> AsyncGenerator[str | Response, None]:
142        """
143        Structured interaction with the model, streaming the response.
144
145        Yields the raw (JSON) `str` chunks as they arrive, followed by a single final
146        `Response` whose message content is the parsed `schema` instance.
147
148        Note the chunks are unvalidated JSON fragments; validation only happens once the stream
149        completes, so a schema mismatch surfaces at the end of the stream.
150
151        Args:
152            messages: The conversation so far, as a `MessageHistory`.
153            schema: The pydantic model the response must conform to.
154
155        Yields:
156            str | Response: raw JSON `str` chunks, then one final complete `Response`.
157        """
158        async for item in self._astream_structured(messages, schema):
159            yield item
160
161    def supports_streamed_tool_calling(self) -> bool:
162        """Whether `astream_chat_with_tools` can stream on this model.
163
164        A model that routes tool calls through an API without incremental support should
165        return False; `rt.astream` then runs the call buffered (with a warning) rather than
166        failing, so the final result is unaffected and only the incremental chunks are lost.
167
168        Returns:
169            bool: True by default — subclasses narrow this where they can tell.
170        """
171        return True
172
173    # ================ END Streaming (per-call) LLM calls ===============
174
175    @abstractmethod
176    def _chat(self, messages: MessageHistory) -> Response:
177        pass
178
179    @abstractmethod
180    def _structured(
181        self, messages: MessageHistory, schema: Type[BaseModel]
182    ) -> Response:
183        pass
184
185    @abstractmethod
186    def _chat_with_tools(self, messages: MessageHistory, tools: List[Tool]) -> Response:
187        pass
188
189    # Note: the _astream_* methods are deliberately NOT abstract so that existing ModelBase
190    # subclasses keep working; subclasses that support streaming should override them with
191    # async generator implementations yielding `str` chunks followed by a final `Response`.
192
193    def _astream_chat(
194        self, messages: MessageHistory
195    ) -> AsyncGenerator[str | Response, None]:
196        raise NotImplementedError(
197            f"{type(self).__name__} does not support streamed chat calls."
198        )
199
200    def _astream_chat_with_tools(
201        self, messages: MessageHistory, tools: List[Tool]
202    ) -> AsyncGenerator[str | Response, None]:
203        raise NotImplementedError(
204            f"{type(self).__name__} does not support streamed tool-calling calls."
205        )
206
207    def _astream_structured(
208        self, messages: MessageHistory, schema: Type[BaseModel]
209    ) -> AsyncGenerator[str | Response, None]:
210        raise NotImplementedError(
211            f"{type(self).__name__} does not support streamed structured calls."
212        )
213
214    @abstractmethod
215    async def _achat(self, messages: MessageHistory) -> Response:
216        pass
217
218    @abstractmethod
219    async def _astructured(
220        self,
221        messages: MessageHistory,
222        schema: Type[BaseModel],
223    ) -> Response:
224        pass
225
226    @abstractmethod
227    async def _achat_with_tools(
228        self, messages: MessageHistory, tools: List[Tool]
229    ) -> Response:
230        pass

25class ProviderError(_ColoredError, Exception):
26    """
27    Base class for failures that happen while talking to a model provider.
28
29    Covers the model misbehaving, an unknown model, and exhausted retries. Errors about
30    *defining* a tool are a separate root -- see
31    :class:`railtracks.llm.tools.tool.ToolCreationError`.
32
33    Raised only by direct model calls. Inside a node these surface as
34    :class:`railtracks.exceptions.LLMError`.
35    """

38class ProviderTimeoutError(ProviderError):
39    """The provider did not answer in time."""

42class ProviderRateLimitError(ProviderError):
43    """The provider rejected the call for rate/quota reasons."""

46class ProviderAuthenticationError(ProviderError):
47    """The provider rejected the credentials. Retrying will not help."""

50class RetryError(ProviderError):
51    """
52    Raised when an error occurs during an LLM call that is being retried.
53    """
54
55    def __init__(
56        self,
57        retry_method: str,
58        message: str,
59        notes: list[str],
60        exception_list: list[Exception],
61    ):
62        full_message = (
63            f"LLM call failed after retries from {retry_method} retry: {message}"
64        )
65        self.message = message
66        self.notes = notes
67        self.exception_list = exception_list
68        super().__init__(full_message)

290class ToolCreationError(_ColoredError, Exception):
291    """Exception raised when a tool cannot be created.
292
293    A separate root from `ProviderError`: a malformed tool is a bug in the caller's
294    code, not a runtime failure of a provider, and it ends the run.
295    """
296
297    def __init__(self, message, notes=None):
298        super().__init__(message)
299        self.notes = notes or []
300
301    def __str__(self):
302        base = super().__str__()
303        if self.notes:
304            notes_str = (
305                "\n"
306                + self._color("Tips to debug:\n", self.GREEN)
307                + "\n".join(self._color(f"- {note}", self.GREEN) for note in self.notes)
308            )
309            return f"\n{self._color(base, self.RED)}{notes_str}"
310        return self._color(base, self.RED)

 6class ModelError(ProviderError):
 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        if not self.message_history:
25            return self._color(base, self.RED)
26
27        try:
28            count = len(self.message_history)
29        except TypeError:
30            count = None
31
32        summary = (
33            f"{count} message(s) redacted; "
34            "call err.format_verbose() to render or read err.message_history"
35            if count is not None
36            else "message history redacted; "
37            "call err.format_verbose() to render or read err.message_history"
38        )
39        detail = self._color("Message History: ", self.BOLD_GREEN) + self._color(
40            summary, self.GREEN
41        )
42        notes_str = "\n" + self._color("Details:\n", self.BOLD_GREEN) + f"  {detail}"
43        return f"\n{self._color(base, self.RED)}{notes_str}"
44
45    def format_verbose(self) -> str:
46        """Render the exception with the full input ``MessageHistory`` embedded."""
47        base = super().__str__()
48        if not self.message_history:
49            return self._color(base, self.RED)
50
51        mh_str = str(self.message_history)
52        indented_mh = "\n".join("    " + line for line in mh_str.splitlines())
53        detail = self._color("Message History:\n", self.BOLD_GREEN) + self._color(
54            indented_mh, self.GREEN
55        )
56        notes_str = "\n" + self._color("Details:\n", self.BOLD_GREEN) + f"  {detail}"
57        return f"\n{self._color(base, self.RED)}{notes_str}"

60class ModelNotFoundError(ProviderError):
61    def __init__(self, reason: str, notes: list[str] = None):
62        self.reason = reason
63        self.notes = notes or []
64        super().__init__(reason)
65
66    def __str__(self):
67        base = super().__str__()
68        if self.notes:
69            notes_str = (
70                "\n"
71                + self._color("Tips to debug:\n", self.GREEN)
72                + "\n".join(self._color(f"- {note}", self.GREEN) for note in self.notes)
73            )
74            return f"\n{self._color(base, self.RED)}{notes_str}"
75        return self._color(base, self.RED)

78class FunctionCallingNotSupportedError(ModelError):
79    """Error raised when a model does not support function calling."""
80
81    def __init__(self, model_name: str):
82        super().__init__(
83            reason=f"Model {model_name} does not support function calling. Chat with tools is not supported."
84        )

87class UnsupportedHyperparameterError(ModelError):
88    """Error raised when a model does not support a given common LLM hyperparameter."""
89
90    def __init__(self, model_name: str, hyperparameter: str, value):
91        super().__init__(
92            reason=(
93                f"Model {model_name} does not support '{hyperparameter}' "
94                f"(got {hyperparameter}={value!r})."
95            )
96        )

 99class MutuallyExclusiveHyperparametersError(ModelError):
100    """Error raised when two or more common hyperparameters cannot be combined for
101    this model."""
102
103    def __init__(self, model_name: str, hyperparameters: list[str], values: dict):
104        joined = " and ".join(f"'{p}'" for p in hyperparameters)
105        super().__init__(
106            reason=(
107                f"Model {model_name} does not support specifying {joined} together "
108                f"(got {values!r}). Use only one."
109            )
110        )

22class ToolCall(BaseModel):
23    """
24    A simple model object that represents a tool call.
25
26    This simple model represents a moment when a tool is called.
27    """
28
29    identifier: str = Field(description="The identifier attatched to this tool call.")
30    name: str = Field(description="The name of the tool being called.")
31    arguments: Dict[str, Any] = Field(
32        description="The arguments provided as input to the tool."
33    )
34
35    def __str__(self):
36        return f"{self.name}({self.arguments})"

39class ToolCalls(List[ToolCall]):
40    """
41    An assistant turn that calls tools: the calls themselves, plus any text the
42    model spoke alongside them.
43
44    Models routinely answer with both at once — "I'll look that up for you."
45    followed by a call to `search`. Holding that text here rather than in a
46    separate field on the message keeps the model's whole response inside
47    `content`.
48
49    This is a `list[ToolCall]`, so anything that already treats tool-call
50    content as a list — iteration, indexing, `isinstance(content, list)` —
51    keeps working, and `AssistantMessage` accepts a plain list as before.
52    """
53
54    def __init__(
55        self, tool_calls: Iterable[ToolCall] = (), text: str | None = None
56    ) -> None:
57        super().__init__(tool_calls)
58        self.text = text
59
60    def __repr__(self):
61        if self.text is None:
62            return f"ToolCalls({list.__repr__(self)})"
63        return f"ToolCalls({list.__repr__(self)}, text={self.text!r})"

66class ToolResponse(BaseModel):
67    """
68    A simple model object that represents a tool response.
69
70    This simple model should be used when adding a response to a tool.
71    """
72
73    identifier: str = Field(
74        description="The identifier attached to this tool response. This should match the identifier of the tool call."
75    )
76    name: str = Field(description="The name of the tool that generated this response.")
77    result: AnyStr = Field(description="The result of the tool call.")
78
79    def __str__(self):
80        return f"{self.name} -> {self.result}"

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

350class SystemMessage(_StringOnlyContent[Role.system]):
351    """
352    A simple class that represents a system message.
353
354    Args:
355        content (str): The content of the system message.
356    """
357
358    def __init__(self, content: str):
359        super().__init__(content=content, role=Role.system)

362class AssistantMessage(Message[_T, Role.assistant], Generic[_T]):
363    """
364    A simple class that represents a message from the assistant.
365
366    Args:
367        content (_T): The content of the assistant message. A tool-calling turn is a
368            `ToolCalls`, which holds both the calls and any text the model spoke
369            alongside them; a plain `list[ToolCall]` is accepted and normalized to one.
370    """
371
372    def __init__(self, content: _T):
373        # Normalizing here means a tool-calling turn is always a ToolCalls, so
374        # nothing downstream has to handle both shapes, while callers that pass a
375        # plain list of tool calls keep working.
376        if isinstance(content, list) and not isinstance(content, ToolCalls):
377            content = cast(_T, ToolCalls(content))
378
379        super().__init__(content=content, role=Role.assistant)
380
381        # Optionally stores the raw litellm message object so providers that
382        # attach extra metadata (e.g. Gemini thought_signature) can round-trip
383        # it back without any manual reconstruction.
384        self.raw_litellm_message: Any | None = None
385
386    def encode(self):
387        encoded = super().encode()
388
389        # A ToolCalls encodes as a bare array, which would drop the text the model
390        # spoke with its calls; surface it the way providers put it on the wire.
391        if isinstance(self.content, ToolCalls) and self.content.text is not None:
392            encoded["text"] = self.content.text
393
394        return encoded

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    ):
198        """
199        A simple class that represents a message that an LLM can read.
200
201        Args:
202            content: The content of the message. It can take on any of the following types:
203                - str: A simple string message.
204                - ToolCalls: Tool calls plus any text the model spoke alongside them.
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        """
211        assert isinstance(role, Role)
212        self.validate_content(content)
213        self._content = content
214        self._role = role
215
216    @classmethod
217    def validate_content(cls, content: _T):
218        pass
219
220    @property
221    def content(self) -> _T:
222        """Collects the content of the message."""
223        return self._content
224
225    @property
226    def role(self) -> _TRole:
227        """Collects the role of the message."""
228        return self._role
229
230    def __str__(self):
231        return f"{self.role.value}: {self.content}"
232
233    def __repr__(self):
234        return str(self)
235
236    def encode(self) -> dict[str, Any]:
237        return {
238            "role": self.role.value,
239            "content": self.content,
240        }
241
242    @property
243    def tool_calls(self):
244        """Gets the tool calls attached to this message, if any. If there are none return and empty list."""
245        tools: list[ToolCall] = []
246        if isinstance(self.content, list):
247            tools.extend(deepcopy(self.content))
248
249        return tools