diff --git a/.github/workflows/integration_tests.yml b/.github/workflows/integration_tests.yml index 90c533c697..4e941215ec 100644 --- a/.github/workflows/integration_tests.yml +++ b/.github/workflows/integration_tests.yml @@ -87,7 +87,7 @@ jobs: SECRET_KEY: ${{ secrets.SECRET_KEY }} run: | poetry run ./manage.py migrate - PLAYWRIGHT_BROWSERS_PATH=.venv scripts/run_integration_tests.sh + PLAYWRIGHT_BROWSERS_PATH=.venv scripts/tests/run_integration_tests.sh - name: Create trace.zip if: always() run: | diff --git a/Dockerfile b/Dockerfile index 2f625b7c4e..12ca0faf23 100644 --- a/Dockerfile +++ b/Dockerfile @@ -45,7 +45,7 @@ RUN export NODE_VERSION=$(cat /tmp/.nvmrc) && cd /tmp/ && curl -fsSLO --compress && ln -s /usr/local/bin/node /usr/local/bin/nodejs - # This is done to copy only the source code from HEAD into the image +# This is done to copy only the source code from HEAD into the image RUN --mount=type=bind,source=.git/,target=/tmp/app/.git/ \ git clone /tmp/app/.git/ /app/ @@ -69,5 +69,5 @@ EXPOSE 3000 ARG ENTRY_SCRIPT_ARG="scripts/prod/startapp.sh" ENV ENTRY_SCRIPT=${ENTRY_SCRIPT_ARG} -CMD ["sh","-c","${ENTRY_SCRIPT}"] +CMD ["sh", "-c", "${ENTRY_SCRIPT}"] diff --git a/fab_credits/__init__.py b/fab_credits/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/fab_credits/admin.py b/fab_credits/admin.py new file mode 100644 index 0000000000..ae2e59c64f --- /dev/null +++ b/fab_credits/admin.py @@ -0,0 +1,10 @@ +from django.contrib import admin + +from . import models + + +@admin.register(models.UserUsage) +class UserUsageAdmin(admin.ModelAdmin): + autocomplete_fields = ("user",) + search_fields = ["user__username", "user__email"] + readonly_fields = ["input_tokens_used", "output_tokens_used"] diff --git a/fab_credits/apps.py b/fab_credits/apps.py new file mode 100644 index 0000000000..6621d888f0 --- /dev/null +++ b/fab_credits/apps.py @@ -0,0 +1,6 @@ +from django.apps import AppConfig + + +class FabCreditsConfig(AppConfig): + default_auto_field = "django.db.models.BigAutoField" + name = "fab_credits" diff --git a/fab_credits/migrations/0001_initial.py b/fab_credits/migrations/0001_initial.py new file mode 100644 index 0000000000..987e1bea0a --- /dev/null +++ b/fab_credits/migrations/0001_initial.py @@ -0,0 +1,70 @@ +# Generated by Django 4.2.11 on 2024-07-18 10:43 + +from django.conf import settings +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + + initial = True + + dependencies = [ + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ] + + operations = [ + migrations.CreateModel( + name="UserUsage", + fields=[ + ( + "id", + models.BigAutoField( + auto_created=True, + primary_key=True, + serialize=False, + verbose_name="ID", + ), + ), + ( + "input_tokens_used", + models.IntegerField( + default=0, + help_text="Keeps track of input token used. Filled in automatically", + ), + ), + ( + "output_tokens_used", + models.IntegerField( + default=0, + help_text="Keeps track of output token used. Filled in automatically", + ), + ), + ( + "total_allowed_tokens", + models.IntegerField( + default=0, + help_text="Used by admins to set the total number of input&output tokens the user can use.", + ), + ), + ( + "platform", + models.CharField( + choices=[("OPENAI", "Openai"), ("ANTHROPIC", "Anthropic")], + default="OPENAI", + max_length=50, + ), + ), + ( + "user", + models.ForeignKey( + on_delete=django.db.models.deletion.CASCADE, + to=settings.AUTH_USER_MODEL, + ), + ), + ], + options={ + "unique_together": {("user", "platform")}, + }, + ), + ] diff --git a/fab_credits/migrations/0002_alter_userusage_unique_together_userusage_model_name_and_more.py b/fab_credits/migrations/0002_alter_userusage_unique_together_userusage_model_name_and_more.py new file mode 100644 index 0000000000..40eb9608ef --- /dev/null +++ b/fab_credits/migrations/0002_alter_userusage_unique_together_userusage_model_name_and_more.py @@ -0,0 +1,41 @@ +# Generated by Django 4.2.11 on 2024-07-18 12:44 + +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ("fab_credits", "0001_initial"), + ] + + operations = [ + migrations.AlterUniqueTogether( + name="userusage", + unique_together=set(), + ), + migrations.AddField( + model_name="userusage", + name="model_name", + field=models.CharField( + choices=[ + ("gpt-4o", "gpt-4o"), + ("gpt-4o-2024-05-13", "gpt-4o-2024-05-13"), + ("gpt-3.5-turbo-0125", "gpt-3.5-turbo-0125"), + ("gpt-3.5-turbo-instruct", "gpt-3.5-turbo-instruct"), + ("claude-3-5-sonnet-20240620", "claude-3-5-sonnet-20240620"), + ("claude-3-opus-20240229", "claude-3-opus-20240229"), + ("claude-3-sonnet-20240229", "claude-3-sonnet-20240229"), + ("claude-3-haiku-20240307", "claude-3-haiku-20240307"), + ], + default="gpt-4o", + max_length=50, + ), + ), + migrations.AlterUniqueTogether( + name="userusage", + unique_together={("user", "platform", "model_name")}, + ), + ] diff --git a/fab_credits/migrations/__init__.py b/fab_credits/migrations/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/fab_credits/models.py b/fab_credits/models.py new file mode 100644 index 0000000000..200e429324 --- /dev/null +++ b/fab_credits/models.py @@ -0,0 +1,45 @@ +from django.db import models + +from users.models import User + +available_models = [ + ("gpt-4o", "gpt-4o"), + ("gpt-4o-2024-05-13", "gpt-4o-2024-05-13"), + ("gpt-3.5-turbo-0125", "gpt-3.5-turbo-0125"), + ("gpt-3.5-turbo-instruct", "gpt-3.5-turbo-instruct"), + ("claude-3-5-sonnet-20240620", "claude-3-5-sonnet-20240620"), + ("claude-3-opus-20240229", "claude-3-opus-20240229"), + ("claude-3-sonnet-20240229", "claude-3-sonnet-20240229"), + ("claude-3-haiku-20240307", "claude-3-haiku-20240307"), +] + + +# Create your models here. +class UserUsage(models.Model): + class UsagePlatform(models.TextChoices): + OpenAI = "OPENAI" + Anthropic = "ANTHROPIC" + + user = models.ForeignKey(User, on_delete=models.CASCADE) + input_tokens_used = models.IntegerField( + default=0, help_text="Keeps track of input token used. Filled in automatically" + ) + output_tokens_used = models.IntegerField( + default=0, help_text="Keeps track of output token used. Filled in automatically" + ) + total_allowed_tokens = models.IntegerField( + default=0, + help_text="Used by admins to set the total number of input&output tokens the user can use.", + ) + platform = models.CharField( + max_length=50, choices=UsagePlatform.choices, default=UsagePlatform.OpenAI + ) + model_name = models.CharField( + max_length=50, default="gpt-4o", choices=available_models + ) + + class Meta: + unique_together = ("user", "platform", "model_name") + + def __str__(self): + return f"{self.user.username}({self.user.email}) - {self.platform}/{self.model_name}" diff --git a/fab_credits/tests.py b/fab_credits/tests.py new file mode 100644 index 0000000000..a39b155ac3 --- /dev/null +++ b/fab_credits/tests.py @@ -0,0 +1 @@ +# Create your tests here. diff --git a/fab_credits/urls.py b/fab_credits/urls.py new file mode 100644 index 0000000000..1da207449e --- /dev/null +++ b/fab_credits/urls.py @@ -0,0 +1,16 @@ +from django.urls import path + +from . import views + +urlpatterns = [ + path( + "anthropic/v1/messages", + views.anthropic_v1_messages, + name="anthropic_v1_messages", + ), + path( + "openai/v1/chat/completions", + views.openai_v1_chat_completions, + name="anthropic_v1_messages", + ), +] diff --git a/fab_credits/views.py b/fab_credits/views.py new file mode 100644 index 0000000000..e99b15ab95 --- /dev/null +++ b/fab_credits/views.py @@ -0,0 +1,224 @@ +import json +import logging +from typing import Any, Callable # noqa: UP035 + +import requests +from django.conf import settings +from django.http import JsonResponse, StreamingHttpResponse +from rest_framework.decorators import api_view, permission_classes +from rest_framework.permissions import IsAuthenticated + +from .models import UserUsage + + +def streaming_response( + url, + headers, + data, + platform: UserUsage.UsagePlatform, + user_usage, + get_io_tokens_streaming_fn: Callable[dict[Any], tuple[int | None, int | None]], +): + + if platform == UserUsage.UsagePlatform.OpenAI: + data["stream_options"] = {"include_usage": True} + + platform_response = requests.post(url, headers=headers, json=data, stream=True) + + platform_response.raise_for_status() + + def resp_iterator(response): + for line in response.iter_lines(): + if line: + line = f"\n{line.decode('utf-8')}\n" + input_tokens, output_tokens = get_io_tokens_streaming_fn(line) + if input_tokens is not None or output_tokens is not None: + user_usage.input_tokens_used += input_tokens or 0 + user_usage.output_tokens_used += output_tokens or 0 + user_usage.save() + yield line + + return StreamingHttpResponse( + resp_iterator(platform_response), content_type="text/event-stream" + ) + + +def normal_response( + url, + headers, + data, + platform: UserUsage.UsagePlatform, + user_usage, + get_io_tokens_fn: Callable[dict[Any], tuple[int, int]], +): + platform_response = requests.post( + url, + headers=headers, + json=data, + ) + + platform_response.raise_for_status() + + response_data = platform_response.json() + + input_tokens, output_tokens = get_io_tokens_fn(response_data) + + user_usage.input_tokens_used += input_tokens + user_usage.output_tokens_used += output_tokens + + user_usage.save() + + return JsonResponse(response_data) + + +def make_request( + request, + platform: UserUsage.UsagePlatform, + model_name: str, + headers: dict[str, str], + url: str, + get_io_tokens_fn: Callable[dict[Any], tuple[int, int]], + get_io_tokens_streaming_fn: Callable[dict[Any], tuple[int, int]], +): + user_usage = UserUsage.objects.filter( + user=request.user, platform=platform, model_name=model_name + ).first() + + if user_usage is None: + return JsonResponse( + { + "error": f"You don't have an allowance for model <{model_name}> on <{platform.label}> ." + }, + status=400, + ) + + if ( + user_usage.input_tokens_used + user_usage.output_tokens_used + >= user_usage.total_allowed_tokens + ): + return JsonResponse( + {"error": "You have exceeded your credit allowance."}, status=400 + ) + + try: + data = json.loads(request.body) + streaming_mode = data.get("stream", False) + + if streaming_mode: + return streaming_response( + url, headers, data, platform, user_usage, get_io_tokens_streaming_fn + ) + else: + return normal_response( + url, headers, data, platform, user_usage, get_io_tokens_fn + ) + + except json.JSONDecodeError: + return JsonResponse({"error": "Invalid JSON data"}, status=400) + except requests.exceptions.RequestException as e: + error_msg = f"Error forwarding request to {'Anthropic' if platform == UserUsage.UsagePlatform.Anthropic else 'OpenAI'} API: {e}" + logging.error(error_msg) + return JsonResponse( + {"error": error_msg}, + status=502, + ) + except Exception as e: + return JsonResponse({"error": f"An unexpected error occurred: {e}"}, status=500) + + +@api_view(["POST"]) +@permission_classes([IsAuthenticated]) +def openai_v1_chat_completions(request): + + headers = {**request.headers} + headers["Authorization"] = f"Bearer {settings.FAB_CREDITS_OPENAI_API_KEY}" + headers.pop("Host") + + data = json.loads(request.body) + model_name = data.get("model", None) + + def get_io_tokens_fn(data): + return ( + data.get("usage").get("prompt_tokens"), + data.get("usage").get("completion_tokens"), + ) + + def get_io_tokens_streaming_fn(line): + line = line.strip() + + data_str = line[len("data: ") :] + + try: + data = json.loads(data_str) + + if data["usage"] is None: + return None, None + + return get_io_tokens_fn(data) + except json.JSONDecodeError: + # last line is not a valid json, but a [DONE] token after Data + return None, None + + response = make_request( + request, + UserUsage.UsagePlatform.OpenAI, + headers=headers, + url="https://api.openai.com/v1/chat/completions", + model_name=model_name, + get_io_tokens_fn=get_io_tokens_fn, + get_io_tokens_streaming_fn=get_io_tokens_streaming_fn, + ) + + return response + + +@api_view(["POST"]) +@permission_classes([IsAuthenticated]) +def anthropic_v1_messages(request): + + headers = {**request.headers} + headers["x-api-key"] = settings.FAB_CREDITS_ANTHROPIC_API_KEY + headers.pop("Authorization") + headers.pop("Host") + + data = json.loads(request.body) + model_name = data.get("model", None) + + def get_io_tokens_fn(data): + return ( + data.get("usage").get("input_tokens"), + data.get("usage").get("output_tokens"), + ) + + def get_io_tokens_streaming_fn(line): + line = line.strip() + if not line.startswith("data: "): + return None, None + + data_str = line[len("data: ") :] + + try: + data = json.loads(data_str) + + usage_data = data.get("usage", None) or data.get("message", {}).get( + "usage", None + ) + if usage_data is None: + return None, None + + return usage_data.get("input_tokens"), usage_data.get("output_tokens") + except json.JSONDecodeError as e: + logging.error("Invalid json data: ", data_str) + raise e + + response = make_request( + request=request, + platform=UserUsage.UsagePlatform.Anthropic, + headers=headers, + url="https://api.anthropic.com/v1/messages", + get_io_tokens_fn=get_io_tokens_fn, + get_io_tokens_streaming_fn=get_io_tokens_streaming_fn, + model_name=model_name, + ) + + return response diff --git a/metaculus_web/settings.py b/metaculus_web/settings.py index 89e81ba7d4..2f2e68dd0c 100644 --- a/metaculus_web/settings.py +++ b/metaculus_web/settings.py @@ -65,6 +65,7 @@ "comments", "notifications", "fab_management", + "fab_credits", ] @@ -293,7 +294,13 @@ "BACKEND": "storages.backends.s3.S3Storage", "OPTIONS": {}, }, - "staticfiles": {"BACKEND": "whitenoise.storage.CompressedManifestStaticFilesStorage"}, + "staticfiles": { + "BACKEND": ( + "whitenoise.storage.CompressedManifestStaticFilesStorage" + if not DEBUG + else "django.contrib.staticfiles.storage.StaticFilesStorage" + ) + }, } ITN_DB_MACHINE_SSH_ADDR = os.environ.get("ITN_DB_MACHINE_SSH_ADDR") @@ -326,6 +333,9 @@ GOOGLE_CREDEBTIALS_FAB_SHEET_B64 = os.environ.get("GOOGLE_CREDEBTIALS_FAB_SHEET_B64") +FAB_CREDITS_ANTHROPIC_API_KEY = os.environ.get("FAB_CREDITS_ANTHROPIC_API_KEY") +FAB_CREDITS_OPENAI_API_KEY = os.environ.get("FAB_CREDITS_OPENAI_API_KEY") + ALLOWED_HOSTS = [".metaculus.com", "localhost", "127.0.0.1"] CSRF_TRUSTED_ORIGINS = [FRONTEND_BASE_URL] diff --git a/metaculus_web/urls.py b/metaculus_web/urls.py index 0829a2f745..36e3f00f5f 100644 --- a/metaculus_web/urls.py +++ b/metaculus_web/urls.py @@ -35,6 +35,7 @@ path("api/", include("comments.urls")), path("api/", include("scoring.urls")), path("api/", include("misc.urls")), + path("proxy/", include("fab_credits.urls")), # Backward compatibility endpoints path("api2/", include(comments.urls.old_api)), path("api2/", include(posts.urls.old_api)), diff --git a/migrator/management/commands/migrate_old_db.py b/migrator/management/commands/migrate_old_db.py index f3192d6fc6..c974f541e0 100644 --- a/migrator/management/commands/migrate_old_db.py +++ b/migrator/management/commands/migrate_old_db.py @@ -22,6 +22,7 @@ from migrator.services.migrate_users import migrate_users from migrator.services.migrate_votes import migrate_votes from migrator.services.post_migrate import post_migrate_calculate_divergence +from migrator.services.migrate_fab_credits import migrate_fab_credits from migrator.utils import reset_sequence from posts.jobs import job_compute_movement from posts.services.common import compute_hotness @@ -60,6 +61,7 @@ def handle(self, *args, site_ids=None, **options): # main model migration migrate_users() print("Migrated users") + migrate_fab_credits() migrate_questions(site_ids=site_ids) print("Migrated questions") migrate_projects(site_ids=site_ids) diff --git a/migrator/services/migrate_fab_credits.py b/migrator/services/migrate_fab_credits.py new file mode 100644 index 0000000000..c14ab06613 --- /dev/null +++ b/migrator/services/migrate_fab_credits.py @@ -0,0 +1,13 @@ +from ..utils import paginated_query + +from fab_credits.models import UserUsage + + +def migrate_fab_credits(): + entries = UserUsage.objects.bulk_create( + [ + UserUsage(**usage) + for usage in paginated_query("SELECT * FROM fab_credits_userusage") + ] + ) + print(f"Migrated {len(entries)} Fab UserUsage entries") diff --git a/scripts/run_integration_tests.sh b/scripts/tests/run_integration_tests.sh similarity index 100% rename from scripts/run_integration_tests.sh rename to scripts/tests/run_integration_tests.sh