From 923789c9fbe62ccae5b7f665a5f26862d5229860 Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Mon, 18 May 2026 00:02:21 +0530 Subject: [PATCH 1/2] feat: add LiteLLM as LLM provider --- linkedin/llm.py | 15 ++++++++++ .../0008_siteconfig_add_litellm_provider.py | 29 +++++++++++++++++++ linkedin/models.py | 1 + 3 files changed, 45 insertions(+) create mode 100644 linkedin/migrations/0008_siteconfig_add_litellm_provider.py diff --git a/linkedin/llm.py b/linkedin/llm.py index c32af490..665c22d9 100644 --- a/linkedin/llm.py +++ b/linkedin/llm.py @@ -140,6 +140,20 @@ def _build_openai_compatible(cfg): )) +def _build_litellm(cfg): + if not cfg.llm_api_base: + raise ValueError("LLM_API_BASE is required for the litellm provider (your LiteLLM proxy URL).") + from openai import AsyncOpenAI + from pydantic_ai.models.openai import OpenAIModel + from pydantic_ai.providers.openai import OpenAIProvider + client = AsyncOpenAI( + api_key=cfg.llm_api_key or "unused", + base_url=cfg.llm_api_base, + max_retries=_MAX_RETRIES, + ) + return OpenAIModel(cfg.ai_model, provider=OpenAIProvider(openai_client=client)) + + _PROVIDER_BUILDERS: dict[str, Callable] = { "openai": _build_openai, "anthropic": _build_anthropic, @@ -148,6 +162,7 @@ def _build_openai_compatible(cfg): "mistral": _build_mistral, "cohere": _build_cohere, "openai_compatible": _build_openai_compatible, + "litellm": _build_litellm, } diff --git a/linkedin/migrations/0008_siteconfig_add_litellm_provider.py b/linkedin/migrations/0008_siteconfig_add_litellm_provider.py new file mode 100644 index 00000000..1e9b88e3 --- /dev/null +++ b/linkedin/migrations/0008_siteconfig_add_litellm_provider.py @@ -0,0 +1,29 @@ +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("linkedin", "0007_siteconfig_llm_provider"), + ] + + operations = [ + migrations.AlterField( + model_name="siteconfig", + name="llm_provider", + field=models.CharField( + choices=[ + ("openai", "OpenAI"), + ("anthropic", "Anthropic"), + ("google", "Google"), + ("groq", "Groq"), + ("mistral", "Mistral"), + ("cohere", "Cohere"), + ("openai_compatible", "OpenAI-compatible"), + ("litellm", "LiteLLM"), + ], + default="openai", + max_length=32, + ), + ), + ] diff --git a/linkedin/models.py b/linkedin/models.py index a0426021..7f0fbd8d 100644 --- a/linkedin/models.py +++ b/linkedin/models.py @@ -28,6 +28,7 @@ class LLMProvider(models.TextChoices): MISTRAL = "mistral", "Mistral" COHERE = "cohere", "Cohere" OPENAI_COMPATIBLE = "openai_compatible", "OpenAI-compatible" + LITELLM = "litellm", "LiteLLM" llm_provider = models.CharField( max_length=32, From 7898466af0dee9841cb11d3f7200da6f9a2082a5 Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Mon, 18 May 2026 00:18:18 +0530 Subject: [PATCH 2/2] fix: skip API key validation for litellm provider + add tests --- linkedin/llm.py | 5 +- tests/test_litellm_provider.py | 107 +++++++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm_provider.py diff --git a/linkedin/llm.py b/linkedin/llm.py index 665c22d9..7db0a878 100644 --- a/linkedin/llm.py +++ b/linkedin/llm.py @@ -168,12 +168,15 @@ def _build_litellm(cfg): # ── Model factory ──────────────────────────────────────────────────── +_PROVIDERS_WITHOUT_API_KEY = {"litellm"} + + def _validated_site_config(): """Load `SiteConfig` and assert the required LLM fields are populated.""" from linkedin.models import SiteConfig cfg = SiteConfig.load() - if not cfg.llm_api_key: + if not cfg.llm_api_key and cfg.llm_provider not in _PROVIDERS_WITHOUT_API_KEY: raise ValueError("LLM_API_KEY is not set in Site Configuration.") if not cfg.ai_model: raise ValueError("AI_MODEL is not set in Site Configuration.") diff --git a/tests/test_litellm_provider.py b/tests/test_litellm_provider.py new file mode 100644 index 00000000..367f1a07 --- /dev/null +++ b/tests/test_litellm_provider.py @@ -0,0 +1,107 @@ +"""Tests for the LiteLLM provider integration.""" +from unittest.mock import MagicMock, patch + +import pytest + + +@pytest.fixture +def site_config(db): + from linkedin.models import SiteConfig + + cfg = SiteConfig.load() + cfg.llm_provider = "litellm" + cfg.ai_model = "anthropic/claude-sonnet-4-20250514" + cfg.llm_api_base = "http://localhost:4000/v1" + cfg.llm_api_key = "" + cfg.save() + return cfg + + +class TestLiteLLMProviderChoice: + def test_litellm_in_choices(self): + from linkedin.models import SiteConfig + + values = [c[0] for c in SiteConfig.LLMProvider.choices] + assert "litellm" in values + + def test_litellm_label(self): + from linkedin.models import SiteConfig + + labels = dict(SiteConfig.LLMProvider.choices) + assert labels["litellm"] == "LiteLLM" + + +class TestBuildLiteLLM: + def test_builds_openai_model(self, site_config): + from linkedin.llm import _build_litellm + + model = _build_litellm(site_config) + assert model is not None + assert model.model_name == "anthropic/claude-sonnet-4-20250514" + + def test_requires_api_base(self, site_config): + from linkedin.llm import _build_litellm + + site_config.llm_api_base = "" + with pytest.raises(ValueError, match="LLM_API_BASE is required"): + _build_litellm(site_config) + + def test_api_key_optional(self, site_config): + from linkedin.llm import _build_litellm + + site_config.llm_api_key = "" + model = _build_litellm(site_config) + assert model is not None + + def test_api_key_forwarded_when_set(self, site_config): + from linkedin.llm import _build_litellm + + site_config.llm_api_key = "sk-litellm-test-key" + model = _build_litellm(site_config) + assert model is not None + + +class TestValidatedSiteConfig: + def test_litellm_skips_api_key_check(self, site_config): + from linkedin.llm import _validated_site_config + + site_config.llm_api_key = "" + site_config.save() + cfg = _validated_site_config() + assert cfg.llm_provider == "litellm" + + def test_litellm_still_requires_model(self, site_config): + from linkedin.llm import _validated_site_config + + site_config.ai_model = "" + site_config.save() + with pytest.raises(ValueError, match="AI_MODEL is not set"): + _validated_site_config() + + def test_openai_still_requires_api_key(self, db): + from linkedin.models import SiteConfig + from linkedin.llm import _validated_site_config + + cfg = SiteConfig.load() + cfg.llm_provider = "openai" + cfg.ai_model = "gpt-4o" + cfg.llm_api_key = "" + cfg.save() + with pytest.raises(ValueError, match="LLM_API_KEY is not set"): + _validated_site_config() + + +class TestGetLLMModel: + def test_litellm_end_to_end(self, site_config): + from linkedin.llm import get_llm_model + + model = get_llm_model() + assert model is not None + + def test_litellm_with_hosted_key(self, site_config): + from linkedin.llm import get_llm_model + + site_config.llm_api_key = "sk-hosted-proxy-key" + site_config.save() + model = get_llm_model() + assert model is not None