diff --git a/pyproject.toml b/pyproject.toml index b9682d2..b523f47 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -121,6 +121,9 @@ distributed = [ "pyspark<4.1", "ray==2.48", ] +litellm = [ + "litellm>=1.80.0,<1.87.0", +] [project.scripts] timecopilot = "timecopilot._cli:main" diff --git a/timecopilot/agent.py b/timecopilot/agent.py index f24e053..0e14203 100644 --- a/timecopilot/agent.py +++ b/timecopilot/agent.py @@ -8,6 +8,7 @@ from pydantic_ai import Agent, ModelRetry, RunContext from pydantic_ai.agent import AgentRunResult from pydantic_ai.models import Model +from pydantic_ai.models.openai import OpenAIModel from rich.console import Console from rich.panel import Panel from rich.table import Table @@ -34,6 +35,35 @@ from .forecaster import Forecaster, TimeCopilotForecaster from .models.adapters.sktime import SKTimeAdapter + + +def _resolve_llm(llm: str | Model) -> str | Model: + """Resolve LLM string, adding LiteLLM support for 100+ providers. + + If the string starts with 'litellm:', creates a pydantic-ai OpenAI model + backed by litellm SDK (no proxy needed). LiteLLM reads provider API keys + from environment variables (ANTHROPIC_API_KEY, OPENAI_API_KEY, etc.). + + Examples: + --llm litellm:anthropic/claude-sonnet-4-6 + --llm litellm:groq/llama-3.3-70b-versatile + --llm litellm:deepseek/deepseek-chat + """ + if not (isinstance(llm, str) and llm.startswith("litellm:")): + return llm + + model_name = llm[len("litellm:"):] + try: + import litellm + + litellm.drop_params = True + client = litellm.AsyncOpenAI() + except ImportError: + raise ImportError( + "litellm package required for litellm: models. " + "Install with: pip install 'timecopilot[litellm]'" + ) + return OpenAIModel(model_name, openai_client=client) from .models.prophet import Prophet from .models.stats import ( ADIDA, @@ -541,7 +571,7 @@ def __init__( "model is not allowed to be passed as a keyword argument" "use `llm` instead" ) - self.llm = llm + self.llm = _resolve_llm(llm) self.forecasting_agent = Agent( deps_type=ExperimentDataset,