diff --git a/airflow-wingman/src/airflow_wingman/llms_models.py b/airflow-wingman/src/airflow_wingman/llms_models.py index fd328ab..2e154b4 100644 --- a/airflow-wingman/src/airflow_wingman/llms_models.py +++ b/airflow-wingman/src/airflow_wingman/llms_models.py @@ -1,4 +1,17 @@ MODELS = { + "openai": { + "name": "OpenAI", + "endpoint": "https://api.openai.com/v1/chat/completions", + "models": [ + { + "id": "gpt-4o", + "name": "GPT-4o", + "default": True, + "context_window": 128000, + "description": "Input $5/M tokens, Output $15/M tokens", + } + ], + }, "anthropic": { "name": "Anthropic", "endpoint": "https://api.anthropic.com/v1/messages", @@ -24,18 +37,11 @@ MODELS = { "endpoint": "https://openrouter.ai/api/v1/chat/completions", "models": [ { - "id": "anthropic/claude-3.5-sonnet", - "name": "Claude 3.5 Sonnet", + "id": "custom", + "name": "Custom Model", "default": False, - "context_window": 200000, - "description": "Input $3/M tokens, Output $15/M tokens", - }, - { - "id": "anthropic/claude-3.5-haiku", - "name": "Claude 3.5 Haiku", - "default": False, - "context_window": 200000, - "description": "Input $0.80/M tokens, Output $4/M tokens", + "context_window": 128000, # Default context window, will be updated based on model + "description": "Enter any model name supported by OpenRouter (e.g., 'anthropic/claude-3-opus', 'meta-llama/llama-2-70b')", }, ], }, diff --git a/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html b/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html index 12d69be..6c6360b 100644 --- a/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html +++ b/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html @@ -21,7 +21,7 @@