Skip to content

Commit 1f89774

Browse files
yash-pouranikyash-pouranik
committed
feat: implement AI schema limits and Schema Canvas Viewer
1 parent 4b52148 commit 1f89774

9 files changed

Lines changed: 144 additions & 66 deletions

File tree

‎apps/dashboard-api/src/controllers/ai.controller.js‎

Lines changed: 5 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -163,34 +163,8 @@ const queryBuilder = async (req, res, next) => {
163163

164164
const ccSessionKey = (projectId, developerId) =>
165165
`ai:cc:session:${projectId}:${developerId}`;
166-
const ccIterationsKey = (projectId, developerId) =>
167-
`ai:cc:iterations:${projectId}:${developerId}`;
168-
169-
const LUA_RESERVE_ITERATION = `
170-
local sessionExists = redis.call('EXISTS', KEYS[2])
171-
local current = redis.call('GET', KEYS[1])
172-
local count = tonumber(current) or 0
173-
174-
if sessionExists == 0 and count > 0 then
175-
count = 0
176-
redis.call('SET', KEYS[1], '0')
177-
end
178-
179-
if count >= tonumber(ARGV[1]) then
180-
return {0, count}
181-
end
182-
183-
local nextVal = redis.call('INCR', KEYS[1])
184-
redis.call('EXPIRE', KEYS[1], ARGV[2])
185-
if sessionExists == 1 then
186-
redis.call('EXPIRE', KEYS[2], ARGV[2])
187-
end
188-
return {1, nextVal}
189-
`;
190166

191167
const collectionCreator = async (req, res, next) => {
192-
let reservedIteration = false;
193-
let iterKey = null;
194168
let hasByok = false;
195169

196170
try {
@@ -221,7 +195,6 @@ const collectionCreator = async (req, res, next) => {
221195

222196
// 2. Load Redis session
223197
const sessionKey = ccSessionKey(projectId, req.user._id);
224-
iterKey = ccIterationsKey(projectId, req.user._id);
225198
const sessionStr = await redis.get(sessionKey);
226199

227200
let session = { messages: [], hasByok: false };
@@ -251,24 +224,11 @@ const collectionCreator = async (req, res, next) => {
251224
const encryptedByok = resolvedKey ? { groqKey: encryptForTransit(resolvedKey) } : null;
252225
hasByok = !!encryptedByok;
253226

254-
// 4. Atomic iteration reservation before calling Python service
255-
const PLATFORM_ITERATION_LIMIT = 3;
256-
let iterations = 0;
257-
258-
if (!hasByok) {
259-
const [allowed, newCount] = await redis.eval(
260-
LUA_RESERVE_ITERATION,
261-
2,
262-
iterKey,
263-
sessionKey,
264-
PLATFORM_ITERATION_LIMIT,
265-
7200
266-
);
267-
if (allowed === 0) {
268-
throw new AppError(403, "AI iteration limit reached (3/3 on platform key). Add your Groq API key in Settings for unlimited usage.");
269-
}
270-
reservedIteration = true;
271-
iterations = newCount;
227+
const PLATFORM_ITERATION_LIMIT = 20;
228+
let iterations = Math.floor(session.messages.length / 2);
229+
230+
if (!hasByok && iterations >= PLATFORM_ITERATION_LIMIT) {
231+
throw new AppError(403, `Session iteration limit reached (${PLATFORM_ITERATION_LIMIT}/${PLATFORM_ITERATION_LIMIT}). Add your Groq API key in Settings for unlimited usage.`);
272232
}
273233

274234
// 5. Update session
@@ -345,14 +305,6 @@ const collectionCreator = async (req, res, next) => {
345305
}, "Agent responded successfully").send(res);
346306

347307
} catch (error) {
348-
// Rollback reserved iteration on error if platform key was used
349-
if (!hasByok && reservedIteration && iterKey) {
350-
try {
351-
await redis.decr(iterKey);
352-
} catch (revertErr) {
353-
console.warn("[CollectionCreator] Failed to rollback reserved iteration counter", revertErr);
354-
}
355-
}
356308
if (error instanceof AppError) {
357309
return next(error);
358310
}

‎apps/landing/fix_circles.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
import re
2+
3+
file_path = r'e:\Majors\WEB\urBackend\apps\landing\src\styles\landing.css'
4+
5+
with open(file_path, 'r') as f:
6+
content = f.read()
7+
8+
# CSS selectors that should be strictly circular
9+
circular_classes = [
10+
r'\.hero-glow',
11+
r'\.byom-glow',
12+
r'\.pulsing-dot',
13+
r'\.hf-dot-moving',
14+
r'\.flow-dot',
15+
r'\.flow-spinner',
16+
r'\.dot',
17+
r'\.orbit-center',
18+
r'\.orbit-sun-logo',
19+
r'\.orbit-sun-glow',
20+
r'\.orbit-sun-ring-dec',
21+
r'\.orbit-ring',
22+
r'\.orbit-dust',
23+
r'\.byom-particle',
24+
r'\.timeline-node-dot',
25+
r'\.loader-dot',
26+
r'\.mini-circle',
27+
r'\.mini-dot',
28+
r'\.lh-engine'
29+
]
30+
31+
for cls in circular_classes:
32+
pattern = re.compile(rf'({cls}(?:,\s*[^{{]+)*\s*{{)([^}}]+)', re.MULTILINE)
33+
34+
def replacer(match):
35+
header = match.group(1)
36+
body = match.group(2)
37+
if 'border-radius: 50%;' not in body:
38+
return header + '\n border-radius: 50%;' + body
39+
return match.group(0)
40+
41+
content = pattern.sub(replacer, content)
42+
43+
with open(file_path, 'w') as f:
44+
f.write(content)
45+
46+
print("Done restoring border-radius to circular elements.")

‎apps/landing/fix_missed_circles.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import re
2+
3+
file_path = r'e:\Majors\WEB\urBackend\apps\landing\src\styles\landing.css'
4+
5+
with open(file_path, 'r') as f:
6+
content = f.read()
7+
8+
# CSS selectors that should be strictly circular
9+
circular_classes = [
10+
r'\.orbit-inner-wrapper',
11+
r'\.orbit-mid-wrapper',
12+
r'\.orbit-outer-wrapper',
13+
r'\.orbit-sun-ring-1',
14+
r'\.orbit-sun-ring-2',
15+
r'\.orbit-connect-line::after',
16+
r'\.byom-pipe-particle'
17+
]
18+
19+
for cls in circular_classes:
20+
# Match the CSS block exactly.
21+
pattern = re.compile(rf'^({cls}(?:,\s*[^{{]+)*\s*{{)([^}}]+)', re.MULTILINE)
22+
23+
def replacer(match):
24+
header = match.group(1)
25+
body = match.group(2)
26+
if 'border-radius: 50%;' not in body:
27+
return header + '\n border-radius: 50%;' + body
28+
return match.group(0)
29+
30+
content = pattern.sub(replacer, content)
31+
32+
with open(file_path, 'w') as f:
33+
f.write(content)
34+
35+
print("Done restoring border-radius to missed elements.")

‎apps/python-service/routers/ai.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,8 @@ async def collection_creator(request: CollectionCreatorRequest, req: Request):
187187
developer_id=request.developer_id,
188188
plan=request.plan,
189189
encrypted_byok=request.encrypted_byok.model_dump() if request.encrypted_byok else None,
190+
feature="collection-creator",
191+
increment=(len(request.messages) == 1),
190192
)
191193

192194
structured_llm = llm.with_structured_output(CollectionCreatorResponse)

‎apps/python-service/services/byok.py‎

Lines changed: 35 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@
1515
FREE_TIER_LIMIT = 5
1616
PRO_TIER_LIMIT = 20
1717

18+
CC_FREE_TIER_LIMIT = 1
19+
CC_PRO_TIER_LIMIT = 5
20+
1821
LUA_RATE_LIMIT = """
1922
local count = redis.call('INCR', KEYS[1])
2023
local ttl = redis.call('TTL', KEYS[1])
@@ -25,11 +28,19 @@
2528
return {count, ttl}
2629
"""
2730

31+
LUA_CHECK_LIMIT = """
32+
local count = tonumber(redis.call('GET', KEYS[1]) or '0')
33+
local ttl = redis.call('TTL', KEYS[1])
34+
return {count, ttl}
35+
"""
36+
2837

2938
async def resolve_ai_client(
3039
developer_id: str | None,
3140
plan: str | None,
3241
encrypted_byok: dict | None,
42+
feature: str = "query-builder",
43+
increment: bool = True,
3344
) -> ChatGroq:
3445
"""Resolve the appropriate AI client based on BYOK and plan.
3546
@@ -42,14 +53,16 @@ async def resolve_ai_client(
4253
developer_id: The developer's MongoDB ObjectId string.
4354
plan: The developer's plan ('free' or 'pro').
4455
encrypted_byok: Transit-encrypted BYOK payload from Node, or None.
56+
feature: The feature being used ("query-builder" or "collection-creator").
57+
increment: Whether to increment the usage count.
4558
4659
Returns:
4760
An initialized ChatGroq client.
4861
4962
Raises:
5063
HTTPException: 401 if BYOK key is invalid, 403 if rate limited, 500 if no platform key.
5164
"""
52-
logger.info("Resolving AI client | developer_id=%s, plan=%s, has_byok=%s", developer_id, plan, bool(encrypted_byok and encrypted_byok.get("groqKey")))
65+
logger.info("Resolving AI client | feature=%s, dev_id=%s, plan=%s, has_byok=%s, increment=%s", feature, developer_id, plan, bool(encrypted_byok and encrypted_byok.get("groqKey")), increment)
5366

5467
# ── 1. Try BYOK ──
5568
if encrypted_byok and encrypted_byok.get("groqKey"):
@@ -83,18 +96,33 @@ async def resolve_ai_client(
8396
detail="developer_id is required for platform key rate limiting"
8497
)
8598

86-
limit = PRO_TIER_LIMIT if plan == "pro" else FREE_TIER_LIMIT
99+
if feature == "collection-creator":
100+
limit = CC_PRO_TIER_LIMIT if plan == "pro" else CC_FREE_TIER_LIMIT
101+
key_prefix = "ai:cc:count"
102+
else:
103+
limit = PRO_TIER_LIMIT if plan == "pro" else FREE_TIER_LIMIT
104+
key_prefix = "ai:gen:count"
105+
87106
month = datetime.datetime.now(datetime.UTC).strftime("%Y-%m")
88-
key = f"ai:gen:count:{developer_id}:{month}"
107+
key = f"{key_prefix}:{developer_id}:{month}"
108+
109+
logger.info("🏢 Using Platform Groq key. Checking Redis rate limit at '%s' (limit=%d/month, increment=%s)...", key, limit, increment)
89110

90-
logger.info("🏢 Using Platform Groq key. Checking Redis rate limit at '%s' (limit=%d/month)...", key, limit)
111+
if increment:
112+
count, ttl = await redis_client.eval(LUA_RATE_LIMIT, 1, key, 32 * 86400)
113+
else:
114+
count, ttl = await redis_client.eval(LUA_CHECK_LIMIT, 1, key)
115+
# If we are just checking, we still enforce the limit but don't increment.
116+
# Actually, if count >= limit, we should block.
117+
# Wait, if count == limit and we check, it's allowed.
118+
# If count > limit, it's blocked.
119+
# So we just use count.
91120

92-
# Atomic script execution: INCR + EXPIRE (if new key)
93-
count, ttl = await redis_client.eval(LUA_RATE_LIMIT, 1, key, 32 * 86400)
94121
logger.info("📊 Monthly usage for developer_id=%s: %d/%d (TTL: %ds remaining in window)", developer_id, count, limit, ttl)
95122

96123
if count > limit:
97124
tier_name = "Pro" if plan == "pro" else "Free"
125+
feature_name = "Schema Designs" if feature == "collection-creator" else "AI usages"
98126
upgrade_msg = (
99127
"Add your own Groq key in Settings to get unlimited usage."
100128
if plan == "pro"
@@ -103,7 +131,7 @@ async def resolve_ai_client(
103131
logger.warning("🚫 Rate limit exceeded for developer_id=%s: %d > %d allowed in %s tier", developer_id, count, limit, tier_name)
104132
raise HTTPException(
105133
status_code=403,
106-
detail=f"{tier_name} tier AI limit reached ({limit}/month). {upgrade_msg}",
134+
detail=f"{tier_name} tier limit reached ({limit} {feature_name}/month). {upgrade_msg}",
107135
)
108136

109137
# ── 4. Return platform client ──

‎apps/web-dashboard/src/components/Database/DatabaseHeader.jsx‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import React from 'react';
22
import {
33
Menu, List as ListIcon, Table as TableIcon, Code,
4-
Filter, RefreshCw, Shield, Plus, Download
4+
Filter, RefreshCw, Shield, Plus, Download, Workflow
55
} from 'lucide-react';
66
import AiQueryBar from './AiQueryBar';
77

@@ -67,7 +67,8 @@ const DatabaseHeader = ({
6767
{[
6868
{ id: 'table', icon: TableIcon, title: 'Table' },
6969
{ id: 'list', icon: ListIcon, title: 'List' },
70-
{ id: 'json', icon: Code, title: 'JSON' }
70+
{ id: 'json', icon: Code, title: 'JSON' },
71+
{ id: 'canvas', icon: Workflow, title: 'Schema Graph' }
7172
].map(mode => (
7273
<button
7374
key={mode.id}

‎apps/web-dashboard/src/components/SchemaCanvasViewer.jsx‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,8 @@ export default function SchemaCanvasViewer({
3737
isInserting,
3838
onInsertAll,
3939
onResetChat,
40-
onNavigateToDb
40+
onNavigateToDb,
41+
readOnly = false
4142
}) {
4243
const [viewMode, setViewMode] = useState('canvas'); // 'canvas' | 'list' | 'json'
4344
const [isConfirmOpen, setIsConfirmOpen] = useState(false);
@@ -209,7 +210,7 @@ export default function SchemaCanvasViewer({
209210

210211
{/* Right Action Controls */}
211212
<div className="flex items-center gap-2.5">
212-
{messages.length > 1 && (
213+
{!readOnly && messages?.length > 1 && (
213214
<Tooltip>
214215
<TooltipTrigger asChild>
215216
<Button
@@ -238,7 +239,7 @@ export default function SchemaCanvasViewer({
238239
</Button>
239240
)}
240241

241-
{schema && schema.length > 0 && (
242+
{schema && schema.length > 0 && !readOnly && (
242243
<Button
243244
variant="default"
244245
size="sm"

‎apps/web-dashboard/src/pages/Database.jsx‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import { PUBLIC_API_URL } from '../config';
1515
import DatabaseHeader from "../components/Database/DatabaseHeader";
1616
import DatabaseFilter from "../components/Database/DatabaseFilter";
1717
import Pagination from "../components/Database/Pagination";
18+
import SchemaCanvasViewer from "../components/SchemaCanvasViewer";
1819

1920
export default function Database() {
2021
const { projectId } = useParams();
@@ -377,6 +378,17 @@ export default function Database() {
377378
recoveringIds={recoveringIds}
378379
isViewer={isViewer}
379380
/>
381+
) : viewMode === "canvas" ? (
382+
<div style={{ position: 'absolute', top: 0, left: 0, right: 0, bottom: 0, background: 'var(--color-bg-main)' }}>
383+
<SchemaCanvasViewer
384+
schema={collections.filter(c => c.name !== 'users').map(c => ({
385+
collection: c.name,
386+
fields: (c.model || []).map(f => ({ ...f, name: f.key }))
387+
}))}
388+
messages={[]}
389+
readOnly={true}
390+
/>
391+
</div>
380392
) : (
381393
<div style={{ height: '100%', overflow: 'auto', padding: '1rem', background: 'var(--color-bg-input)', color: 'var(--color-primary)', fontFamily: 'monospace', fontSize: '0.75rem' }}>
382394
<pre>{JSON.stringify(data, null, 2)}</pre>

‎apps/web-dashboard/src/utils/api.js‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,9 @@ api.interceptors.response.use(
100100

101101
const isVerificationGate = VERIFICATION_KEYWORDS.some((kw) => message.toLowerCase().includes(kw));
102102
const isPlanError = UPGRADE_KEYWORDS.some((kw) => message.toLowerCase().includes(kw));
103+
const requiresByok = message.toLowerCase().includes('add your own groq key') || message.toLowerCase().includes('settings for unlimited');
103104

104-
if (isPlanError && !isVerificationGate) {
105+
if (isPlanError && !isVerificationGate && !requiresByok) {
105106
toast.error("Plan limit reached. Please upgrade to continue.");
106107
if (window.location.href !== 'https://urbackend.bitbros.in/pricing') {
107108
window.location.assign('https://urbackend.bitbros.in/pricing');

0 commit comments

Comments
 (0)