Merge branch 'main' of https://git.igny8.com/salman/igny8
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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')},
|
||||
},
|
||||
),
|
||||
]
|
||||
@@ -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'),
|
||||
),
|
||||
]
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user