This commit is contained in:
alorig
2025-12-24 01:58:22 +05:00
60 changed files with 12275 additions and 1272 deletions

View File

@@ -4,7 +4,58 @@ Admin configuration for AI models
from django.contrib import admin
from unfold.admin import ModelAdmin
from igny8_core.admin.base import Igny8ModelAdmin
from igny8_core.ai.models import AITaskLog
from igny8_core.ai.models import AITaskLog, IntegrationState
from import_export.admin import ExportMixin
from import_export import resources
@admin.register(IntegrationState)
class IntegrationStateAdmin(Igny8ModelAdmin):
"""Admin interface for Integration States"""
list_display = [
'account',
'is_openai_enabled',
'is_runware_enabled',
'is_image_generation_enabled',
'updated_at',
]
list_filter = [
'is_openai_enabled',
'is_runware_enabled',
'is_image_generation_enabled',
]
search_fields = [
'account__name',
'account__domain',
]
readonly_fields = ['created_at', 'updated_at']
fieldsets = (
('Account', {
'fields': ('account',)
}),
('Integration States', {
'fields': (
'is_openai_enabled',
'is_runware_enabled',
'is_image_generation_enabled',
)
}),
('Timestamps', {
'fields': ('created_at', 'updated_at'),
'classes': ('collapse',)
}),
)
class AITaskLogResource(resources.ModelResource):
"""Resource class for exporting AI Task Logs"""
class Meta:
model = AITaskLog
fields = ('id', 'function_name', 'account__name', 'status', 'phase',
'cost', 'tokens', 'duration', 'created_at')
export_order = fields
from import_export.admin import ExportMixin

View File

@@ -376,16 +376,32 @@ class AIEngine:
# Map function name to operation type
operation_type = self._get_operation_type(function_name)
# Get actual token usage from response (AI returns 'input_tokens' and 'output_tokens')
# Extract token usage from AI response (standardize key names)
tokens_input = raw_response.get('input_tokens', 0)
tokens_output = raw_response.get('output_tokens', 0)
total_tokens = tokens_input + tokens_output
# Deduct credits based on actual token usage
# Get model_config for token-based billing
model_config = None
if tokens_input > 0 or tokens_output > 0:
# Get model from response or use account default
model_config = CreditService.get_model_for_operation(
account=self.account,
operation_type=operation_type,
task_model_override=None # TODO: Support task-level model override
)
# Calculate actual amount based on results (for non-token operations)
actual_amount = self._get_actual_amount(function_name, save_result, parsed, data)
# Deduct credits using token-based calculation if tokens available
CreditService.deduct_credits_for_operation(
account=self.account,
operation_type=operation_type,
amount=actual_amount, # Fallback for non-token operations
tokens_input=tokens_input,
tokens_output=tokens_output,
model_config=model_config,
cost_usd=raw_response.get('cost'),
model_used=raw_response.get('model', ''),
related_object_type=self._get_related_object_type(function_name),
@@ -399,10 +415,7 @@ class AIEngine:
}
)
logger.info(
f"[AIEngine] Credits deducted: {operation_type}, "
f"tokens: {tokens_input + tokens_output} ({tokens_input} in, {tokens_output} out)"
)
logger.info(f"[AIEngine] Credits deducted: {operation_type}, tokens: {total_tokens} ({tokens_input} in, {tokens_output} out), model: {model_config.model_name if model_config else 'legacy'}")
except InsufficientCreditsError as e:
# This shouldn't happen since we checked before, but log it
logger.error(f"[AIEngine] Insufficient credits during deduction: {e}")

View File

@@ -0,0 +1,33 @@
# Generated by Django 5.2.9 on 2025-12-23 12:51
import django.db.models.deletion
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
('ai', '0002_initial'),
('igny8_core_auth', '0018_add_country_remove_intent_seedkeyword'),
]
operations = [
migrations.CreateModel(
name='IntegrationState',
fields=[
('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
('integration_type', models.CharField(choices=[('openai', 'OpenAI'), ('runware', 'Runware'), ('image_generation', 'Image Generation Service')], help_text='Type of integration (openai, runware, image_generation)', max_length=50)),
('is_enabled', models.BooleanField(default=True, help_text='Whether this integration is enabled for this account')),
('created_at', models.DateTimeField(auto_now_add=True)),
('updated_at', models.DateTimeField(auto_now=True)),
('account', models.ForeignKey(help_text='Account that owns this integration state', on_delete=django.db.models.deletion.CASCADE, related_name='integration_states', to='igny8_core_auth.account')),
],
options={
'verbose_name': 'Integration State',
'verbose_name_plural': 'Integration States',
'db_table': 'ai_integration_state',
'indexes': [models.Index(fields=['account', 'integration_type'], name='ai_integrat_account_667460_idx'), models.Index(fields=['integration_type', 'is_enabled'], name='ai_integrat_integra_22ddc7_idx')],
'unique_together': {('account', 'integration_type')},
},
),
]

View File

@@ -0,0 +1,155 @@
# Generated manually on 2025-12-23
from django.db import migrations, models
import django.db.models.deletion
def migrate_data_forward(apps, schema_editor):
"""Convert multiple records per account to single record with 3 fields"""
IntegrationState = apps.get_model('ai', 'IntegrationState')
db_alias = schema_editor.connection.alias
# Get all accounts with integration states
accounts = {}
for state in IntegrationState.objects.using(db_alias).all():
account_id = state.account_id
if account_id not in accounts:
accounts[account_id] = {
'account': state.account,
'is_openai_enabled': True,
'is_runware_enabled': True,
'is_image_generation_enabled': True,
}
# Set the appropriate field based on integration_type
if state.integration_type == 'openai':
accounts[account_id]['is_openai_enabled'] = state.is_enabled
elif state.integration_type == 'runware':
accounts[account_id]['is_runware_enabled'] = state.is_enabled
elif state.integration_type == 'image_generation':
accounts[account_id]['is_image_generation_enabled'] = state.is_enabled
# Store the data for later
return accounts
def migrate_data_backward(apps, schema_editor):
"""Convert single record back to multiple records"""
pass # We'll lose data on rollback, but that's acceptable
class Migration(migrations.Migration):
dependencies = [
('ai', '0003_add_integration_state_model'),
('igny8_core_auth', '0001_initial'),
]
operations = [
# First, remove indexes and constraints
migrations.RemoveIndex(
model_name='integrationstate',
name='ai_integrat_account_667460_idx',
),
migrations.RemoveIndex(
model_name='integrationstate',
name='ai_integrat_integra_22ddc7_idx',
),
migrations.AlterUniqueTogether(
name='integrationstate',
unique_together=set(),
),
# Add new fields (nullable for now)
migrations.AddField(
model_name='integrationstate',
name='is_image_generation_enabled',
field=models.BooleanField(default=True, help_text='Whether Image Generation Service is enabled for this account', null=True),
),
migrations.AddField(
model_name='integrationstate',
name='is_openai_enabled',
field=models.BooleanField(default=True, help_text='Whether OpenAI integration is enabled for this account', null=True),
),
migrations.AddField(
model_name='integrationstate',
name='is_runware_enabled',
field=models.BooleanField(default=True, help_text='Whether Runware integration is enabled for this account', null=True),
),
# Migrate data using SQL
migrations.RunSQL(
sql="""
-- Delete all records, we'll recreate them properly
TRUNCATE TABLE ai_integration_state CASCADE;
""",
reverse_sql=migrations.RunSQL.noop,
),
# Remove old fields
migrations.RemoveField(
model_name='integrationstate',
name='integration_type',
),
migrations.RemoveField(
model_name='integrationstate',
name='is_enabled',
),
# Drop the old primary key
migrations.RunSQL(
sql='ALTER TABLE ai_integration_state DROP CONSTRAINT IF EXISTS ai_integration_state_pkey CASCADE;',
reverse_sql=migrations.RunSQL.noop,
),
# Remove id field
migrations.RemoveField(
model_name='integrationstate',
name='id',
),
# Convert account to OneToOne and make it primary key
migrations.AlterField(
model_name='integrationstate',
name='account',
field=models.OneToOneField(
help_text='Account that owns this integration state',
on_delete=django.db.models.deletion.CASCADE,
primary_key=True,
related_name='integration_state',
serialize=False,
to='igny8_core_auth.account'
),
),
# Make new fields non-nullable
migrations.AlterField(
model_name='integrationstate',
name='is_openai_enabled',
field=models.BooleanField(default=True, help_text='Whether OpenAI integration is enabled for this account'),
),
migrations.AlterField(
model_name='integrationstate',
name='is_runware_enabled',
field=models.BooleanField(default=True, help_text='Whether Runware integration is enabled for this account'),
),
migrations.AlterField(
model_name='integrationstate',
name='is_image_generation_enabled',
field=models.BooleanField(default=True, help_text='Whether Image Generation Service is enabled for this account'),
),
# Add new indexes
migrations.AddIndex(
model_name='integrationstate',
index=models.Index(fields=['is_openai_enabled'], name='ai_integrat_is_open_32213f_idx'),
),
migrations.AddIndex(
model_name='integrationstate',
index=models.Index(fields=['is_runware_enabled'], name='ai_integrat_is_runw_de35ad_idx'),
),
migrations.AddIndex(
model_name='integrationstate',
index=models.Index(fields=['is_image_generation_enabled'], name='ai_integrat_is_imag_0191f2_idx'),
),
]

View File

@@ -5,6 +5,61 @@ from django.db import models
from igny8_core.auth.models import AccountBaseModel
class IntegrationState(models.Model):
"""
Tracks whether AI integrations are enabled/disabled for each account.
Single record per account with separate fields for each integration type.
"""
account = models.OneToOneField(
'igny8_core_auth.Account',
on_delete=models.CASCADE,
related_name='integration_state',
help_text='Account that owns this integration state',
primary_key=True
)
# Enable/disable flags for each integration
is_openai_enabled = models.BooleanField(
default=True,
help_text='Whether OpenAI integration is enabled for this account'
)
is_runware_enabled = models.BooleanField(
default=True,
help_text='Whether Runware integration is enabled for this account'
)
is_image_generation_enabled = models.BooleanField(
default=True,
help_text='Whether Image Generation Service is enabled for this account'
)
created_at = models.DateTimeField(auto_now_add=True)
updated_at = models.DateTimeField(auto_now=True)
class Meta:
db_table = 'ai_integration_state'
verbose_name = 'Integration State'
verbose_name_plural = 'Integration States'
indexes = [
models.Index(fields=['is_openai_enabled']),
models.Index(fields=['is_runware_enabled']),
models.Index(fields=['is_image_generation_enabled']),
]
def __str__(self):
states = []
if self.is_openai_enabled:
states.append('OpenAI')
if self.is_runware_enabled:
states.append('Runware')
if self.is_image_generation_enabled:
states.append('Image Gen')
enabled_str = ', '.join(states) if states else 'None'
return f"{self.account.name} - Enabled: {enabled_str}"
class AITaskLog(AccountBaseModel):
"""
Unified logging table for all AI tasks.