diff --git a/.env.example b/.env.example index 291b350d..06adc87a 100644 --- a/.env.example +++ b/.env.example @@ -29,6 +29,17 @@ 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. +# Provider API keys are read from standard env vars (OPENAI_API_KEY, ANTHROPIC_API_KEY, etc.) +LITELLM_MODEL="gpt-4o-mini" + # 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..99653d21 100644 --- a/Backend/gpt.py +++ b/Backend/gpt.py @@ -16,6 +16,15 @@ 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_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: return Client(host=OLLAMA_BASE_URL, timeout=OLLAMA_TIMEOUT) @@ -31,7 +40,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,22 +106,47 @@ def list_ollama_models() -> Tuple[List[str], str]: return unique_names, default_model +def _litellm_generate_response(prompt: str, model_name: str) -> str: + """Generate a response using the LiteLLM SDK.""" + import litellm + + try: + 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 + + 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 + + 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. - """ - model_name = (ai_model or "").strip() or OLLAMA_MODEL + default_model = LITELLM_MODEL if LLM_PROVIDER == "litellm" else OLLAMA_MODEL + model_name = (ai_model or "").strip() or default_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, } ) 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",