Feature: Implement Daily AI Usage Limit with Cryptographic Challenge Verification #24
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,23 @@ | ||
| import type { Knex } from 'knex' | ||
|
|
||
| export async function up(knex: Knex): Promise<void> { | ||
| await knex.schema.alterTable('users', table => { | ||
| // Track daily AI usage count | ||
| table.integer('ai_usage_count').defaultTo(0).notNullable() | ||
| // Store the last reset date for the AI usage count | ||
| table.date('ai_usage_reset_date').nullable() | ||
| // Store the last generated challenge for AI usage | ||
| table.string('ai_challenge_uuid', 36).nullable() | ||
| // Store the timestamp when the challenge was generated | ||
| table.timestamp('ai_challenge_created_at').nullable() | ||
| }) | ||
| } | ||
|
|
||
| export async function down(knex: Knex): Promise<void> { | ||
| await knex.schema.alterTable('users', table => { | ||
| table.dropColumn('ai_usage_count') | ||
| table.dropColumn('ai_usage_reset_date') | ||
| table.dropColumn('ai_challenge_uuid') | ||
| table.dropColumn('ai_challenge_created_at') | ||
| }) | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,122 +1,267 @@ | ||
| /* eslint-disable @typescript-eslint/no-unsafe-argument */ | ||
| import express, { Request, Response } from 'express' | ||
| import { Knex } from 'knex' | ||
| import { AiChallengeVerifyDTO } from '../types/ai' | ||
| import { AiPromptRequest } from '../types/ai' | ||
| import { AiService } from '../utils/ai-service' | ||
| import { AiUsageService } from '../utils/ai-usage' | ||
| import { requireAuth } from '../utils/auth' | ||
|
|
||
| // Define a custom Request type that includes the address property | ||
| interface CustomRequest extends Request { | ||
| address?: string | ||
| } | ||
|
|
||
| // Get API key from environment variables | ||
| const OPENAI_API_KEY = process.env.OPENAI_API_KEY || '' | ||
|
|
||
| /** | ||
| * Creates an Express router for AI-related endpoints | ||
| * @param db - Knex database instance | ||
| * @param customAiService - Optional custom AI service for testing | ||
| * @param customAiUsageService - Optional custom AI usage service for testing | ||
| * @returns Express router with AI endpoints | ||
| */ | ||
| export function createAiRouter( | ||
| db: Knex, | ||
| customAiService: AiService | null = null, | ||
| customAiUsageService: AiUsageService | null = null, | ||
| ): express.Router { | ||
| const router = express.Router() | ||
|
|
||
| // Create AI service if not provided | ||
| const aiService = customAiService || (OPENAI_API_KEY ? new AiService({ apiKey: OPENAI_API_KEY }) : null) | ||
|
|
||
| // Create AI usage service if not provided | ||
| const aiUsageService = customAiUsageService || new AiUsageService(db) | ||
|
|
||
| /** | ||
| * Generate a challenge for AI usage | ||
| * GET /api/ai/challenge | ||
| */ | ||
| router.get('/challenge', requireAuth, async (req: CustomRequest, res: Response) => { | ||
| try { | ||
| const address = req.address | ||
|
|
||
| if (!address) { | ||
| return res.status(401).json({ | ||
| success: false, | ||
| error: 'Authentication required', | ||
| }) | ||
| } | ||
|
|
||
| // Generate a challenge | ||
| const challenge = await aiUsageService.generateChallenge(address) | ||
|
|
||
| return res.json({ | ||
| success: true, | ||
| data: challenge, | ||
| }) | ||
| } catch (error) { | ||
| console.error('Error generating AI challenge:', error instanceof Error ? error.message : 'Unknown error') | ||
| return res.status(500).json({ | ||
| success: false, | ||
| error: 'Failed to generate challenge', | ||
| }) | ||
| } | ||
| }) | ||
|
|
||
| /** | ||
| * Verify a challenge for AI usage | ||
| * POST /api/ai/verify-challenge | ||
| */ | ||
| router.post('/verify-challenge', requireAuth, async (req: CustomRequest, res: Response) => { | ||
| try { | ||
| const { address, challenge, signature } = req.body as AiChallengeVerifyDTO | ||
|
|
||
| // Verify the challenge | ||
| const result = await aiUsageService.verifyChallenge(address, challenge, signature) | ||
|
|
||
| return res.json({ | ||
| success: result.success, | ||
| data: { | ||
| remaining_attempts: result.remaining_attempts, | ||
| max_attempts: result.max_attempts, | ||
| }, | ||
| }) | ||
| } catch (error) { | ||
| console.error('Error verifying AI challenge:', error instanceof Error ? error.message : 'Unknown error') | ||
| return res.status(500).json({ | ||
| success: false, | ||
| error: 'Failed to verify challenge', | ||
| }) | ||
|
Comment on lines
+88
to
92
There was a problem hiding this comment. Similar to the previous comment, the error message here is also generic. Consider adding more context to the error message to help the user understand why the challenge verification failed. For example, you could check if the challenge is valid before attempting to verify it and return a more informative error message if the challenge is invalid. console.error('Error verifying AI challenge:', error instanceof Error ? error.message : 'Unknown error')
const errorMessage = error instanceof Error ? error.message : 'Unknown error'
return res.status(500).json({
success: false,
error: `Failed to verify challenge: ${errorMessage}`, // More specific error message
}) |
||
| } | ||
| }) | ||
|
|
||
| /** | ||
| * Get remaining AI requests for a user | ||
| * GET /api/ai/remaining-requests | ||
| */ | ||
| router.get('/remaining-requests', requireAuth, async (req: CustomRequest, res: Response) => { | ||
| try { | ||
| const address = req.address | ||
|
|
||
| if (!address) { | ||
| return res.status(401).json({ | ||
| success: false, | ||
| error: 'Authentication required', | ||
| }) | ||
| } | ||
|
|
||
| // Get remaining requests | ||
| const remaining = await aiUsageService.getRemainingRequests(address) | ||
|
|
||
| return res.json({ | ||
| success: true, | ||
| data: remaining, | ||
| }) | ||
| } catch (error) { | ||
| console.error('Error getting remaining AI requests:', error instanceof Error ? error.message : 'Unknown error') | ||
| return res.status(500).json({ | ||
| success: false, | ||
| error: 'Failed to get remaining requests', | ||
| }) | ||
|
Comment on lines
+119
to
+123
There was a problem hiding this comment. The error message here is also generic. Consider adding more context to the error message to help the user understand why getting the remaining AI requests failed. For example, you could check if the user exists before attempting to get the remaining requests and return a more informative error message if the user is not found. console.error('Error getting remaining AI requests:', error instanceof Error ? error.message : 'Unknown error')
const errorMessage = error instanceof Error ? error.message : 'Unknown error'
return res.status(500).json({
success: false,
error: `Failed to get remaining requests: ${errorMessage}`, // More specific error message
}) |
||
| } | ||
| }) | ||
|
|
||
| /** | ||
| * Process a prompt with a template | ||
| * POST /api/ai/process-prompt | ||
| */ | ||
| router.post( | ||
| '/process-prompt', | ||
| async (req: Request, res: Response, next: () => void) => { | ||
| // Check if AI service is available first | ||
| if (!aiService) { | ||
| return res.status(503).json({ | ||
| success: false, | ||
| error: 'AI service is not available', | ||
| }) | ||
| } | ||
|
|
||
| // If AI service is available, proceed with authentication | ||
| requireAuth(req as CustomRequest, res, next) | ||
| }, | ||
| async (req: CustomRequest, res: Response) => { | ||
| try { | ||
| const { prompt, templateId, challenge, signature } = req.body as AiPromptRequest | ||
| const address = req.address | ||
|
|
||
| if (!address) { | ||
| return res.status(401).json({ | ||
| success: false, | ||
| error: 'Authentication required', | ||
| }) | ||
| } | ||
|
|
||
| // Validate required fields | ||
| if (!prompt || !templateId || !challenge || !signature) { | ||
| return res.status(400).json({ | ||
| success: false, | ||
| error: 'Missing required fields', | ||
| }) | ||
| } | ||
|
|
||
| // Verify the challenge | ||
| try { | ||
| const verificationResult = await aiUsageService.verifyChallenge(address, challenge, signature) | ||
|
|
||
| // If verification failed due to usage limits | ||
| if (!verificationResult.success) { | ||
| return res.status(403).json({ | ||
| success: false, | ||
| error: 'Daily AI usage limit reached', | ||
| data: { | ||
| remaining_attempts: verificationResult.remaining_attempts, | ||
| max_attempts: verificationResult.max_attempts, | ||
| }, | ||
| }) | ||
| } | ||
| } catch (error) { | ||
| const errorMessage = error instanceof Error ? error.message : 'Unknown error' | ||
| console.error('Error verifying challenge:', errorMessage) | ||
| return res.status(403).json({ | ||
| success: false, | ||
| error: 'Challenge verification failed', | ||
| }) | ||
|
Comment on lines
+183
to
+186
There was a problem hiding this comment. Returning a 403 status code here might not be the most appropriate choice. A 403 status code typically indicates that the server understands the request but refuses to authorize it. In this case, the user is authorized, but they have exceeded their daily usage limit. A more appropriate status code might be 429 (Too Many Requests). return res.status(429).json({
success: false,
error: 'Daily AI usage limit reached',
data: {
remaining_attempts: verificationResult.remaining_attempts,
max_attempts: verificationResult.max_attempts,
},
}) |
||
| } | ||
|
|
||
| // Get the template | ||
| const template = await db('templates').where({ id: templateId }).first() | ||
| if (!template) { | ||
| return res.status(404).json({ | ||
| success: false, | ||
| error: 'Template not found', | ||
| }) | ||
| } | ||
|
|
||
| // Check if template has JSON data | ||
| if (template.json_data === '') { | ||
| console.error('Database error: Empty JSON data') | ||
| return res.status(500).json({ | ||
| success: false, | ||
| error: 'Internal server error', | ||
| }) | ||
| } | ||
|
|
||
| if (!template.json_data || template.json_data === '{}') { | ||
| return res.status(400).json({ | ||
| success: false, | ||
| error: 'Template JSON data is missing', | ||
| }) | ||
| } | ||
|
|
||
| // Parse the template JSON data | ||
| let schema: Record<string, unknown> | ||
| let systemPrompt: string | ||
| try { | ||
| const jsonData = JSON.parse(template.json_data) | ||
| if (typeof jsonData !== 'object' || !jsonData || !jsonData.schema || typeof jsonData.schema !== 'object') { | ||
| return res.status(400).json({ | ||
| success: false, | ||
| error: 'Invalid template JSON data: missing or invalid schema', | ||
| }) | ||
| } | ||
| schema = jsonData.schema as Record<string, unknown> | ||
| systemPrompt = | ||
| typeof jsonData.systemPrompt === 'string' ? jsonData.systemPrompt : 'You are a helpful assistant.' | ||
| } catch (error) { | ||
| // Invalid JSON format | ||
| console.error('Error parsing template JSON data:', error instanceof Error ? error.message : 'Unknown error') | ||
| return res.status(400).json({ | ||
| success: false, | ||
| error: 'Invalid template JSON data', | ||
| }) | ||
| } | ||
|
|
||
| // Check if AI service is available | ||
| if (!aiService) { | ||
| return res.status(503).json({ | ||
| success: false, | ||
| error: 'AI service is not available', | ||
| }) | ||
| } | ||
|
|
||
| // Process the prompt with the template | ||
| const result = await aiService.processTemplatePrompt(prompt, schema, systemPrompt) | ||
|
|
||
| // Return the response | ||
| return res.json({ | ||
| success: true, | ||
| data: { | ||
| result: result.parsedData || {}, | ||
| requiredValidation: !result.isValid, | ||
| }, | ||
| }) | ||
| } catch (error) { | ||
| console.error('Error processing prompt:', error instanceof Error ? error.message : 'Unknown error') | ||
| return res.status(500).json({ | ||
| success: false, | ||
| error: 'Internal server error', | ||
| }) | ||
| } | ||
| }, | ||
| ) | ||
|
|
||
| return router | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The error message here is a bit generic. Consider providing more specific feedback to the user about why the challenge generation failed. For example, you could check if the user exists before attempting to generate the challenge and return a more informative error message if the user is not found.