railtracks.llm

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

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}"

178class UserMessage(_StringOnlyContent[Role.user]):
179    """
180    Note that we only support string input
181
182    Args:
183        content: The content of the user message.
184        attachment: The file attachment(s) for the user message. Can be a single string or a list of strings,
185                    containing file paths, URLs, or data URIs. Defaults to None.
186        inject_prompt: Whether to inject prompt with context variables. Defaults to True.
187    """
188
189    def __init__(
190        self,
191        content: str | None = None,
192        attachment: str | list[str] | None = None,
193        inject_prompt: bool = True,
194    ):
195        if attachment is not None:
196            if isinstance(attachment, list):
197                self.attachment = [Attachment(att) for att in attachment]
198            else:
199                self.attachment = [Attachment(attachment)]
200
201            if content is None:
202                logger.warning(
203                    "UserMessage initialized without content, setting to empty string."
204                )
205                content = ""
206        else:
207            self.attachment = None
208
209        if content is None:
210            raise ValueError(
211                "UserMessage must have content if no attachment is provided."
212            )
213        super().__init__(content=content, role=Role.user, inject_prompt=inject_prompt)

216class SystemMessage(_StringOnlyContent[Role.system]):
217    """
218    A simple class that represents a system message.
219
220    Args:
221        content (str): The content of the system message.
222        inject_prompt (bool, optional): Whether to inject prompt with context  variables. Defaults to True.
223    """
224
225    def __init__(self, content: str, inject_prompt: bool = True):
226        super().__init__(content=content, role=Role.system, inject_prompt=inject_prompt)

229class AssistantMessage(Message[_T, Role.assistant], Generic[_T]):
230    """
231    A simple class that represents a message from the assistant.
232
233    Args:
234        content (_T): The content of the assistant message.
235        inject_prompt (bool, optional): Whether to inject prompt with context  variables. Defaults to True.
236    """
237
238    def __init__(self, content: _T, inject_prompt: bool = True):
239        super().__init__(
240            content=content, role=Role.assistant, inject_prompt=inject_prompt
241        )
242
243        # Optionally stores the raw litellm message object so providers that
244        # attach extra metadata (e.g. Gemini thought_signature) can round-trip
245        # it back without any manual reconstruction.
246        self.raw_litellm_message = None

 85class Message(Generic[_T, _TRole]):
 86    """
 87    A base class that represents a message that an LLM can read.
 88
 89    Note the content may take on a variety of allowable types.
 90    """
 91
 92    def __init__(
 93        self,
 94        content: _T,
 95        role: _TRole,
 96        inject_prompt: bool = True,
 97    ):
 98        """
 99        A simple class that represents a message that an LLM can read.
100
101        Args:
102            content: The content of the message. It can take on any of the following types:
103                - str: A simple string message.
104                - List[ToolCall]: A list of tool calls.
105                - ToolResponse: A tool response.
106                - BaseModel: A custom base model object.
107                - Stream: A stream object with a final_message and a generator.
108            role: The role of the message (assistant, user, system, tool, etc.).
109            inject_prompt (bool, optional): Whether to inject prompt with context variables. Defaults to True.
110        """
111        assert isinstance(role, Role)
112        self.validate_content(content)
113        self._content = content
114        self._role = role
115        self._inject_prompt = inject_prompt
116
117    @classmethod
118    def validate_content(cls, content: _T):
119        pass
120
121    @property
122    def content(self) -> _T:
123        """Collects the content of the message."""
124        return self._content
125
126    @property
127    def role(self) -> _TRole:
128        """Collects the role of the message."""
129        return self._role
130
131    @property
132    def inject_prompt(self) -> bool:
133        """
134        A boolean that indicates whether this message should be injected into from context.
135        """
136        return self._inject_prompt
137
138    @inject_prompt.setter
139    def inject_prompt(self, value: bool):
140        """
141        Sets the inject_prompt property.
142        """
143        self._inject_prompt = value
144
145    def __str__(self):
146        return f"{self.role.value}: {self.content}"
147
148    def __repr__(self):
149        return str(self)
150
151    @property
152    def tool_calls(self):
153        """Gets the tool calls attached to this message, if any. If there are none return and empty list."""
154        tools: list[ToolCall] = []
155        if isinstance(self.content, list):
156            tools.extend(deepcopy(self.content))
157
158        return tools

250class ToolMessage(Message[ToolResponse, Role.tool]):
251    """
252    A simple class that represents a message that is a tool call answer.
253
254    Args:
255        content (ToolResponse): The tool response content for the message.
256    """
257
258    def __init__(self, content: ToolResponse):
259        if not isinstance(content, ToolResponse):
260            raise TypeError(
261                f"A {self.__class__.__name__} needs a ToolResponse but got {type(content)}. Check the invoke function of the OutputLessToolCallLLM node. That is the only place to return a ToolMessage."
262            )
263        super().__init__(content=content, role=Role.tool)

 9class MessageHistory(List[Message]):
10    """
11    A basic object that represents a history of messages. The object has all the same capability as a list such as
12    `.remove()`, `.append()`, etc.
13    """
14
15    def __str__(self):
16        return "\n".join([str(message) for message in self])
17
18    def removed_system_messages(self) -> MessageHistory:
19        """
20        Returns a new MessageHistory object with all SystemMessages removed.
21        """
22        return MessageHistory([msg for msg in self if msg.role != Role.system])

13class ModelProvider(str, Enum):
14    """
15    Enum of supported LLM model providers for RailTracks.
16
17    Attributes:
18        OPENAI: OpenAI models (e.g., GPT-3, GPT-4).
19        ANTHROPIC: Anthropic models (e.g., Claude).
20        GEMINI: Google Gemini models.
21        HUGGINGFACE: HuggingFace-hosted models.
22        AZUREAI: Azure OpenAI Service models.
23        OLLAMA: Ollama local LLMs.
24        COHERE: Cohere models.
25    """
26
27    OPENAI = "OpenAI"
28    ANTHROPIC = "Anthropic"
29    GEMINI = "Vertex_AI"
30    HUGGINGFACE = "HuggingFace"
31    AZUREAI = "AzureAI"
32    OLLAMA = "Ollama"
33    COHERE = "cohere_chat"
34    TELUS = "Telus"
35    PORTKEY = "PortKey"
36    UNKNOWN = "Unknown"

 29class Tool:
 30    """
 31    A quasi-immutable class designed to represent a single Tool object.
 32    You pass in key details (name, description, and required parameters).
 33    """
 34
 35    def __init__(
 36        self,
 37        name: str,
 38        detail: str,
 39        parameters: Iterable[Parameter] | Dict[str, Any] | None = None,
 40    ):
 41        """
 42        Creates a new Tool instance.
 43
 44        Args:
 45            name: The name of the tool.
 46            detail: A detailed description of the tool.
 47            parameters: Parameters attached to this tool; a set of Parameter objects, or a dict.
 48        """
 49
 50        if (
 51            isinstance(parameters, dict) and len(parameters) > 0
 52        ):  # if parameters is a JSON-output_schema, convert into Parameter objects (Checks should be done in validate_tool_params)
 53            props = parameters.get("properties")
 54            required_fields = list(parameters.get("required", []))
 55            param_objs: List[Parameter] = []
 56            for name, prop in props.items():
 57                param_objs.append(
 58                    parse_json_schema_to_parameter(name, prop, name in required_fields)
 59                )
 60            parameters = param_objs
 61
 62        self._name = name
 63        self._detail = detail
 64        self._parameters = parameters
 65
 66    @property
 67    def name(self) -> str:
 68        """Get the name of the tool."""
 69        return self._name
 70
 71    @property
 72    def detail(self) -> str:
 73        """Returns the detailed description for this tool."""
 74        return self._detail
 75
 76    @property
 77    def parameters(self) -> List[Parameter] | None:
 78        """Gets the parameters attached to this tool (if any)."""
 79        return self._parameters
 80
 81    def __str__(self) -> str:
 82        """String representation of the tool."""
 83        if self._parameters:
 84            params_str = "{" + ", ".join(str(p) for p in self._parameters) + "}"
 85        return f"Tool(name={self._name}, detail={self._detail}, parameters={params_str if self._parameters else 'None'})"
 86
 87    @classmethod
 88    def from_function(
 89        cls,
 90        func: Callable,
 91        /,
 92        *,
 93        name: str | None = None,
 94        details: str | None = None,
 95        params: Type[BaseModel] | Dict[str, Any] | List[Parameter] | None = None,
 96    ) -> Self:
 97        """
 98        Creates a Tool from a Python callable.
 99        Uses the function's docstring and type annotations to extract details and parameter info.
100
101        KEY NOTE: No checking is done to ensure that the inserted params match the function signature
102
103        Args:
104            func: The function to create a tool from.
105            name: Optional name for the tool. If not provided, uses the function's name.
106            details: Optional detailed description for the tool. If not provided, extracts from the function's docstring.
107            params: Optional parameters for the tool. If not provided, infers from the function's signature and docstring.
108
109        Returns:
110            A Tool instance representing the function.
111        """
112        # TODO: add set verification to ensure that the params match the function signature
113        # Check if the function is a method in a class
114        in_class = bool(func.__qualname__ and "." in func.__qualname__)
115
116        # Parse the docstring to get parameter descriptions
117        arg_descriptions = parse_docstring_args(func.__doc__ or "")
118
119        try:
120            # Get the function signature
121            signature = inspect.signature(func)
122        except ValueError:
123            raise ToolCreationError(
124                message="Cannot convert kwargs for builtin functions.",
125                notes=[
126                    "Please use a cutom made function.",
127                    "Eg.- \ndef my_custom_function(a: int, b: str):\n    pass",
128                ],
129            )
130
131        if name is not None:
132            # TODO: add some checking here to ensure that the name is valid snake case.
133            function_name = name
134        else:
135            function_name = func.__name__
136
137        docstring = func.__doc__.strip() if func.__doc__ else ""
138
139        if params is not None:
140            parameters = params
141        else:
142            # Check for multiple Args sections (warning)
143            # Only need to do this if we need to.
144            if docstring.count("Args:") > 1:
145                warnings.warn("Multiple 'Args:' sections found in the docstring.")
146            # Create parameter handlers
147            handlers: List[ParameterHandler] = [
148                PydanticModelHandler(),
149                SequenceParameterHandler(),
150                UnionParameterHandler(),
151                DefaultParameterHandler(),
152            ]
153
154            parameters: List[Parameter] = []
155
156            for param in signature.parameters.values():
157                # Skip 'self' parameter for class methods
158                if in_class and (param.name == "self" or param.name == "cls"):
159                    continue
160
161                description = arg_descriptions.get(param.name, "")
162
163                # Check if the parameter is required
164                required = param.default == inspect.Parameter.empty
165
166                handler = next(h for h in handlers if h.can_handle(param.annotation))
167
168                param_obj = handler.create_parameter(
169                    param.name, param.annotation, description, required
170                )
171
172                parameters.append(param_obj)
173
174        if details is not None:
175            main_description = details
176        else:
177            main_description = extract_main_description(docstring)
178
179        tool_info = Tool(
180            name=function_name,
181            detail=main_description,
182            parameters=parameters,
183        )
184        return tool_info
185
186    @classmethod
187    def from_mcp(cls, tool) -> Self:
188        """
189        Creates a Tool from an MCP tool object.
190
191        Args:
192            tool: The MCP tool to create a Tool from.
193
194        Returns:
195            A Tool instance representing the MCP tool.
196        """
197        input_schema = getattr(tool, "inputSchema", None)
198        if not input_schema or input_schema["type"] != "object":
199            raise ToolCreationError(
200                message="The inputSchema for an MCP Tool must be 'object'. ",
201                notes=[
202                    "If an MCP tool has a different output_schema, create a GitHub issue and support will be added."
203                ],
204            )
205
206        properties = input_schema.get("properties", {})
207        required_fields = set(input_schema.get("required", []))
208        param_objs = set()
209        for name, prop in properties.items():
210            required = name in required_fields
211            param_objs.add(parse_json_schema_to_parameter(name, prop, required))
212
213        return cls(name=tool.name, detail=tool.description, parameters=param_objs)