From 8ef35be49e30f582f8bce7fb9e4a7eae0e6ba721 Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Tue, 2 Jun 2026 02:34:33 +0530 Subject: [PATCH 1/2] feat: add LiteLLM as LLM provider --- .env.example | 13 ++++++++++ Backend/gpt.py | 69 +++++++++++++++++++++++++++++++++++++++++++++++++- pyproject.toml | 3 +++ 3 files changed, 84 insertions(+), 1 deletion(-) diff --git a/.env.example b/.env.example index 291b350d..2ac52f7c 100644 --- a/.env.example +++ b/.env.example @@ -29,6 +29,19 @@ OLLAMA_BASE_URL="http://localhost:11434" # Pull a model first: ollama pull llama3.1:8b OLLAMA_MODEL="llama3.1:8b" +# LLM Provider +# Set to "litellm" to use LiteLLM SDK instead of Ollama. +# Supports 100+ providers: OpenAI, Anthropic, Bedrock, Vertex, Groq, etc. +# Install with: pip install moneyprinter[litellm] +LLM_PROVIDER="ollama" + +# LiteLLM Settings (only used when LLM_PROVIDER=litellm) +# Model uses LiteLLM naming: "gpt-4o-mini", "anthropic/claude-sonnet-4-20250514", etc. +LITELLM_MODEL="gpt-4o-mini" +# API key for your chosen provider (or leave empty to read from provider env vars +# like OPENAI_API_KEY, ANTHROPIC_API_KEY, etc.) +LITELLM_API_KEY="" + # AssemblyAI API Key # Sign up at https://www.assemblyai.com/ to receive an API key. ASSEMBLY_AI_API_KEY="" diff --git a/Backend/gpt.py b/Backend/gpt.py index d2bef57d..4b594527 100644 --- a/Backend/gpt.py +++ b/Backend/gpt.py @@ -16,6 +16,11 @@ OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3.1:8b") OLLAMA_TIMEOUT = float(os.getenv("OLLAMA_TIMEOUT", "180")) +# LLM provider: "ollama" (default) or "litellm" +LLM_PROVIDER = os.getenv("LLM_PROVIDER", "ollama").lower() +LITELLM_MODEL = os.getenv("LITELLM_MODEL", "gpt-4o-mini") +LITELLM_API_KEY = os.getenv("LITELLM_API_KEY") + def _ollama_client() -> Client: return Client(host=OLLAMA_BASE_URL, timeout=OLLAMA_TIMEOUT) @@ -31,7 +36,40 @@ def _extract_model_name(model_obj) -> str: return "" -def list_ollama_models() -> Tuple[List[str], str]: +def list_models() -> Tuple[List[str], str]: + """ + Returns available model names and configured default model. + Delegates to Ollama or LiteLLM depending on LLM_PROVIDER. + + Returns: + Tuple[List[str], str]: (available model names, default model) + """ + if LLM_PROVIDER == "litellm": + return _list_litellm_models() + return _list_ollama_models() + + +def _list_litellm_models() -> Tuple[List[str], str]: + """Returns a curated list of popular LiteLLM model identifiers.""" + models = [ + LITELLM_MODEL, + "gpt-4.1", + "gpt-4.1-mini", + "anthropic/claude-sonnet-4-20250514", + "anthropic/claude-haiku-4-5-20251001", + "groq/llama-3.3-70b-versatile", + "bedrock/anthropic.claude-sonnet-4-20250514-v1:0", + "vertex_ai/gemini-2.5-flash", + ] + unique = list(dict.fromkeys(models)) + return unique, LITELLM_MODEL + + +# Keep old name as alias for backward compatibility +list_ollama_models = list_models + + +def _list_ollama_models() -> Tuple[List[str], str]: """ Returns available Ollama model names and configured default model. @@ -64,6 +102,32 @@ def list_ollama_models() -> Tuple[List[str], str]: return unique_names, default_model +def _litellm_generate_response(prompt: str, ai_model: str) -> str: + """Generate a response using the LiteLLM SDK.""" + import litellm + + model_name = (ai_model or "").strip() or LITELLM_MODEL + + params = { + "model": model_name, + "messages": [{"role": "user", "content": prompt}], + "stream": False, + "drop_params": True, + } + if LITELLM_API_KEY: + params["api_key"] = LITELLM_API_KEY + + try: + response = litellm.completion(**params) + except Exception as err: + raise RuntimeError(f"LiteLLM request failed: {err}") from err + + content = response.choices[0].message.content + if not content or not content.strip(): + raise RuntimeError("LiteLLM returned an empty response.") + return content.strip() + + def generate_response(prompt: str, ai_model: str) -> str: """ Generate a script for a video, depending on the subject of the video. @@ -79,6 +143,9 @@ def generate_response(prompt: str, ai_model: str) -> str: """ + if LLM_PROVIDER == "litellm": + return _litellm_generate_response(prompt, ai_model) + model_name = (ai_model or "").strip() or OLLAMA_MODEL try: diff --git a/pyproject.toml b/pyproject.toml index de7414a4..feb234c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,9 @@ dependencies = [ "psycopg[binary]==3.2.3", ] +[project.optional-dependencies] +litellm = ["litellm>=1.80.0,<1.87.0"] + [dependency-groups] dev = [ "pytest==8.4.1", From b92f9331c4989e33aef288abd921e1ef89a0dc44 Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Tue, 2 Jun 2026 02:43:28 +0530 Subject: [PATCH 2/2] fix: add timeout, bounds check, provider validation, clean up main.py --- .env.example | 4 +--- Backend/gpt.py | 52 ++++++++++++++++++++++++------------------------- Backend/main.py | 13 +++++++------ 3 files changed, 34 insertions(+), 35 deletions(-) diff --git a/.env.example b/.env.example index 2ac52f7c..06adc87a 100644 --- a/.env.example +++ b/.env.example @@ -37,10 +37,8 @@ LLM_PROVIDER="ollama" # LiteLLM Settings (only used when LLM_PROVIDER=litellm) # Model uses LiteLLM naming: "gpt-4o-mini", "anthropic/claude-sonnet-4-20250514", etc. +# Provider API keys are read from standard env vars (OPENAI_API_KEY, ANTHROPIC_API_KEY, etc.) LITELLM_MODEL="gpt-4o-mini" -# API key for your chosen provider (or leave empty to read from provider env vars -# like OPENAI_API_KEY, ANTHROPIC_API_KEY, etc.) -LITELLM_API_KEY="" # AssemblyAI API Key # Sign up at https://www.assemblyai.com/ to receive an API key. diff --git a/Backend/gpt.py b/Backend/gpt.py index 4b594527..99653d21 100644 --- a/Backend/gpt.py +++ b/Backend/gpt.py @@ -19,7 +19,11 @@ # LLM provider: "ollama" (default) or "litellm" LLM_PROVIDER = os.getenv("LLM_PROVIDER", "ollama").lower() LITELLM_MODEL = os.getenv("LITELLM_MODEL", "gpt-4o-mini") -LITELLM_API_KEY = os.getenv("LITELLM_API_KEY") +LITELLM_TIMEOUT = float(os.getenv("LITELLM_TIMEOUT", str(OLLAMA_TIMEOUT))) + +if LLM_PROVIDER not in ("ollama", "litellm"): + log(f"[!] Unknown LLM_PROVIDER '{LLM_PROVIDER}', falling back to 'ollama'", "warning") + LLM_PROVIDER = "ollama" def _ollama_client() -> Client: @@ -102,51 +106,47 @@ def _list_ollama_models() -> Tuple[List[str], str]: return unique_names, default_model -def _litellm_generate_response(prompt: str, ai_model: str) -> str: +def _litellm_generate_response(prompt: str, model_name: str) -> str: """Generate a response using the LiteLLM SDK.""" import litellm - model_name = (ai_model or "").strip() or LITELLM_MODEL - - params = { - "model": model_name, - "messages": [{"role": "user", "content": prompt}], - "stream": False, - "drop_params": True, - } - if LITELLM_API_KEY: - params["api_key"] = LITELLM_API_KEY - try: - response = litellm.completion(**params) + response = litellm.completion( + model=model_name, + messages=[{"role": "user", "content": prompt}], + stream=False, + drop_params=True, + timeout=LITELLM_TIMEOUT, + ) except Exception as err: raise RuntimeError(f"LiteLLM request failed: {err}") from err - content = response.choices[0].message.content - if not content or not content.strip(): + if not response.choices: + raise RuntimeError("LiteLLM returned no choices.") + + content = (response.choices[0].message.content or "").strip() + if not content: raise RuntimeError("LiteLLM returned an empty response.") - return content.strip() + return content def generate_response(prompt: str, ai_model: str) -> str: """ - Generate a script for a video, depending on the subject of the video. + Generate a response from the configured LLM provider. Args: - video_subject (str): The subject of the video. - ai_model (str): The AI model to use for generation. - + prompt (str): The prompt to send to the model. + ai_model (str): The model identifier to use. Returns: - str: The response from the AI model. - """ - if LLM_PROVIDER == "litellm": - return _litellm_generate_response(prompt, ai_model) + default_model = LITELLM_MODEL if LLM_PROVIDER == "litellm" else OLLAMA_MODEL + model_name = (ai_model or "").strip() or default_model - model_name = (ai_model or "").strip() or OLLAMA_MODEL + if LLM_PROVIDER == "litellm": + return _litellm_generate_response(prompt, model_name) try: client = _ollama_client() diff --git a/Backend/main.py b/Backend/main.py index 7f46c7ae..6c3bdbd5 100644 --- a/Backend/main.py +++ b/Backend/main.py @@ -6,7 +6,7 @@ from sqlalchemy import and_, case, select from db import SessionLocal, init_db -from gpt import list_ollama_models +from gpt import list_models, LLM_PROVIDER, LITELLM_MODEL from logstream import log from repository import create_job, get_job, list_job_events, request_cancel from utils import ENV_FILE, SONGS_DIR, check_env_vars, clean_dir @@ -26,7 +26,7 @@ @app.route("/api/models", methods=["GET"]) def models(): try: - available_models, default_model = list_ollama_models() + available_models, default_model = list_models() return jsonify( { "status": "success", @@ -35,13 +35,14 @@ def models(): } ) except Exception as err: - log(f"[-] Error fetching Ollama models: {str(err)}", "error") + log(f"[-] Error fetching models: {str(err)}", "error") + fallback = LITELLM_MODEL if LLM_PROVIDER == "litellm" else os.getenv("OLLAMA_MODEL", "llama3.1:8b") return jsonify( { "status": "error", - "message": "Could not fetch Ollama models. Is Ollama running?", - "models": [os.getenv("OLLAMA_MODEL", "llama3.1:8b")], - "default": os.getenv("OLLAMA_MODEL", "llama3.1:8b"), + "message": f"Could not fetch models from {LLM_PROVIDER}.", + "models": [fallback], + "default": fallback, } )