diff --git a/docs/documentation/getting_started/llm_setup.md b/docs/documentation/getting_started/llm_setup.md index 5df97a39a..244048be7 100644 --- a/docs/documentation/getting_started/llm_setup.md +++ b/docs/documentation/getting_started/llm_setup.md @@ -85,3 +85,26 @@ Not sure which to use? A few pointers: | Need multimodal (vision) | OpenAI, Anthropic, or Gemini | For a full breakdown of supported providers and how to configure each one, see [Supported Providers](../../integrations/llms/providers.md). + +--- + +## Testing without an API key + +`EchoLLM` is an offline model for tests and local examples. Constructing it and calling it needs no API key and no network access. Pass it to an [agent node](../agent_design/overview.md#agent-node) the same way you would pass any other model, then run that agent through a [Flow](../invocation/flows.md). + +On every chat call the model returns one assistant message: the full message history it received, one message per line in each message's `str()` form (`role: content`), then a final line that reads `This is an echo of the test model.` It never requests a tool call, so an agent that has tools finishes in a single turn and the reply is still that echo. The async methods return the same text as the sync ones. [Structured output](../agent_design/structured_extraction.md) is not supported; `structured` and `astructured` raise `NotImplementedError` with the message `EchoLLM does not support structured output.` + +```python +--8<-- "docs/scripts/documentation/echo_llm.py:echo" +``` + +Running that example prints: + +```text +system: You are a helper +user: hi +This is an echo of the test model. +``` + +The response reports zero input tokens, zero output tokens, zero cost, and zero latency, and the model name is `echo`. + diff --git a/docs/scripts/documentation/echo_llm.py b/docs/scripts/documentation/echo_llm.py new file mode 100644 index 000000000..5124a0980 --- /dev/null +++ b/docs/scripts/documentation/echo_llm.py @@ -0,0 +1,12 @@ +# --8<-- [start: echo] +import railtracks as rt + +agent = rt.agent_node( + llm=rt.llm.EchoLLM(), + system_message="You are a helper", +) + +flow = rt.Flow("Echo example", entry_point=agent) +result = flow.invoke("hi") +print(result.content) +# --8<-- [end: echo] diff --git a/packages/railtracks/src/railtracks/llm/__init__.py b/packages/railtracks/src/railtracks/llm/__init__.py index f55160c84..28493be19 100644 --- a/packages/railtracks/src/railtracks/llm/__init__.py +++ b/packages/railtracks/src/railtracks/llm/__init__.py @@ -14,6 +14,7 @@ AnthropicLLM, AppleFMLLM, AzureAILLM, + EchoLLM, GeminiLLM, HuggingFaceLLM, OllamaLLM, @@ -69,6 +70,7 @@ "AnthropicLLM", "AppleFMLLM", "AzureAILLM", + "EchoLLM", "HuggingFaceLLM", "OpenAILLM", "GeminiLLM", diff --git a/packages/railtracks/src/railtracks/llm/models/__init__.py b/packages/railtracks/src/railtracks/llm/models/__init__.py index 0c6b6ebe3..cb7a52e46 100644 --- a/packages/railtracks/src/railtracks/llm/models/__init__.py +++ b/packages/railtracks/src/railtracks/llm/models/__init__.py @@ -8,6 +8,7 @@ from .cloud import AzureAILLM, PortKeyLLM from .local.apple_fm import AppleFMLLM from .local.ollama import OllamaLLM +from .testing import EchoLLM __all__ = [ OpenAILLM, @@ -18,5 +19,6 @@ AppleFMLLM, HuggingFaceLLM, PortKeyLLM, + EchoLLM, "OpenAICompatibleProvider", ] diff --git a/packages/railtracks/src/railtracks/llm/models/testing/__init__.py b/packages/railtracks/src/railtracks/llm/models/testing/__init__.py new file mode 100644 index 000000000..a62850ab2 --- /dev/null +++ b/packages/railtracks/src/railtracks/llm/models/testing/__init__.py @@ -0,0 +1,5 @@ +from .echo import EchoLLM + +__all__ = [ + "EchoLLM", +] diff --git a/packages/railtracks/src/railtracks/llm/models/testing/echo.py b/packages/railtracks/src/railtracks/llm/models/testing/echo.py new file mode 100644 index 000000000..22b1d0de8 --- /dev/null +++ b/packages/railtracks/src/railtracks/llm/models/testing/echo.py @@ -0,0 +1,219 @@ +"""Offline echo model for tests that must not call a provider.""" + +from __future__ import annotations + +import asyncio +from typing import AsyncGenerator, List, Type + +from pydantic import BaseModel + +from ...history import MessageHistory +from ...message import AssistantMessage +from ...model import ModelBase +from ...providers import ModelProvider +from ...response import MessageInfo, Response +from ...tools import Tool + +_ECHO_LINE = "This is an echo of the test model." +_STREAM_CHUNKS = 5 +_STREAM_SECONDS = 5.0 + + +def _split_evenly(text: str, parts: int) -> List[str]: + """Split `text` into `parts` slices whose lengths differ by at most one. + + Earlier slices receive any leftover characters. A text shorter than + `parts` still returns `parts` slices; the later ones are empty. + + Args: + text: The text to split. + parts: How many slices to return. + + Returns: + List[str]: `parts` slices that join back to `text`. + """ + base, extra = divmod(len(text), parts) + chunks: List[str] = [] + start = 0 + for index in range(parts): + size = base + (1 if index < extra else 0) + chunks.append(text[start : start + size]) + start += size + return chunks + + +class EchoLLM(ModelBase): + """Offline model that returns the message history it was given. + + Each chat call returns one assistant message. The text is every received + message, in order, one per line in its ``str()`` form (``role: content``), + followed by ``This is an echo of the test model.`` The model never requests + a tool call. A streamed chat call yields that same text as five slices, + one per second, and then the full response. ``structured`` and + ``astructured`` raise ``NotImplementedError``. + + Creating or calling this model does not need an API key or network access. + The reported usage is zero tokens, zero cost, and zero latency, and the + model name is ``echo``. + """ + + def model_name(self) -> str: + """Return ``"echo"``. + + Returns: + str: The name ``"echo"``. + """ + return "echo" + + def model_provider(self) -> ModelProvider: + """Return `ModelProvider.UNKNOWN`. + + The echo model is not served by a provider. + + Returns: + ModelProvider: `ModelProvider.UNKNOWN`. + """ + return ModelProvider.UNKNOWN + + @classmethod + def model_gateway(cls) -> ModelProvider: + """Return `ModelProvider.UNKNOWN`. + + The echo model is not reached through a gateway. + + Returns: + ModelProvider: `ModelProvider.UNKNOWN`. + """ + return ModelProvider.UNKNOWN + + def _message_info(self) -> MessageInfo: + return MessageInfo( + input_tokens=0, + output_tokens=0, + latency=0.0, + model_name=self.model_name(), + total_cost=0.0, + ) + + def _echo(self, messages: MessageHistory) -> Response: + lines = [str(message) for message in messages] + lines.append(_ECHO_LINE) + return Response( + message=AssistantMessage("\n".join(lines)), + message_info=self._message_info(), + ) + + def _chat(self, messages: MessageHistory) -> Response: + """Return the echoed history as an assistant message. + + Args: + messages: The conversation the caller sent. + + Returns: + Response: One assistant message containing the compacted history. + """ + return self._echo(messages) + + def _chat_with_tools(self, messages: MessageHistory, tools: List[Tool]) -> Response: + """Return the echoed history and do not request a tool call. + + Args: + messages: The conversation the caller sent. + tools: Tools offered to the model. None of them is requested. + + Returns: + Response: The same echo `_chat` returns for `messages`. + """ + return self._echo(messages) + + def _structured( + self, messages: MessageHistory, schema: Type[BaseModel] + ) -> Response: + """Reject structured output. + + Args: + messages: The conversation the caller sent. + schema: The pydantic model requested for the response. + + Raises: + NotImplementedError: EchoLLM does not support structured output. + """ + raise NotImplementedError("EchoLLM does not support structured output.") + + async def _achat(self, messages: MessageHistory) -> Response: + """Asynchronous form of `_chat`. + + Args: + messages: The conversation the caller sent. + + Returns: + Response: The same response `_chat` returns. + """ + return self._chat(messages) + + async def _achat_with_tools( + self, messages: MessageHistory, tools: List[Tool] + ) -> Response: + """Asynchronous form of `_chat_with_tools`. + + Args: + messages: The conversation the caller sent. + tools: Tools offered to the model. None of them is requested. + + Returns: + Response: The same response `_chat_with_tools` returns. + """ + return self._chat_with_tools(messages, tools) + + async def _astructured( + self, messages: MessageHistory, schema: Type[BaseModel] + ) -> Response: + """Asynchronous form of `_structured`. + + Args: + messages: The conversation the caller sent. + schema: The pydantic model requested for the response. + + Raises: + NotImplementedError: EchoLLM does not support structured output. + """ + return self._structured(messages, schema) + + async def _astream_chat( + self, messages: MessageHistory + ) -> AsyncGenerator[str | Response, None]: + """Stream the echoed history over five seconds. + + The text `_chat` returns is split into five slices and one is yielded + each second. Earlier slices take any leftover characters, so the + lengths differ by at most one. A text shorter than five characters + still yields five slices; the later ones are empty. Joining the slices + reproduces the response text. The last item is that same response, and + the reported usage stays at zero. + + Args: + messages: The conversation the caller sent. + + Yields: + str | Response: Five text slices, then the `Response` `_chat` returns. + """ + response = self._echo(messages) + for chunk in _split_evenly(response.text, _STREAM_CHUNKS): + await asyncio.sleep(_STREAM_SECONDS / _STREAM_CHUNKS) + yield chunk + yield response + + async def _astream_chat_with_tools( + self, messages: MessageHistory, tools: List[Tool] + ) -> AsyncGenerator[str | Response, None]: + """Stream the echoed history and do not request a tool call. + + Args: + messages: The conversation the caller sent. + tools: Tools offered to the model. None of them is requested. + + Yields: + str | Response: The same slices and final `Response` `_astream_chat` yields. + """ + async for item in self._astream_chat(messages): + yield item diff --git a/packages/railtracks/tests/unit_tests/llm/test_echo.py b/packages/railtracks/tests/unit_tests/llm/test_echo.py new file mode 100644 index 000000000..46e78a2a6 --- /dev/null +++ b/packages/railtracks/tests/unit_tests/llm/test_echo.py @@ -0,0 +1,384 @@ +import socket +from collections.abc import AsyncIterable + +import pytest +import railtracks as rt +from pydantic import BaseModel +from railtracks.exceptions import LLMError +from railtracks.llm import ( + AssistantMessage, + EchoLLM, + MessageHistory, + ModelProvider, + Response, + SystemMessage, + ToolMessage, + ToolResponse, + UserMessage, +) + +_ECHO_LINE = "This is an echo of the test model." +_SPEC_HISTORY = "system: You are a helper\nuser: hi\nThis is an echo of the test model." + + +def _history() -> MessageHistory: + return MessageHistory( + [ + SystemMessage("You are a helper"), + UserMessage("hi"), + AssistantMessage("ok"), + ToolMessage(ToolResponse(identifier="1", name="add", result="3")), + ] + ) + + +def _expected_echo(messages: MessageHistory) -> str: + lines = [str(message) for message in messages] + lines.append(_ECHO_LINE) + return "\n".join(lines) + + +def test_echo_llm_is_exported_from_railtracks_llm(): + model = rt.llm.EchoLLM() + assert isinstance(model, EchoLLM) + assert "EchoLLM" in rt.llm.__all__ + + +def test_identity_and_zero_usage(): + model = EchoLLM() + response = model.chat( + MessageHistory([SystemMessage("You are a helper"), UserMessage("hi")]) + ) + + assert model.model_name() == "echo" + assert model.model_provider() is ModelProvider.UNKNOWN + assert model.model_gateway() is ModelProvider.UNKNOWN + assert response.message.role == "assistant" + assert response.message.content == _SPEC_HISTORY + info = response.message_info + assert info.input_tokens == 0 + assert info.output_tokens == 0 + assert info.total_tokens == 0 + assert info.total_cost == 0 + assert info.latency == 0 + assert info.model_name == "echo" + + +def test_chat_echoes_every_role(): + messages = _history() + response = EchoLLM().chat(messages) + assert response.message.content == ( + "system: You are a helper\n" + "user: hi\n" + "assistant: ok\n" + "tool: add -> 3\n" + "This is an echo of the test model." + ) + assert response.message.content == _expected_echo(messages) + + +def test_empty_history_is_only_the_echo_line(): + response = EchoLLM().chat(MessageHistory()) + assert response.message.content == _ECHO_LINE + + +def test_chat_with_tools_never_requests_a_tool(): + messages = _history() + tool = rt.llm.Tool("add", "Add two numbers.", None) + response = EchoLLM().chat_with_tools(messages, [tool]) + + assert response.message.content == _expected_echo(messages) + assert response.message.tool_calls == [] + + +async def test_async_methods_match_sync(): + model = EchoLLM() + messages = _history() + tool = rt.llm.Tool("add", "Add two numbers.", None) + + chat = await model.achat(messages) + with_tools = await model.achat_with_tools(messages, [tool]) + + assert chat.message.content == model.chat(messages).message.content + assert ( + with_tools.message.content + == model.chat_with_tools(messages, [tool]).message.content + ) + assert chat.message_info.model_name == "echo" + assert chat.message_info.input_tokens == 0 + assert chat.message_info.output_tokens == 0 + assert chat.message_info.total_cost == 0 + + +def test_structured_output_raises(): + model = EchoLLM() + messages = MessageHistory([UserMessage("hi")]) + + class Answer(BaseModel): + text: str + + with pytest.raises( + NotImplementedError, match="EchoLLM does not support structured output" + ): + model.structured(messages, Answer) + + +async def test_astructured_output_raises(): + model = EchoLLM() + messages = MessageHistory([UserMessage("hi")]) + + class Answer(BaseModel): + text: str + + with pytest.raises( + NotImplementedError, match="EchoLLM does not support structured output" + ): + await model.astructured(messages, Answer) + + +def test_call_needs_no_api_key_or_network(monkeypatch): + for name in ( + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "GEMINI_API_KEY", + "AZURE_API_KEY", + "HUGGINGFACE_API_KEY", + ): + monkeypatch.delenv(name, raising=False) + + def fail_connect(*args, **kwargs): + raise AssertionError(f"EchoLLM opened a connection: {args!r} {kwargs!r}") + + monkeypatch.setattr(socket, "create_connection", fail_connect) + monkeypatch.setattr(socket.socket, "connect", fail_connect) + + model = EchoLLM() + response = model.chat(MessageHistory([UserMessage("hi")])) + assert response.message.content == f"user: hi\n{_ECHO_LINE}" + + +@rt.function_node +def add(a: int, b: int) -> int: + """Add two numbers. + + Args: + a: The first number. + b: The second number. + + Returns: + The sum. + """ + return a + b + + +def test_flow_invoke_without_tools(): + agent = rt.agent_node( + llm=rt.llm.EchoLLM(), + system_message="You are a helper", + ) + result = rt.Flow("echo without tools", entry_point=agent).invoke("hi") + assert result.content == _SPEC_HISTORY + assert result.tool_invocations == [] + + +async def test_flow_ainvoke_without_tools(): + agent = rt.agent_node( + llm=rt.llm.EchoLLM(), + system_message="You are a helper", + ) + result = await rt.Flow("echo without tools async", entry_point=agent).ainvoke("hi") + assert result.content == _SPEC_HISTORY + + +def test_flow_invoke_with_tools_finishes_in_one_turn(): + agent = rt.agent_node( + llm=rt.llm.EchoLLM(), + system_message="You are a helper", + tool_nodes=[add], + ) + result = rt.Flow("echo with tools", entry_point=agent).invoke("hi") + assert result.content == _SPEC_HISTORY + assert result.tool_invocations == [] + + +async def test_flow_ainvoke_with_tools_finishes_in_one_turn(): + agent = rt.agent_node( + llm=rt.llm.EchoLLM(), + system_message="You are a helper", + tool_nodes=[add], + ) + result = await rt.Flow("echo with tools async", entry_point=agent).ainvoke("hi") + assert result.content == _SPEC_HISTORY + assert result.tool_invocations == [] + + +def test_flow_structured_output_raises(): + class Answer(BaseModel): + text: str + + agent = rt.agent_node( + llm=rt.llm.EchoLLM(), + system_message="You are a helper", + output_schema=Answer, + ) + with pytest.raises(LLMError, match="EchoLLM does not support structured output"): + rt.Flow("echo structured", entry_point=agent).invoke("hi") + + +def _record_sleeps(monkeypatch: pytest.MonkeyPatch) -> list[float]: + sleeps: list[float] = [] + + async def fake_sleep(seconds: float) -> None: + sleeps.append(seconds) + + monkeypatch.setattr("railtracks.llm.models.testing.echo.asyncio.sleep", fake_sleep) + return sleeps + + +async def _collect( + stream: AsyncIterable[str | Response], +) -> tuple[list[str], Response]: + chunks: list[str] = [] + final: Response | None = None + async for item in stream: + if isinstance(item, str): + assert final is None + chunks.append(item) + else: + assert isinstance(item, Response) + assert final is None + final = item + assert final is not None + return chunks, final + + +def _assert_even_chunks(chunks: list[str], content: str) -> None: + assert len(chunks) == 5 + assert "".join(chunks) == content + sizes = [len(chunk) for chunk in chunks] + assert max(sizes) - min(sizes) <= 1 + assert sizes == sorted(sizes, reverse=True) + + +async def test_astream_chat_yields_five_chunks_then_the_chat_response(monkeypatch): + sleeps = _record_sleeps(monkeypatch) + model = EchoLLM() + messages = _history() + chat = model.chat(messages) + + chunks, final = await _collect(model.astream_chat(messages)) + + _assert_even_chunks(chunks, chat.message.content) + assert sleeps == [1.0, 1.0, 1.0, 1.0, 1.0] + assert sum(sleeps) == 5 + assert final.message.content == chat.message.content + assert final.message.tool_calls == [] + assert final.message_info.input_tokens == 0 + assert final.message_info.output_tokens == 0 + assert final.message_info.total_tokens == 0 + assert final.message_info.total_cost == 0 + assert final.message_info.latency == 0 + assert final.message_info.model_name == "echo" + + +async def test_astream_chat_with_tools_matches_chat_and_requests_nothing(monkeypatch): + sleeps = _record_sleeps(monkeypatch) + model = EchoLLM() + messages = _history() + tool = rt.llm.Tool("add", "Add two numbers.", None) + chat = model.chat_with_tools(messages, [tool]) + + chunks, final = await _collect(model.astream_chat_with_tools(messages, [tool])) + + _assert_even_chunks(chunks, chat.message.content) + assert sum(sleeps) == 5 + assert final.message.content == chat.message.content + assert final.message.tool_calls == [] + assert final.message_info.input_tokens == 0 + assert final.message_info.output_tokens == 0 + assert final.message_info.latency == 0 + assert final.message_info.model_name == "echo" + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + ("", ["", "", "", "", ""]), + ("a", ["a", "", "", "", ""]), + ("abcd", ["a", "b", "c", "d", ""]), + ("abcde", ["a", "b", "c", "d", "e"]), + ("abcdef", ["ab", "c", "d", "e", "f"]), + ("ab\ncd", ["a", "b", "\n", "c", "d"]), + ("123456789012", ["123", "456", "78", "90", "12"]), + ("🙂🙂", ["🙂", "🙂", "", "", ""]), + ], +) +async def test_astream_chat_splits_short_and_uneven_text( + monkeypatch: pytest.MonkeyPatch, text: str, expected: list[str] +) -> None: + sleeps = _record_sleeps(monkeypatch) + model = EchoLLM() + sent = Response( + message=AssistantMessage(text), + message_info=model._message_info(), + ) + + def fake_echo(self, messages: MessageHistory) -> Response: + return sent + + monkeypatch.setattr(EchoLLM, "_echo", fake_echo) + + chunks, final = await _collect( + model.astream_chat(MessageHistory([UserMessage("hi")])) + ) + + assert chunks == expected + assert "".join(chunks) == sent.message.content + assert final is sent + assert sleeps == [1.0] * 5 + assert sum(sleeps) == 5 + + +async def test_astream_structured_raises(): + model = EchoLLM() + messages = MessageHistory([UserMessage("hi")]) + + class Answer(BaseModel): + text: str + + with pytest.raises( + NotImplementedError, + match="EchoLLM does not support streamed structured calls", + ): + async for _item in model.astream_structured(messages, Answer): + pass + + +async def test_rt_astream_without_tools(monkeypatch): + _record_sleeps(monkeypatch) + agent = rt.agent_node( + llm=rt.llm.EchoLLM(), + system_message="You are a helper", + ) + stream = rt.astream(agent, user_input="hi") + chunks = [chunk async for chunk in stream] + + _assert_even_chunks(chunks, _SPEC_HISTORY) + assert stream.result.content == _SPEC_HISTORY + assert stream.result.text == _SPEC_HISTORY + assert stream.result.tool_invocations == [] + + +async def test_rt_astream_with_tools(monkeypatch): + _record_sleeps(monkeypatch) + agent = rt.agent_node( + llm=rt.llm.EchoLLM(), + system_message="You are a helper", + tool_nodes=[add], + ) + stream = rt.astream(agent, user_input="hi") + chunks = [chunk async for chunk in stream] + + _assert_even_chunks(chunks, _SPEC_HISTORY) + assert "".join(chunks) == stream.result.content + assert stream.result.tool_invocations == []