From 7a2f3a7554a85626e54ef2cb31bcfdbdcf5654f6 Mon Sep 17 00:00:00 2001 From: abhinav Date: Sun, 2 Aug 2026 12:25:03 -0700 Subject: [PATCH] Add LiteLLM activity sample --- README.md | 1 + litellm_activity/README.md | 51 +++++++++++++++++++++++++ litellm_activity/__init__.py | 1 + litellm_activity/activities.py | 27 +++++++++++++ litellm_activity/shared.py | 10 +++++ litellm_activity/starter.py | 34 +++++++++++++++++ litellm_activity/worker.py | 31 +++++++++++++++ litellm_activity/workflow.py | 26 +++++++++++++ pyproject.toml | 2 + tests/litellm_activity/__init__.py | 0 tests/litellm_activity/activity_test.py | 44 +++++++++++++++++++++ tests/litellm_activity/workflow_test.py | 34 +++++++++++++++++ uv.lock | 4 ++ 13 files changed, 265 insertions(+) create mode 100644 litellm_activity/README.md create mode 100644 litellm_activity/__init__.py create mode 100644 litellm_activity/activities.py create mode 100644 litellm_activity/shared.py create mode 100644 litellm_activity/starter.py create mode 100644 litellm_activity/worker.py create mode 100644 litellm_activity/workflow.py create mode 100644 tests/litellm_activity/__init__.py create mode 100644 tests/litellm_activity/activity_test.py create mode 100644 tests/litellm_activity/workflow_test.py diff --git a/README.md b/README.md index 26264b8b1..0a823b400 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,7 @@ Some examples require extra dependencies. See each sample's directory for specif * [langfuse_tracing](langfuse_tracing) - Trace Temporal workflows in Langfuse with the OpenTelemetry plugin and OTLP export. * [langgraph_plugin](langgraph_plugin) - Run LangGraph workflows as durable Temporal workflows (Graph API and Functional API). * [langsmith_tracing](langsmith_tracing) - Trace Temporal workflows with LangSmith via the LangSmith plugin. +* [litellm_activity](litellm_activity) - Call LLM providers through LiteLLM from a Temporal Activity. * [message_passing/introduction](message_passing/introduction/) - Introduction to queries, signals, and updates. * [message_passing/safe_message_handlers](message_passing/safe_message_handlers/) - Safely handling updates and signals. * [message_passing/update_with_start/lazy_initialization](message_passing/update_with_start/lazy_initialization/) - Use update-with-start to update a Shopping Cart, starting it if it does not exist. diff --git a/litellm_activity/README.md b/litellm_activity/README.md new file mode 100644 index 000000000..2d3d11a14 --- /dev/null +++ b/litellm_activity/README.md @@ -0,0 +1,51 @@ +# LiteLLM Activity + +This sample calls an LLM provider through [LiteLLM](https://docs.litellm.ai/) from a Temporal Activity. + +LLM calls perform network I/O and return nondeterministic results, so they must not run in Workflow code. The Workflow only schedules the Activity and records its result, keeping replay deterministic. + +## Prerequisites + +Follow the [repository prerequisites](../README.md), then install the sample's dependencies: + +```bash +uv sync --group litellm +``` + +Set the API key expected by your provider. This example uses OpenAI by default: + +```bash +export OPENAI_API_KEY="your-api-key" +``` + +To use another [LiteLLM-supported provider](https://docs.litellm.ai/docs/providers), set its credentials and model name. For example: + +```bash +export ANTHROPIC_API_KEY="your-api-key" +export LITELLM_MODEL="anthropic/claude-sonnet-4-5-20250929" +``` + +Provider credentials stay in the Worker environment; they are not passed through the Workflow or stored in Event History. + +## Run the sample + +Start a local Temporal server, then run these commands in separate terminals: + +```bash +# Terminal 1: run the Worker +uv run --group litellm python -m litellm_activity.worker + +# Terminal 2: start a Workflow +uv run --group litellm python -m litellm_activity.starter \ + "Why should LLM calls run in Temporal Activities?" +``` + +The Activity gives each provider call a 30-second client timeout. The Workflow gives each Activity attempt 45 seconds, limits the entire Activity execution to two minutes, and retries failures up to three times with exponential backoff. LiteLLM's own retries are disabled so Temporal records and controls every attempt. + +## Tests + +The tests replace the provider call and Activity with deterministic fakes, so they do not require an API key or make live LLM requests: + +```bash +uv run --group litellm pytest tests/litellm_activity +``` diff --git a/litellm_activity/__init__.py b/litellm_activity/__init__.py new file mode 100644 index 000000000..ed0fc3745 --- /dev/null +++ b/litellm_activity/__init__.py @@ -0,0 +1 @@ +"""Call LiteLLM from a Temporal Activity.""" diff --git a/litellm_activity/activities.py b/litellm_activity/activities.py new file mode 100644 index 000000000..696bd0962 --- /dev/null +++ b/litellm_activity/activities.py @@ -0,0 +1,27 @@ +from litellm import ModelResponse, acompletion +from temporalio import activity + +from litellm_activity.shared import LLMRequest + + +@activity.defn +async def call_litellm(request: LLMRequest) -> str: + """Make the nondeterministic network call outside Workflow code.""" + response = await acompletion( + model=request.model, + messages=[ + {"role": "system", "content": request.system_prompt}, + {"role": "user", "content": request.prompt}, + ], + timeout=30, + # Let Temporal own retries so every attempt is visible in Event History. + num_retries=0, + ) + + if not isinstance(response, ModelResponse): + raise TypeError("Expected a non-streaming LiteLLM response") + + content = response.choices[0].message.content + if not content: + raise ValueError("LiteLLM returned an empty response") + return content diff --git a/litellm_activity/shared.py b/litellm_activity/shared.py new file mode 100644 index 000000000..008d0b1de --- /dev/null +++ b/litellm_activity/shared.py @@ -0,0 +1,10 @@ +from dataclasses import dataclass + + +@dataclass +class LLMRequest: + """Serializable input shared by the client, Workflow, and Activity.""" + + prompt: str + model: str = "openai/gpt-4o-mini" + system_prompt: str = "You are a helpful assistant." diff --git a/litellm_activity/starter.py b/litellm_activity/starter.py new file mode 100644 index 000000000..434c1c937 --- /dev/null +++ b/litellm_activity/starter.py @@ -0,0 +1,34 @@ +import asyncio +import os +import sys +import uuid + +from temporalio.client import Client +from temporalio.envconfig import ClientConfig + +from litellm_activity.shared import LLMRequest +from litellm_activity.worker import TASK_QUEUE +from litellm_activity.workflow import LiteLLMWorkflow + + +async def main() -> None: + config = ClientConfig.load_client_connect_config() + config.setdefault("target_host", "localhost:7233") + client = await Client.connect(**config) + + prompt = " ".join(sys.argv[1:]) or "Explain Temporal in one sentence." + request = LLMRequest( + prompt=prompt, + model=os.getenv("LITELLM_MODEL", "openai/gpt-4o-mini"), + ) + result = await client.execute_workflow( + LiteLLMWorkflow.run, + request, + id=f"litellm-activity-{uuid.uuid4()}", + task_queue=TASK_QUEUE, + ) + print(result) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/litellm_activity/worker.py b/litellm_activity/worker.py new file mode 100644 index 000000000..b8db9481b --- /dev/null +++ b/litellm_activity/worker.py @@ -0,0 +1,31 @@ +import asyncio +import logging + +from temporalio.client import Client +from temporalio.envconfig import ClientConfig +from temporalio.worker import Worker + +from litellm_activity.activities import call_litellm +from litellm_activity.workflow import LiteLLMWorkflow + +TASK_QUEUE = "litellm-activity-task-queue" + + +async def main() -> None: + logging.basicConfig(level=logging.INFO) + + config = ClientConfig.load_client_connect_config() + config.setdefault("target_host", "localhost:7233") + client = await Client.connect(**config) + + worker = Worker( + client, + task_queue=TASK_QUEUE, + workflows=[LiteLLMWorkflow], + activities=[call_litellm], + ) + await worker.run() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/litellm_activity/workflow.py b/litellm_activity/workflow.py new file mode 100644 index 000000000..c91e6696f --- /dev/null +++ b/litellm_activity/workflow.py @@ -0,0 +1,26 @@ +from datetime import timedelta + +from temporalio import workflow +from temporalio.common import RetryPolicy + +from litellm_activity.shared import LLMRequest + +with workflow.unsafe.imports_passed_through(): + from litellm_activity.activities import call_litellm + + +@workflow.defn +class LiteLLMWorkflow: + @workflow.run + async def run(self, request: LLMRequest) -> str: + return await workflow.execute_activity( + call_litellm, + request, + start_to_close_timeout=timedelta(seconds=45), + schedule_to_close_timeout=timedelta(minutes=2), + retry_policy=RetryPolicy( + initial_interval=timedelta(seconds=1), + maximum_interval=timedelta(seconds=10), + maximum_attempts=3, + ), + ) diff --git a/pyproject.toml b/pyproject.toml index e77dcd459..c3e945afa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,7 @@ langgraph = [ "langchain-anthropic>=0.3.0", "temporalio[langgraph,langsmith]>=1.30.0", ] +litellm = ["litellm>=1.85.0,<2"] nexus = ["nexus-rpc>=1.1.0,<2"] open-telemetry = [ "temporalio[opentelemetry]", @@ -119,6 +120,7 @@ packages = [ "langfuse_tracing", "langgraph_plugin", "langsmith_tracing", + "litellm_activity", "message_passing", "nexus", "open_telemetry", diff --git a/tests/litellm_activity/__init__.py b/tests/litellm_activity/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/litellm_activity/activity_test.py b/tests/litellm_activity/activity_test.py new file mode 100644 index 000000000..7cd39dc76 --- /dev/null +++ b/tests/litellm_activity/activity_test.py @@ -0,0 +1,44 @@ +from typing import Any + +from litellm import ModelResponse + +from litellm_activity import activities +from litellm_activity.shared import LLMRequest + + +async def test_call_litellm(monkeypatch: Any) -> None: + captured: dict[str, Any] = {} + + async def mock_acompletion(**kwargs: Any) -> ModelResponse: + captured.update(kwargs) + return ModelResponse( + model="test-model", + choices=[ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "Hello from LiteLLM"}, + } + ], + ) + + monkeypatch.setattr(activities, "acompletion", mock_acompletion) + + result = await activities.call_litellm( + LLMRequest( + prompt="Hello", + model="test/model", + system_prompt="Be concise.", + ) + ) + + assert result == "Hello from LiteLLM" + assert captured == { + "model": "test/model", + "messages": [ + {"role": "system", "content": "Be concise."}, + {"role": "user", "content": "Hello"}, + ], + "timeout": 30, + "num_retries": 0, + } diff --git a/tests/litellm_activity/workflow_test.py b/tests/litellm_activity/workflow_test.py new file mode 100644 index 000000000..739de2c7f --- /dev/null +++ b/tests/litellm_activity/workflow_test.py @@ -0,0 +1,34 @@ +import uuid + +from temporalio import activity +from temporalio.client import Client +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import Worker + +from litellm_activity.shared import LLMRequest +from litellm_activity.workflow import LiteLLMWorkflow + + +async def test_litellm_workflow(client: Client, env: WorkflowEnvironment) -> None: + expected = "Temporal makes LLM calls durable." + + @activity.defn(name="call_litellm") + async def mock_call_litellm(request: LLMRequest) -> str: + assert request.prompt == "What does Temporal add to LLM calls?" + return expected + + task_queue = f"test-litellm-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[LiteLLMWorkflow], + activities=[mock_call_litellm], + ): + result = await client.execute_workflow( + LiteLLMWorkflow.run, + LLMRequest(prompt="What does Temporal add to LLM calls?"), + id=f"test-litellm-{uuid.uuid4()}", + task_queue=task_queue, + ) + + assert result == expected diff --git a/uv.lock b/uv.lock index 2205dd8c9..ee53b7b6f 100644 --- a/uv.lock +++ b/uv.lock @@ -5140,6 +5140,9 @@ langsmith-tracing = [ { name = "openai" }, { name = "temporalio", extra = ["langsmith", "pydantic"] }, ] +litellm = [ + { name = "litellm" }, +] nexus = [ { name = "nexus-rpc" }, ] @@ -5235,6 +5238,7 @@ langsmith-tracing = [ { name = "openai", specifier = ">=1.4.0" }, { name = "temporalio", extras = ["pydantic", "langsmith"], specifier = ">=1.30.0" }, ] +litellm = [{ name = "litellm", specifier = ">=1.85.0,<2" }] nexus = [{ name = "nexus-rpc", specifier = ">=1.1.0,<2" }] open-telemetry = [ { name = "opentelemetry-exporter-otlp-proto-grpc" },