mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-01-02 22:32:27 +08:00
Compare commits
5 Commits
refactor/d
...
v0.4.0
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c42efdc702 | ||
|
|
dd027f1856 | ||
|
|
869391a029 | ||
|
|
8b9336466f | ||
|
|
ee514efa9e |
@@ -85,9 +85,11 @@ Here are some example prompts and their generated diagrams:
|
||||
|
||||
- **LLM-Powered Diagram Creation**: Leverage Large Language Models to create and manipulate draw.io diagrams directly through natural language commands
|
||||
- **Image-Based Diagram Replication**: Upload existing diagrams or images and have the AI replicate and enhance them automatically
|
||||
- **PDF & Text File Upload**: Upload PDF documents and text files to extract content and generate diagrams from existing documents
|
||||
- **AI Reasoning Display**: View the AI's thinking process for supported models (OpenAI o1/o3, Gemini, Claude, etc.)
|
||||
- **Diagram History**: Comprehensive version control that tracks all changes, allowing you to view and restore previous versions of your diagrams before the AI editing.
|
||||
- **Interactive Chat Interface**: Communicate with AI to refine your diagrams in real-time
|
||||
- **AWS Architecture Diagram Support**: Specialized support for generating AWS architecture diagrams
|
||||
- **Cloud Architecture Diagram Support**: Specialized support for generating cloud architecture diagrams (AWS, GCP, Azure)
|
||||
- **Animated Connectors**: Create dynamic and animated connectors between diagram elements for better visualization
|
||||
|
||||
## Getting Started
|
||||
|
||||
@@ -17,39 +17,21 @@ import { FaGithub } from "react-icons/fa"
|
||||
import { Toaster, toast } from "sonner"
|
||||
import { ButtonWithTooltip } from "@/components/button-with-tooltip"
|
||||
import { ChatInput } from "@/components/chat-input"
|
||||
import { QuotaLimitToast } from "@/components/quota-limit-toast"
|
||||
import {
|
||||
SettingsDialog,
|
||||
STORAGE_ACCESS_CODE_KEY,
|
||||
STORAGE_AI_API_KEY_KEY,
|
||||
STORAGE_AI_BASE_URL_KEY,
|
||||
STORAGE_AI_MODEL_KEY,
|
||||
STORAGE_AI_PROVIDER_KEY,
|
||||
} from "@/components/settings-dialog"
|
||||
import { SettingsDialog } from "@/components/settings-dialog"
|
||||
import { useDiagram } from "@/contexts/diagram-context"
|
||||
import { getAIConfig } from "@/lib/ai-config"
|
||||
import { findCachedResponse } from "@/lib/cached-responses"
|
||||
import { isPdfFile, isTextFile } from "@/lib/pdf-utils"
|
||||
import { type FileData, useFileProcessor } from "@/lib/use-file-processor"
|
||||
import { useQuotaManager } from "@/lib/use-quota-manager"
|
||||
import { formatXML, wrapWithMxFile } from "@/lib/utils"
|
||||
import { ChatMessageDisplay } from "./chat-message-display"
|
||||
|
||||
// localStorage keys for persistence
|
||||
const STORAGE_MESSAGES_KEY = "next-ai-draw-io-messages"
|
||||
const STORAGE_XML_SNAPSHOTS_KEY = "next-ai-draw-io-xml-snapshots"
|
||||
const STORAGE_SESSION_ID_KEY = "next-ai-draw-io-session-id"
|
||||
export const STORAGE_DIAGRAM_XML_KEY = "next-ai-draw-io-diagram-xml"
|
||||
const STORAGE_REQUEST_COUNT_KEY = "next-ai-draw-io-request-count"
|
||||
const STORAGE_REQUEST_DATE_KEY = "next-ai-draw-io-request-date"
|
||||
const STORAGE_TOKEN_COUNT_KEY = "next-ai-draw-io-token-count"
|
||||
const STORAGE_TOKEN_DATE_KEY = "next-ai-draw-io-token-date"
|
||||
const STORAGE_TPM_COUNT_KEY = "next-ai-draw-io-tpm-count"
|
||||
const STORAGE_TPM_MINUTE_KEY = "next-ai-draw-io-tpm-minute"
|
||||
|
||||
import { useDiagram } from "@/contexts/diagram-context"
|
||||
import { findCachedResponse } from "@/lib/cached-responses"
|
||||
import {
|
||||
extractPdfText,
|
||||
extractTextFileContent,
|
||||
isPdfFile,
|
||||
isTextFile,
|
||||
MAX_EXTRACTED_CHARS,
|
||||
} from "@/lib/pdf-utils"
|
||||
import { formatXML, wrapWithMxFile } from "@/lib/utils"
|
||||
import { ChatMessageDisplay } from "./chat-message-display"
|
||||
|
||||
// Type for message parts (tool calls and their states)
|
||||
interface MessagePart {
|
||||
@@ -186,11 +168,9 @@ export default function ChatPanel({
|
||||
])
|
||||
}
|
||||
|
||||
const [files, setFiles] = useState<File[]>([])
|
||||
// Store extracted PDF text with extraction status
|
||||
const [pdfData, setPdfData] = useState<
|
||||
Map<File, { text: string; charCount: number; isExtracting: boolean }>
|
||||
>(new Map())
|
||||
// File processing using extracted hook
|
||||
const { files, pdfData, handleFileChange, setFiles } = useFileProcessor()
|
||||
|
||||
const [showHistory, setShowHistory] = useState(false)
|
||||
const [showSettingsDialog, setShowSettingsDialog] = useState(false)
|
||||
const [, setAccessCodeRequired] = useState(false)
|
||||
@@ -212,200 +192,12 @@ export default function ChatPanel({
|
||||
.catch(() => setAccessCodeRequired(false))
|
||||
}, [])
|
||||
|
||||
// Helper to check daily request limit
|
||||
// Check if user has their own API key configured (bypass limits)
|
||||
const hasOwnApiKey = useCallback((): boolean => {
|
||||
const provider = localStorage.getItem(STORAGE_AI_PROVIDER_KEY)
|
||||
const apiKey = localStorage.getItem(STORAGE_AI_API_KEY_KEY)
|
||||
return !!(provider && apiKey)
|
||||
}, [])
|
||||
|
||||
const checkDailyLimit = useCallback((): {
|
||||
allowed: boolean
|
||||
remaining: number
|
||||
used: number
|
||||
} => {
|
||||
// Skip limit if user has their own API key
|
||||
if (hasOwnApiKey()) return { allowed: true, remaining: -1, used: 0 }
|
||||
if (dailyRequestLimit <= 0)
|
||||
return { allowed: true, remaining: -1, used: 0 }
|
||||
|
||||
const today = new Date().toDateString()
|
||||
const storedDate = localStorage.getItem(STORAGE_REQUEST_DATE_KEY)
|
||||
let count = parseInt(
|
||||
localStorage.getItem(STORAGE_REQUEST_COUNT_KEY) || "0",
|
||||
10,
|
||||
)
|
||||
|
||||
if (storedDate !== today) {
|
||||
count = 0
|
||||
localStorage.setItem(STORAGE_REQUEST_DATE_KEY, today)
|
||||
localStorage.setItem(STORAGE_REQUEST_COUNT_KEY, "0")
|
||||
}
|
||||
|
||||
return {
|
||||
allowed: count < dailyRequestLimit,
|
||||
remaining: dailyRequestLimit - count,
|
||||
used: count,
|
||||
}
|
||||
}, [dailyRequestLimit, hasOwnApiKey])
|
||||
|
||||
// Helper to increment request count
|
||||
const incrementRequestCount = useCallback((): void => {
|
||||
const count = parseInt(
|
||||
localStorage.getItem(STORAGE_REQUEST_COUNT_KEY) || "0",
|
||||
10,
|
||||
)
|
||||
localStorage.setItem(STORAGE_REQUEST_COUNT_KEY, String(count + 1))
|
||||
}, [])
|
||||
|
||||
// Helper to show quota limit toast (request-based)
|
||||
const showQuotaLimitToast = useCallback(() => {
|
||||
toast.custom(
|
||||
(t) => (
|
||||
<QuotaLimitToast
|
||||
used={dailyRequestLimit}
|
||||
limit={dailyRequestLimit}
|
||||
onDismiss={() => toast.dismiss(t)}
|
||||
/>
|
||||
),
|
||||
{ duration: 15000 },
|
||||
)
|
||||
}, [dailyRequestLimit, hasOwnApiKey])
|
||||
|
||||
// Helper to check daily token limit (checks if already over limit)
|
||||
const checkTokenLimit = useCallback((): {
|
||||
allowed: boolean
|
||||
remaining: number
|
||||
used: number
|
||||
} => {
|
||||
// Skip limit if user has their own API key
|
||||
if (hasOwnApiKey()) return { allowed: true, remaining: -1, used: 0 }
|
||||
if (dailyTokenLimit <= 0)
|
||||
return { allowed: true, remaining: -1, used: 0 }
|
||||
|
||||
const today = new Date().toDateString()
|
||||
const storedDate = localStorage.getItem(STORAGE_TOKEN_DATE_KEY)
|
||||
let count = parseInt(
|
||||
localStorage.getItem(STORAGE_TOKEN_COUNT_KEY) || "0",
|
||||
10,
|
||||
)
|
||||
|
||||
// Guard against NaN (e.g., if "NaN" was stored)
|
||||
if (Number.isNaN(count)) count = 0
|
||||
|
||||
if (storedDate !== today) {
|
||||
count = 0
|
||||
localStorage.setItem(STORAGE_TOKEN_DATE_KEY, today)
|
||||
localStorage.setItem(STORAGE_TOKEN_COUNT_KEY, "0")
|
||||
}
|
||||
|
||||
return {
|
||||
allowed: count < dailyTokenLimit,
|
||||
remaining: dailyTokenLimit - count,
|
||||
used: count,
|
||||
}
|
||||
}, [dailyTokenLimit, hasOwnApiKey])
|
||||
|
||||
// Helper to increment token count
|
||||
const incrementTokenCount = useCallback((tokens: number): void => {
|
||||
// Guard against NaN tokens
|
||||
if (!Number.isFinite(tokens) || tokens <= 0) return
|
||||
|
||||
let count = parseInt(
|
||||
localStorage.getItem(STORAGE_TOKEN_COUNT_KEY) || "0",
|
||||
10,
|
||||
)
|
||||
// Guard against NaN count
|
||||
if (Number.isNaN(count)) count = 0
|
||||
|
||||
localStorage.setItem(STORAGE_TOKEN_COUNT_KEY, String(count + tokens))
|
||||
}, [])
|
||||
|
||||
// Helper to show token limit toast
|
||||
const showTokenLimitToast = useCallback(
|
||||
(used: number) => {
|
||||
toast.custom(
|
||||
(t) => (
|
||||
<QuotaLimitToast
|
||||
type="token"
|
||||
used={used}
|
||||
limit={dailyTokenLimit}
|
||||
onDismiss={() => toast.dismiss(t)}
|
||||
/>
|
||||
),
|
||||
{ duration: 15000 },
|
||||
)
|
||||
},
|
||||
[dailyTokenLimit],
|
||||
)
|
||||
|
||||
// Helper to check TPM (tokens per minute) limit
|
||||
// Note: This only READS, doesn't write. incrementTPMCount handles writes.
|
||||
const checkTPMLimit = useCallback((): {
|
||||
allowed: boolean
|
||||
remaining: number
|
||||
used: number
|
||||
} => {
|
||||
// Skip limit if user has their own API key
|
||||
if (hasOwnApiKey()) return { allowed: true, remaining: -1, used: 0 }
|
||||
if (tpmLimit <= 0) return { allowed: true, remaining: -1, used: 0 }
|
||||
|
||||
const currentMinute = Math.floor(Date.now() / 60000).toString()
|
||||
const storedMinute = localStorage.getItem(STORAGE_TPM_MINUTE_KEY)
|
||||
let count = parseInt(
|
||||
localStorage.getItem(STORAGE_TPM_COUNT_KEY) || "0",
|
||||
10,
|
||||
)
|
||||
|
||||
// Guard against NaN
|
||||
if (Number.isNaN(count)) count = 0
|
||||
|
||||
// If we're in a new minute, treat count as 0 (will be reset on next increment)
|
||||
if (storedMinute !== currentMinute) {
|
||||
count = 0
|
||||
}
|
||||
|
||||
return {
|
||||
allowed: count < tpmLimit,
|
||||
remaining: tpmLimit - count,
|
||||
used: count,
|
||||
}
|
||||
}, [tpmLimit, hasOwnApiKey])
|
||||
|
||||
// Helper to increment TPM count
|
||||
const incrementTPMCount = useCallback((tokens: number): void => {
|
||||
// Guard against NaN tokens
|
||||
if (!Number.isFinite(tokens) || tokens <= 0) return
|
||||
|
||||
const currentMinute = Math.floor(Date.now() / 60000).toString()
|
||||
const storedMinute = localStorage.getItem(STORAGE_TPM_MINUTE_KEY)
|
||||
let count = parseInt(
|
||||
localStorage.getItem(STORAGE_TPM_COUNT_KEY) || "0",
|
||||
10,
|
||||
)
|
||||
|
||||
// Guard against NaN
|
||||
if (Number.isNaN(count)) count = 0
|
||||
|
||||
// Reset if we're in a new minute
|
||||
if (storedMinute !== currentMinute) {
|
||||
count = 0
|
||||
localStorage.setItem(STORAGE_TPM_MINUTE_KEY, currentMinute)
|
||||
}
|
||||
|
||||
localStorage.setItem(STORAGE_TPM_COUNT_KEY, String(count + tokens))
|
||||
}, [])
|
||||
|
||||
// Helper to show TPM limit toast
|
||||
const showTPMLimitToast = useCallback(() => {
|
||||
const limitDisplay =
|
||||
tpmLimit >= 1000 ? `${tpmLimit / 1000}k` : String(tpmLimit)
|
||||
toast.error(
|
||||
`Rate limit reached (${limitDisplay} tokens/min). Please wait 60 seconds before sending another request.`,
|
||||
{ duration: 8000 },
|
||||
)
|
||||
}, [tpmLimit])
|
||||
// Quota management using extracted hook
|
||||
const quotaManager = useQuotaManager({
|
||||
dailyRequestLimit,
|
||||
dailyTokenLimit,
|
||||
tpmLimit,
|
||||
})
|
||||
|
||||
// Generate a unique session ID for Langfuse tracing (restore from localStorage if available)
|
||||
const [sessionId, setSessionId] = useState(() => {
|
||||
@@ -470,14 +262,6 @@ export default function ChatPanel({
|
||||
validationError,
|
||||
)
|
||||
// Return error to model - sendAutomaticallyWhen will trigger retry
|
||||
const errorMessage = `${validationError}
|
||||
|
||||
Please fix the XML issues and call display_diagram again with corrected XML.
|
||||
|
||||
Your failed XML:
|
||||
\`\`\`xml
|
||||
${xml}
|
||||
\`\`\``
|
||||
if (DEBUG) {
|
||||
console.log(
|
||||
"[display_diagram] Adding tool output with state: output-error",
|
||||
@@ -487,7 +271,14 @@ ${xml}
|
||||
tool: "display_diagram",
|
||||
toolCallId: toolCall.toolCallId,
|
||||
state: "output-error",
|
||||
errorText: errorMessage,
|
||||
errorText: `${validationError}
|
||||
|
||||
Please fix the XML issues and call display_diagram again with corrected XML.
|
||||
|
||||
Your failed XML:
|
||||
\`\`\`xml
|
||||
${xml}
|
||||
\`\`\``,
|
||||
})
|
||||
} else {
|
||||
// Success - diagram will be rendered by chat-message-display
|
||||
@@ -645,8 +436,8 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
|
||||
: 0
|
||||
const actualTokens = inputTokens + outputTokens
|
||||
if (actualTokens > 0) {
|
||||
incrementTokenCount(actualTokens)
|
||||
incrementTPMCount(actualTokens)
|
||||
quotaManager.incrementTokenCount(actualTokens)
|
||||
quotaManager.incrementTPMCount(actualTokens)
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -824,20 +615,11 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
|
||||
const toolCallId = `cached-${Date.now()}`
|
||||
|
||||
// Build user message text including any file content
|
||||
let userText = input
|
||||
for (const file of files) {
|
||||
if (isPdfFile(file)) {
|
||||
const extracted = pdfData.get(file)
|
||||
if (extracted?.text) {
|
||||
userText += `\n\n[PDF: ${file.name}]\n${extracted.text}`
|
||||
}
|
||||
} else if (isTextFile(file)) {
|
||||
const extracted = pdfData.get(file)
|
||||
if (extracted?.text) {
|
||||
userText += `\n\n[File: ${file.name}]\n${extracted.text}`
|
||||
}
|
||||
}
|
||||
}
|
||||
const userText = await processFilesAndAppendContent(
|
||||
input,
|
||||
files,
|
||||
pdfData,
|
||||
)
|
||||
|
||||
setMessages([
|
||||
{
|
||||
@@ -875,43 +657,14 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
|
||||
|
||||
// Build user text by concatenating input with pre-extracted text
|
||||
// (Backend only reads first text part, so we must combine them)
|
||||
let userText = input
|
||||
const parts: any[] = []
|
||||
|
||||
if (files.length > 0) {
|
||||
for (const file of files) {
|
||||
if (isPdfFile(file)) {
|
||||
// Use pre-extracted PDF text from pdfData
|
||||
const extracted = pdfData.get(file)
|
||||
if (extracted?.text) {
|
||||
userText += `\n\n[PDF: ${file.name}]\n${extracted.text}`
|
||||
}
|
||||
} else if (isTextFile(file)) {
|
||||
// Use pre-extracted text file content from pdfData
|
||||
const extracted = pdfData.get(file)
|
||||
if (extracted?.text) {
|
||||
userText += `\n\n[File: ${file.name}]\n${extracted.text}`
|
||||
}
|
||||
} else {
|
||||
// Handle as image
|
||||
const reader = new FileReader()
|
||||
const dataUrl = await new Promise<string>(
|
||||
(resolve) => {
|
||||
reader.onload = () =>
|
||||
resolve(reader.result as string)
|
||||
reader.readAsDataURL(file)
|
||||
},
|
||||
const userText = await processFilesAndAppendContent(
|
||||
input,
|
||||
files,
|
||||
pdfData,
|
||||
parts,
|
||||
)
|
||||
|
||||
parts.push({
|
||||
type: "file",
|
||||
url: dataUrl,
|
||||
mediaType: file.type,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Add the combined text as the first part
|
||||
parts.unshift({ type: "text", text: userText })
|
||||
|
||||
@@ -929,56 +682,11 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
|
||||
xmlSnapshotsRef.current.set(messageIndex, chartXml)
|
||||
saveXmlSnapshots()
|
||||
|
||||
// Check daily limit
|
||||
const limitCheck = checkDailyLimit()
|
||||
if (!limitCheck.allowed) {
|
||||
showQuotaLimitToast()
|
||||
return
|
||||
}
|
||||
// Check all quota limits
|
||||
if (!checkAllQuotaLimits()) return
|
||||
|
||||
// Check daily token limit (actual usage tracked after response)
|
||||
const tokenLimitCheck = checkTokenLimit()
|
||||
if (!tokenLimitCheck.allowed) {
|
||||
showTokenLimitToast(tokenLimitCheck.used)
|
||||
return
|
||||
}
|
||||
sendChatMessage(parts, chartXml, previousXml, sessionId)
|
||||
|
||||
// Check TPM (tokens per minute) limit
|
||||
const tpmCheck = checkTPMLimit()
|
||||
if (!tpmCheck.allowed) {
|
||||
showTPMLimitToast()
|
||||
return
|
||||
}
|
||||
|
||||
const accessCode =
|
||||
localStorage.getItem(STORAGE_ACCESS_CODE_KEY) || ""
|
||||
const aiProvider =
|
||||
localStorage.getItem(STORAGE_AI_PROVIDER_KEY) || ""
|
||||
const aiBaseUrl =
|
||||
localStorage.getItem(STORAGE_AI_BASE_URL_KEY) || ""
|
||||
const aiApiKey =
|
||||
localStorage.getItem(STORAGE_AI_API_KEY_KEY) || ""
|
||||
const aiModel = localStorage.getItem(STORAGE_AI_MODEL_KEY) || ""
|
||||
|
||||
sendMessage(
|
||||
{ parts },
|
||||
{
|
||||
body: {
|
||||
xml: chartXml,
|
||||
previousXml,
|
||||
sessionId,
|
||||
},
|
||||
headers: {
|
||||
"x-access-code": accessCode,
|
||||
...(aiProvider && { "x-ai-provider": aiProvider }),
|
||||
...(aiBaseUrl && { "x-ai-base-url": aiBaseUrl }),
|
||||
...(aiApiKey && { "x-ai-api-key": aiApiKey }),
|
||||
...(aiModel && { "x-ai-model": aiModel }),
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
incrementRequestCount()
|
||||
// Token count is tracked in onFinish with actual server usage
|
||||
setInput("")
|
||||
setFiles([])
|
||||
@@ -994,81 +702,122 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
|
||||
setInput(e.target.value)
|
||||
}
|
||||
|
||||
const handleFileChange = async (newFiles: File[]) => {
|
||||
setFiles(newFiles)
|
||||
|
||||
// Extract text immediately for new PDF/text files
|
||||
for (const file of newFiles) {
|
||||
const needsExtraction =
|
||||
(isPdfFile(file) || isTextFile(file)) && !pdfData.has(file)
|
||||
if (needsExtraction) {
|
||||
// Mark as extracting
|
||||
setPdfData((prev) => {
|
||||
const next = new Map(prev)
|
||||
next.set(file, {
|
||||
text: "",
|
||||
charCount: 0,
|
||||
isExtracting: true,
|
||||
})
|
||||
return next
|
||||
})
|
||||
|
||||
// Extract text asynchronously
|
||||
try {
|
||||
let text: string
|
||||
if (isPdfFile(file)) {
|
||||
text = await extractPdfText(file)
|
||||
} else {
|
||||
text = await extractTextFileContent(file)
|
||||
// Helper functions for message actions (regenerate/edit)
|
||||
// Extract previous XML snapshot before a given message index
|
||||
const getPreviousXml = (beforeIndex: number): string => {
|
||||
const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys())
|
||||
.filter((k) => k < beforeIndex)
|
||||
.sort((a, b) => b - a)
|
||||
return snapshotKeys.length > 0
|
||||
? xmlSnapshotsRef.current.get(snapshotKeys[0]) || ""
|
||||
: ""
|
||||
}
|
||||
|
||||
// Check character limit
|
||||
if (text.length > MAX_EXTRACTED_CHARS) {
|
||||
const limitK = MAX_EXTRACTED_CHARS / 1000
|
||||
toast.error(
|
||||
`${file.name}: Content exceeds ${limitK}k character limit (${(text.length / 1000).toFixed(1)}k chars)`,
|
||||
// Restore diagram from snapshot and update ref
|
||||
const restoreDiagramFromSnapshot = (savedXml: string) => {
|
||||
onDisplayChart(savedXml, true) // Skip validation for trusted snapshots
|
||||
chartXMLRef.current = savedXml
|
||||
}
|
||||
|
||||
// Clean up snapshots after a given message index
|
||||
const cleanupSnapshotsAfter = (messageIndex: number) => {
|
||||
for (const key of xmlSnapshotsRef.current.keys()) {
|
||||
if (key > messageIndex) {
|
||||
xmlSnapshotsRef.current.delete(key)
|
||||
}
|
||||
}
|
||||
saveXmlSnapshots()
|
||||
}
|
||||
|
||||
// Check all quota limits (daily requests, tokens, TPM)
|
||||
const checkAllQuotaLimits = (): boolean => {
|
||||
const limitCheck = quotaManager.checkDailyLimit()
|
||||
if (!limitCheck.allowed) {
|
||||
quotaManager.showQuotaLimitToast()
|
||||
return false
|
||||
}
|
||||
|
||||
const tokenLimitCheck = quotaManager.checkTokenLimit()
|
||||
if (!tokenLimitCheck.allowed) {
|
||||
quotaManager.showTokenLimitToast(tokenLimitCheck.used)
|
||||
return false
|
||||
}
|
||||
|
||||
const tpmCheck = quotaManager.checkTPMLimit()
|
||||
if (!tpmCheck.allowed) {
|
||||
quotaManager.showTPMLimitToast()
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// Send chat message with headers and increment quota
|
||||
const sendChatMessage = (
|
||||
parts: any,
|
||||
xml: string,
|
||||
previousXml: string,
|
||||
sessionId: string,
|
||||
) => {
|
||||
const config = getAIConfig()
|
||||
|
||||
sendMessage(
|
||||
{ parts },
|
||||
{
|
||||
body: { xml, previousXml, sessionId },
|
||||
headers: {
|
||||
"x-access-code": config.accessCode,
|
||||
...(config.aiProvider && {
|
||||
"x-ai-provider": config.aiProvider,
|
||||
}),
|
||||
...(config.aiBaseUrl && {
|
||||
"x-ai-base-url": config.aiBaseUrl,
|
||||
}),
|
||||
...(config.aiApiKey && { "x-ai-api-key": config.aiApiKey }),
|
||||
...(config.aiModel && { "x-ai-model": config.aiModel }),
|
||||
},
|
||||
},
|
||||
)
|
||||
setPdfData((prev) => {
|
||||
const next = new Map(prev)
|
||||
next.delete(file)
|
||||
return next
|
||||
})
|
||||
// Remove the file from the list
|
||||
setFiles((prev) => prev.filter((f) => f !== file))
|
||||
continue
|
||||
quotaManager.incrementRequestCount()
|
||||
}
|
||||
|
||||
setPdfData((prev) => {
|
||||
const next = new Map(prev)
|
||||
next.set(file, {
|
||||
text,
|
||||
charCount: text.length,
|
||||
isExtracting: false,
|
||||
})
|
||||
return next
|
||||
})
|
||||
} catch (error) {
|
||||
console.error("Failed to extract text:", error)
|
||||
toast.error(`Failed to read file: ${file.name}`)
|
||||
setPdfData((prev) => {
|
||||
const next = new Map(prev)
|
||||
next.delete(file)
|
||||
return next
|
||||
})
|
||||
// Process files and append content to user text (handles PDF, text, and optionally images)
|
||||
const processFilesAndAppendContent = async (
|
||||
baseText: string,
|
||||
files: File[],
|
||||
pdfData: Map<File, FileData>,
|
||||
imageParts?: any[],
|
||||
): Promise<string> => {
|
||||
let userText = baseText
|
||||
|
||||
for (const file of files) {
|
||||
if (isPdfFile(file)) {
|
||||
const extracted = pdfData.get(file)
|
||||
if (extracted?.text) {
|
||||
userText += `\n\n[PDF: ${file.name}]\n${extracted.text}`
|
||||
}
|
||||
} else if (isTextFile(file)) {
|
||||
const extracted = pdfData.get(file)
|
||||
if (extracted?.text) {
|
||||
userText += `\n\n[File: ${file.name}]\n${extracted.text}`
|
||||
}
|
||||
} else if (imageParts) {
|
||||
// Handle as image (only if imageParts array provided)
|
||||
const reader = new FileReader()
|
||||
const dataUrl = await new Promise<string>((resolve) => {
|
||||
reader.onload = () => resolve(reader.result as string)
|
||||
reader.readAsDataURL(file)
|
||||
})
|
||||
|
||||
imageParts.push({
|
||||
type: "file",
|
||||
url: dataUrl,
|
||||
mediaType: file.type,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Clean up pdfData for removed files
|
||||
setPdfData((prev) => {
|
||||
const next = new Map(prev)
|
||||
for (const key of prev.keys()) {
|
||||
if (!newFiles.includes(key)) {
|
||||
next.delete(key)
|
||||
}
|
||||
}
|
||||
return next
|
||||
})
|
||||
return userText
|
||||
}
|
||||
|
||||
const handleRegenerate = async (messageIndex: number) => {
|
||||
@@ -1103,28 +852,12 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
|
||||
return
|
||||
}
|
||||
|
||||
// Get previous XML (snapshot before the one being regenerated)
|
||||
const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys())
|
||||
.filter((k) => k < userMessageIndex)
|
||||
.sort((a, b) => b - a)
|
||||
const previousXml =
|
||||
snapshotKeys.length > 0
|
||||
? xmlSnapshotsRef.current.get(snapshotKeys[0]) || ""
|
||||
: ""
|
||||
|
||||
// Restore the diagram to the saved state (skip validation for trusted snapshots)
|
||||
onDisplayChart(savedXml, true)
|
||||
|
||||
// Update ref directly to ensure edit_diagram has the correct XML
|
||||
chartXMLRef.current = savedXml
|
||||
// Get previous XML and restore diagram state
|
||||
const previousXml = getPreviousXml(userMessageIndex)
|
||||
restoreDiagramFromSnapshot(savedXml)
|
||||
|
||||
// Clean up snapshots for messages after the user message (they will be removed)
|
||||
for (const key of xmlSnapshotsRef.current.keys()) {
|
||||
if (key > userMessageIndex) {
|
||||
xmlSnapshotsRef.current.delete(key)
|
||||
}
|
||||
}
|
||||
saveXmlSnapshots()
|
||||
cleanupSnapshotsAfter(userMessageIndex)
|
||||
|
||||
// Remove the user message AND assistant message onwards (sendMessage will re-add the user message)
|
||||
// Use flushSync to ensure state update is processed synchronously before sending
|
||||
@@ -1133,53 +866,12 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
|
||||
setMessages(newMessages)
|
||||
})
|
||||
|
||||
// Check daily limit
|
||||
const limitCheck = checkDailyLimit()
|
||||
if (!limitCheck.allowed) {
|
||||
showQuotaLimitToast()
|
||||
return
|
||||
}
|
||||
|
||||
// Check daily token limit (actual usage tracked after response)
|
||||
const tokenLimitCheck = checkTokenLimit()
|
||||
if (!tokenLimitCheck.allowed) {
|
||||
showTokenLimitToast(tokenLimitCheck.used)
|
||||
return
|
||||
}
|
||||
|
||||
// Check TPM (tokens per minute) limit
|
||||
const tpmCheck = checkTPMLimit()
|
||||
if (!tpmCheck.allowed) {
|
||||
showTPMLimitToast()
|
||||
return
|
||||
}
|
||||
// Check all quota limits
|
||||
if (!checkAllQuotaLimits()) return
|
||||
|
||||
// Now send the message after state is guaranteed to be updated
|
||||
const accessCode = localStorage.getItem(STORAGE_ACCESS_CODE_KEY) || ""
|
||||
const aiProvider = localStorage.getItem(STORAGE_AI_PROVIDER_KEY) || ""
|
||||
const aiBaseUrl = localStorage.getItem(STORAGE_AI_BASE_URL_KEY) || ""
|
||||
const aiApiKey = localStorage.getItem(STORAGE_AI_API_KEY_KEY) || ""
|
||||
const aiModel = localStorage.getItem(STORAGE_AI_MODEL_KEY) || ""
|
||||
sendChatMessage(userParts, savedXml, previousXml, sessionId)
|
||||
|
||||
sendMessage(
|
||||
{ parts: userParts },
|
||||
{
|
||||
body: {
|
||||
xml: savedXml,
|
||||
previousXml,
|
||||
sessionId,
|
||||
},
|
||||
headers: {
|
||||
"x-access-code": accessCode,
|
||||
...(aiProvider && { "x-ai-provider": aiProvider }),
|
||||
...(aiBaseUrl && { "x-ai-base-url": aiBaseUrl }),
|
||||
...(aiApiKey && { "x-ai-api-key": aiApiKey }),
|
||||
...(aiModel && { "x-ai-model": aiModel }),
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
incrementRequestCount()
|
||||
// Token count is tracked in onFinish with actual server usage
|
||||
}
|
||||
|
||||
@@ -1200,28 +892,12 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
|
||||
return
|
||||
}
|
||||
|
||||
// Get previous XML (snapshot before the one being edited)
|
||||
const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys())
|
||||
.filter((k) => k < messageIndex)
|
||||
.sort((a, b) => b - a)
|
||||
const previousXml =
|
||||
snapshotKeys.length > 0
|
||||
? xmlSnapshotsRef.current.get(snapshotKeys[0]) || ""
|
||||
: ""
|
||||
|
||||
// Restore the diagram to the saved state (skip validation for trusted snapshots)
|
||||
onDisplayChart(savedXml, true)
|
||||
|
||||
// Update ref directly to ensure edit_diagram has the correct XML
|
||||
chartXMLRef.current = savedXml
|
||||
// Get previous XML and restore diagram state
|
||||
const previousXml = getPreviousXml(messageIndex)
|
||||
restoreDiagramFromSnapshot(savedXml)
|
||||
|
||||
// Clean up snapshots for messages after the user message (they will be removed)
|
||||
for (const key of xmlSnapshotsRef.current.keys()) {
|
||||
if (key > messageIndex) {
|
||||
xmlSnapshotsRef.current.delete(key)
|
||||
}
|
||||
}
|
||||
saveXmlSnapshots()
|
||||
cleanupSnapshotsAfter(messageIndex)
|
||||
|
||||
// Create new parts with updated text
|
||||
const newParts = message.parts?.map((part: any) => {
|
||||
@@ -1238,53 +914,11 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
|
||||
setMessages(newMessages)
|
||||
})
|
||||
|
||||
// Check daily limit
|
||||
const limitCheck = checkDailyLimit()
|
||||
if (!limitCheck.allowed) {
|
||||
showQuotaLimitToast()
|
||||
return
|
||||
}
|
||||
|
||||
// Check daily token limit (actual usage tracked after response)
|
||||
const tokenLimitCheck = checkTokenLimit()
|
||||
if (!tokenLimitCheck.allowed) {
|
||||
showTokenLimitToast(tokenLimitCheck.used)
|
||||
return
|
||||
}
|
||||
|
||||
// Check TPM (tokens per minute) limit
|
||||
const tpmCheck = checkTPMLimit()
|
||||
if (!tpmCheck.allowed) {
|
||||
showTPMLimitToast()
|
||||
return
|
||||
}
|
||||
// Check all quota limits
|
||||
if (!checkAllQuotaLimits()) return
|
||||
|
||||
// Now send the edited message after state is guaranteed to be updated
|
||||
const accessCode = localStorage.getItem(STORAGE_ACCESS_CODE_KEY) || ""
|
||||
const aiProvider = localStorage.getItem(STORAGE_AI_PROVIDER_KEY) || ""
|
||||
const aiBaseUrl = localStorage.getItem(STORAGE_AI_BASE_URL_KEY) || ""
|
||||
const aiApiKey = localStorage.getItem(STORAGE_AI_API_KEY_KEY) || ""
|
||||
const aiModel = localStorage.getItem(STORAGE_AI_MODEL_KEY) || ""
|
||||
|
||||
sendMessage(
|
||||
{ parts: newParts },
|
||||
{
|
||||
body: {
|
||||
xml: savedXml,
|
||||
previousXml,
|
||||
sessionId,
|
||||
},
|
||||
headers: {
|
||||
"x-access-code": accessCode,
|
||||
...(aiProvider && { "x-ai-provider": aiProvider }),
|
||||
...(aiBaseUrl && { "x-ai-base-url": aiBaseUrl }),
|
||||
...(aiApiKey && { "x-ai-api-key": aiApiKey }),
|
||||
...(aiModel && { "x-ai-model": aiModel }),
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
incrementRequestCount()
|
||||
sendChatMessage(newParts, savedXml, previousXml, sessionId)
|
||||
// Token count is tracked in onFinish with actual server usage
|
||||
}
|
||||
|
||||
|
||||
@@ -81,9 +81,11 @@ https://github.com/user-attachments/assets/b2eef5f3-b335-4e71-a755-dc2e80931979
|
||||
|
||||
- **LLM驱动的图表创建**:利用大语言模型通过自然语言命令直接创建和操作draw.io图表
|
||||
- **基于图像的图表复制**:上传现有图表或图像,让AI自动复制和增强
|
||||
- **PDF和文本文件上传**:上传PDF文档和文本文件,提取内容并从现有文档生成图表
|
||||
- **AI推理过程显示**:查看支持模型的AI思考过程(OpenAI o1/o3、Gemini、Claude等)
|
||||
- **图表历史记录**:全面的版本控制,跟踪所有更改,允许您查看和恢复AI编辑前的图表版本
|
||||
- **交互式聊天界面**:与AI实时对话来完善您的图表
|
||||
- **AWS架构图支持**:专门支持生成AWS架构图
|
||||
- **云架构图支持**:专门支持生成云架构图(AWS、GCP、Azure)
|
||||
- **动画连接器**:在图表元素之间创建动态动画连接器,实现更好的可视化效果
|
||||
|
||||
## 快速开始
|
||||
|
||||
@@ -81,9 +81,11 @@ https://github.com/user-attachments/assets/b2eef5f3-b335-4e71-a755-dc2e80931979
|
||||
|
||||
- **LLM搭載のダイアグラム作成**:大規模言語モデルを活用して、自然言語コマンドで直接draw.ioダイアグラムを作成・操作
|
||||
- **画像ベースのダイアグラム複製**:既存のダイアグラムや画像をアップロードし、AIが自動的に複製・強化
|
||||
- **PDFとテキストファイルのアップロード**:PDFドキュメントやテキストファイルをアップロードして、既存のドキュメントからコンテンツを抽出し、ダイアグラムを生成
|
||||
- **AI推論プロセス表示**:サポートされているモデル(OpenAI o1/o3、Gemini、Claudeなど)のAIの思考プロセスを表示
|
||||
- **ダイアグラム履歴**:すべての変更を追跡する包括的なバージョン管理。AI編集前のダイアグラムの以前のバージョンを表示・復元可能
|
||||
- **インタラクティブなチャットインターフェース**:AIとリアルタイムでコミュニケーションしてダイアグラムを改善
|
||||
- **AWSアーキテクチャダイアグラムサポート**:AWSアーキテクチャダイアグラムの生成を専門的にサポート
|
||||
- **クラウドアーキテクチャダイアグラムサポート**:クラウドアーキテクチャダイアグラムの生成を専門的にサポート(AWS、GCP、Azure)
|
||||
- **アニメーションコネクタ**:より良い可視化のためにダイアグラム要素間に動的でアニメーション化されたコネクタを作成
|
||||
|
||||
## はじめに
|
||||
|
||||
@@ -41,9 +41,13 @@ AI_MODEL=global.anthropic.claude-sonnet-4-5-20250929-v1:0
|
||||
# GOOGLE_THINKING_LEVEL=high # Optional: Gemini 3 thinking level (low/high)
|
||||
|
||||
# Azure OpenAI Configuration
|
||||
# Configure endpoint using ONE of these methods:
|
||||
# 1. AZURE_RESOURCE_NAME - SDK constructs: https://{name}.openai.azure.com/openai/v1{path}
|
||||
# 2. AZURE_BASE_URL - SDK appends /v1{path} to your URL
|
||||
# If both are set, AZURE_BASE_URL takes precedence.
|
||||
# AZURE_RESOURCE_NAME=your-resource-name
|
||||
# AZURE_API_KEY=...
|
||||
# AZURE_BASE_URL=https://your-resource.openai.azure.com # Optional: Custom endpoint (overrides resourceName)
|
||||
# AZURE_BASE_URL=https://your-resource.openai.azure.com/openai # Alternative: Custom endpoint
|
||||
# AZURE_REASONING_EFFORT=low # Optional: Azure reasoning effort (low, medium, high)
|
||||
# AZURE_REASONING_SUMMARY=detailed
|
||||
|
||||
@@ -86,3 +90,4 @@ AI_MODEL=global.anthropic.claude-sonnet-4-5-20250929-v1:0
|
||||
# Enable PDF file upload to extract text and generate diagrams
|
||||
# Enabled by default. Set to "false" to disable.
|
||||
# ENABLE_PDF_INPUT=true
|
||||
# NEXT_PUBLIC_MAX_EXTRACTED_CHARS=150000 # Max characters for PDF/text extraction (default: 150000)
|
||||
|
||||
26
lib/ai-config.ts
Normal file
26
lib/ai-config.ts
Normal file
@@ -0,0 +1,26 @@
|
||||
import { STORAGE_KEYS } from "./storage"
|
||||
|
||||
/**
|
||||
* Get AI configuration from localStorage.
|
||||
* Returns API keys and settings for custom AI providers.
|
||||
* Used to override server defaults when user provides their own API key.
|
||||
*/
|
||||
export function getAIConfig() {
|
||||
if (typeof window === "undefined") {
|
||||
return {
|
||||
accessCode: "",
|
||||
aiProvider: "",
|
||||
aiBaseUrl: "",
|
||||
aiApiKey: "",
|
||||
aiModel: "",
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
accessCode: localStorage.getItem(STORAGE_KEYS.accessCode) || "",
|
||||
aiProvider: localStorage.getItem(STORAGE_KEYS.aiProvider) || "",
|
||||
aiBaseUrl: localStorage.getItem(STORAGE_KEYS.aiBaseUrl) || "",
|
||||
aiApiKey: localStorage.getItem(STORAGE_KEYS.aiApiKey) || "",
|
||||
aiModel: localStorage.getItem(STORAGE_KEYS.aiModel) || "",
|
||||
}
|
||||
}
|
||||
@@ -572,10 +572,15 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
|
||||
case "azure": {
|
||||
const apiKey = overrides?.apiKey || process.env.AZURE_API_KEY
|
||||
const baseURL = overrides?.baseUrl || process.env.AZURE_BASE_URL
|
||||
if (baseURL || overrides?.apiKey) {
|
||||
const resourceName = process.env.AZURE_RESOURCE_NAME
|
||||
// Azure requires either baseURL or resourceName to construct the endpoint
|
||||
// resourceName constructs: https://{resourceName}.openai.azure.com/openai/v1{path}
|
||||
if (baseURL || resourceName || overrides?.apiKey) {
|
||||
const customAzure = createAzure({
|
||||
apiKey,
|
||||
// baseURL takes precedence over resourceName per SDK behavior
|
||||
...(baseURL && { baseURL }),
|
||||
...(!baseURL && resourceName && { resourceName }),
|
||||
})
|
||||
model = customAzure(modelId)
|
||||
} else {
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
import { extractText, getDocumentProxy } from "unpdf"
|
||||
|
||||
// Maximum characters allowed for extracted text
|
||||
export const MAX_EXTRACTED_CHARS = 150000 // 150k chars
|
||||
// Maximum characters allowed for extracted text (configurable via env)
|
||||
const DEFAULT_MAX_EXTRACTED_CHARS = 150000 // 150k chars
|
||||
export const MAX_EXTRACTED_CHARS =
|
||||
Number(process.env.NEXT_PUBLIC_MAX_EXTRACTED_CHARS) ||
|
||||
DEFAULT_MAX_EXTRACTED_CHARS
|
||||
|
||||
// Text file extensions we support
|
||||
const TEXT_EXTENSIONS = [
|
||||
|
||||
27
lib/storage.ts
Normal file
27
lib/storage.ts
Normal file
@@ -0,0 +1,27 @@
|
||||
// Centralized localStorage keys
|
||||
// Consolidates all storage keys from chat-panel.tsx and settings-dialog.tsx
|
||||
|
||||
export const STORAGE_KEYS = {
|
||||
// Chat data
|
||||
messages: "next-ai-draw-io-messages",
|
||||
xmlSnapshots: "next-ai-draw-io-xml-snapshots",
|
||||
diagramXml: "next-ai-draw-io-diagram-xml",
|
||||
sessionId: "next-ai-draw-io-session-id",
|
||||
|
||||
// Quota tracking
|
||||
requestCount: "next-ai-draw-io-request-count",
|
||||
requestDate: "next-ai-draw-io-request-date",
|
||||
tokenCount: "next-ai-draw-io-token-count",
|
||||
tokenDate: "next-ai-draw-io-token-date",
|
||||
tpmCount: "next-ai-draw-io-tpm-count",
|
||||
tpmMinute: "next-ai-draw-io-tpm-minute",
|
||||
|
||||
// Settings
|
||||
accessCode: "next-ai-draw-io-access-code",
|
||||
closeProtection: "next-ai-draw-io-close-protection",
|
||||
accessCodeRequired: "next-ai-draw-io-access-code-required",
|
||||
aiProvider: "next-ai-draw-io-ai-provider",
|
||||
aiBaseUrl: "next-ai-draw-io-ai-base-url",
|
||||
aiApiKey: "next-ai-draw-io-ai-api-key",
|
||||
aiModel: "next-ai-draw-io-ai-model",
|
||||
} as const
|
||||
110
lib/use-file-processor.tsx
Normal file
110
lib/use-file-processor.tsx
Normal file
@@ -0,0 +1,110 @@
|
||||
"use client"
|
||||
|
||||
import { useState } from "react"
|
||||
import { toast } from "sonner"
|
||||
import {
|
||||
extractPdfText,
|
||||
extractTextFileContent,
|
||||
isPdfFile,
|
||||
isTextFile,
|
||||
MAX_EXTRACTED_CHARS,
|
||||
} from "@/lib/pdf-utils"
|
||||
|
||||
export interface FileData {
|
||||
text: string
|
||||
charCount: number
|
||||
isExtracting: boolean
|
||||
}
|
||||
|
||||
/**
|
||||
* Hook for processing file uploads, especially PDFs and text files.
|
||||
* Handles text extraction, character limit validation, and cleanup.
|
||||
*/
|
||||
export function useFileProcessor() {
|
||||
const [files, setFiles] = useState<File[]>([])
|
||||
const [pdfData, setPdfData] = useState<Map<File, FileData>>(new Map())
|
||||
|
||||
const handleFileChange = async (newFiles: File[]) => {
|
||||
setFiles(newFiles)
|
||||
|
||||
// Extract text immediately for new PDF/text files
|
||||
for (const file of newFiles) {
|
||||
const needsExtraction =
|
||||
(isPdfFile(file) || isTextFile(file)) && !pdfData.has(file)
|
||||
if (needsExtraction) {
|
||||
// Mark as extracting
|
||||
setPdfData((prev) => {
|
||||
const next = new Map(prev)
|
||||
next.set(file, {
|
||||
text: "",
|
||||
charCount: 0,
|
||||
isExtracting: true,
|
||||
})
|
||||
return next
|
||||
})
|
||||
|
||||
// Extract text asynchronously
|
||||
try {
|
||||
let text: string
|
||||
if (isPdfFile(file)) {
|
||||
text = await extractPdfText(file)
|
||||
} else {
|
||||
text = await extractTextFileContent(file)
|
||||
}
|
||||
|
||||
// Check character limit
|
||||
if (text.length > MAX_EXTRACTED_CHARS) {
|
||||
const limitK = MAX_EXTRACTED_CHARS / 1000
|
||||
toast.error(
|
||||
`${file.name}: Content exceeds ${limitK}k character limit (${(text.length / 1000).toFixed(1)}k chars)`,
|
||||
)
|
||||
setPdfData((prev) => {
|
||||
const next = new Map(prev)
|
||||
next.delete(file)
|
||||
return next
|
||||
})
|
||||
// Remove the file from the list
|
||||
setFiles((prev) => prev.filter((f) => f !== file))
|
||||
continue
|
||||
}
|
||||
|
||||
setPdfData((prev) => {
|
||||
const next = new Map(prev)
|
||||
next.set(file, {
|
||||
text,
|
||||
charCount: text.length,
|
||||
isExtracting: false,
|
||||
})
|
||||
return next
|
||||
})
|
||||
} catch (error) {
|
||||
console.error("Failed to extract text:", error)
|
||||
toast.error(`Failed to read file: ${file.name}`)
|
||||
setPdfData((prev) => {
|
||||
const next = new Map(prev)
|
||||
next.delete(file)
|
||||
return next
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Clean up pdfData for removed files
|
||||
setPdfData((prev) => {
|
||||
const next = new Map(prev)
|
||||
for (const key of prev.keys()) {
|
||||
if (!newFiles.includes(key)) {
|
||||
next.delete(key)
|
||||
}
|
||||
}
|
||||
return next
|
||||
})
|
||||
}
|
||||
|
||||
return {
|
||||
files,
|
||||
pdfData,
|
||||
handleFileChange,
|
||||
setFiles, // Export for external control (e.g., clearing files)
|
||||
}
|
||||
}
|
||||
247
lib/use-quota-manager.tsx
Normal file
247
lib/use-quota-manager.tsx
Normal file
@@ -0,0 +1,247 @@
|
||||
"use client"
|
||||
|
||||
import { useCallback, useMemo } from "react"
|
||||
import { toast } from "sonner"
|
||||
import { QuotaLimitToast } from "@/components/quota-limit-toast"
|
||||
import { STORAGE_KEYS } from "@/lib/storage"
|
||||
|
||||
export interface QuotaConfig {
|
||||
dailyRequestLimit: number
|
||||
dailyTokenLimit: number
|
||||
tpmLimit: number
|
||||
}
|
||||
|
||||
export interface QuotaCheckResult {
|
||||
allowed: boolean
|
||||
remaining: number
|
||||
used: number
|
||||
}
|
||||
|
||||
/**
|
||||
* Hook for managing request/token quotas and rate limiting.
|
||||
* Handles three types of limits:
|
||||
* - Daily request limit
|
||||
* - Daily token limit
|
||||
* - Tokens per minute (TPM) rate limit
|
||||
*
|
||||
* Users with their own API key bypass all limits.
|
||||
*/
|
||||
export function useQuotaManager(config: QuotaConfig): {
|
||||
hasOwnApiKey: () => boolean
|
||||
checkDailyLimit: () => QuotaCheckResult
|
||||
checkTokenLimit: () => QuotaCheckResult
|
||||
checkTPMLimit: () => QuotaCheckResult
|
||||
incrementRequestCount: () => void
|
||||
incrementTokenCount: (tokens: number) => void
|
||||
incrementTPMCount: (tokens: number) => void
|
||||
showQuotaLimitToast: () => void
|
||||
showTokenLimitToast: (used: number) => void
|
||||
showTPMLimitToast: () => void
|
||||
} {
|
||||
const { dailyRequestLimit, dailyTokenLimit, tpmLimit } = config
|
||||
|
||||
// Check if user has their own API key configured (bypass limits)
|
||||
const hasOwnApiKey = useCallback((): boolean => {
|
||||
const provider = localStorage.getItem(STORAGE_KEYS.aiProvider)
|
||||
const apiKey = localStorage.getItem(STORAGE_KEYS.aiApiKey)
|
||||
return !!(provider && apiKey)
|
||||
}, [])
|
||||
|
||||
// Generic helper: Parse count from localStorage with NaN guard
|
||||
const parseStorageCount = (key: string): number => {
|
||||
const count = parseInt(localStorage.getItem(key) || "0", 10)
|
||||
return Number.isNaN(count) ? 0 : count
|
||||
}
|
||||
|
||||
// Generic helper: Create quota checker factory
|
||||
const createQuotaChecker = useCallback(
|
||||
(
|
||||
getTimeKey: () => string,
|
||||
timeStorageKey: string,
|
||||
countStorageKey: string,
|
||||
limit: number,
|
||||
) => {
|
||||
return (): QuotaCheckResult => {
|
||||
if (hasOwnApiKey())
|
||||
return { allowed: true, remaining: -1, used: 0 }
|
||||
if (limit <= 0) return { allowed: true, remaining: -1, used: 0 }
|
||||
|
||||
const currentTime = getTimeKey()
|
||||
const storedTime = localStorage.getItem(timeStorageKey)
|
||||
let count = parseStorageCount(countStorageKey)
|
||||
|
||||
if (storedTime !== currentTime) {
|
||||
count = 0
|
||||
localStorage.setItem(timeStorageKey, currentTime)
|
||||
localStorage.setItem(countStorageKey, "0")
|
||||
}
|
||||
|
||||
return {
|
||||
allowed: count < limit,
|
||||
remaining: limit - count,
|
||||
used: count,
|
||||
}
|
||||
}
|
||||
},
|
||||
[hasOwnApiKey],
|
||||
)
|
||||
|
||||
// Generic helper: Create quota incrementer factory
|
||||
const createQuotaIncrementer = useCallback(
|
||||
(
|
||||
getTimeKey: () => string,
|
||||
timeStorageKey: string,
|
||||
countStorageKey: string,
|
||||
validateInput: boolean = false,
|
||||
) => {
|
||||
return (tokens: number = 1): void => {
|
||||
if (validateInput && (!Number.isFinite(tokens) || tokens <= 0))
|
||||
return
|
||||
|
||||
const currentTime = getTimeKey()
|
||||
const storedTime = localStorage.getItem(timeStorageKey)
|
||||
let count = parseStorageCount(countStorageKey)
|
||||
|
||||
if (storedTime !== currentTime) {
|
||||
count = 0
|
||||
localStorage.setItem(timeStorageKey, currentTime)
|
||||
}
|
||||
|
||||
localStorage.setItem(countStorageKey, String(count + tokens))
|
||||
}
|
||||
},
|
||||
[],
|
||||
)
|
||||
|
||||
// Check daily request limit
|
||||
const checkDailyLimit = useMemo(
|
||||
() =>
|
||||
createQuotaChecker(
|
||||
() => new Date().toDateString(),
|
||||
STORAGE_KEYS.requestDate,
|
||||
STORAGE_KEYS.requestCount,
|
||||
dailyRequestLimit,
|
||||
),
|
||||
[createQuotaChecker, dailyRequestLimit],
|
||||
)
|
||||
|
||||
// Increment request count
|
||||
const incrementRequestCount = useMemo(
|
||||
() =>
|
||||
createQuotaIncrementer(
|
||||
() => new Date().toDateString(),
|
||||
STORAGE_KEYS.requestDate,
|
||||
STORAGE_KEYS.requestCount,
|
||||
false,
|
||||
),
|
||||
[createQuotaIncrementer],
|
||||
)
|
||||
|
||||
// Show quota limit toast (request-based)
|
||||
const showQuotaLimitToast = useCallback(() => {
|
||||
toast.custom(
|
||||
(t) => (
|
||||
<QuotaLimitToast
|
||||
used={dailyRequestLimit}
|
||||
limit={dailyRequestLimit}
|
||||
onDismiss={() => toast.dismiss(t)}
|
||||
/>
|
||||
),
|
||||
{ duration: 15000 },
|
||||
)
|
||||
}, [dailyRequestLimit])
|
||||
|
||||
// Check daily token limit
|
||||
const checkTokenLimit = useMemo(
|
||||
() =>
|
||||
createQuotaChecker(
|
||||
() => new Date().toDateString(),
|
||||
STORAGE_KEYS.tokenDate,
|
||||
STORAGE_KEYS.tokenCount,
|
||||
dailyTokenLimit,
|
||||
),
|
||||
[createQuotaChecker, dailyTokenLimit],
|
||||
)
|
||||
|
||||
// Increment token count
|
||||
const incrementTokenCount = useMemo(
|
||||
() =>
|
||||
createQuotaIncrementer(
|
||||
() => new Date().toDateString(),
|
||||
STORAGE_KEYS.tokenDate,
|
||||
STORAGE_KEYS.tokenCount,
|
||||
true, // Validate input tokens
|
||||
),
|
||||
[createQuotaIncrementer],
|
||||
)
|
||||
|
||||
// Show token limit toast
|
||||
const showTokenLimitToast = useCallback(
|
||||
(used: number) => {
|
||||
toast.custom(
|
||||
(t) => (
|
||||
<QuotaLimitToast
|
||||
type="token"
|
||||
used={used}
|
||||
limit={dailyTokenLimit}
|
||||
onDismiss={() => toast.dismiss(t)}
|
||||
/>
|
||||
),
|
||||
{ duration: 15000 },
|
||||
)
|
||||
},
|
||||
[dailyTokenLimit],
|
||||
)
|
||||
|
||||
// Check TPM (tokens per minute) limit
|
||||
const checkTPMLimit = useMemo(
|
||||
() =>
|
||||
createQuotaChecker(
|
||||
() => Math.floor(Date.now() / 60000).toString(),
|
||||
STORAGE_KEYS.tpmMinute,
|
||||
STORAGE_KEYS.tpmCount,
|
||||
tpmLimit,
|
||||
),
|
||||
[createQuotaChecker, tpmLimit],
|
||||
)
|
||||
|
||||
// Increment TPM count
|
||||
const incrementTPMCount = useMemo(
|
||||
() =>
|
||||
createQuotaIncrementer(
|
||||
() => Math.floor(Date.now() / 60000).toString(),
|
||||
STORAGE_KEYS.tpmMinute,
|
||||
STORAGE_KEYS.tpmCount,
|
||||
true, // Validate input tokens
|
||||
),
|
||||
[createQuotaIncrementer],
|
||||
)
|
||||
|
||||
// Show TPM limit toast
|
||||
const showTPMLimitToast = useCallback(() => {
|
||||
const limitDisplay =
|
||||
tpmLimit >= 1000 ? `${tpmLimit / 1000}k` : String(tpmLimit)
|
||||
toast.error(
|
||||
`Rate limit reached (${limitDisplay} tokens/min). Please wait 60 seconds before sending another request.`,
|
||||
{ duration: 8000 },
|
||||
)
|
||||
}, [tpmLimit])
|
||||
|
||||
return {
|
||||
// Check functions
|
||||
hasOwnApiKey,
|
||||
checkDailyLimit,
|
||||
checkTokenLimit,
|
||||
checkTPMLimit,
|
||||
|
||||
// Increment functions
|
||||
incrementRequestCount,
|
||||
incrementTokenCount,
|
||||
incrementTPMCount,
|
||||
|
||||
// Toast functions
|
||||
showQuotaLimitToast,
|
||||
showTokenLimitToast,
|
||||
showTPMLimitToast,
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "next-ai-draw-io",
|
||||
"version": "0.3.0",
|
||||
"version": "0.4.0",
|
||||
"license": "Apache-2.0",
|
||||
"private": true,
|
||||
"scripts": {
|
||||
|
||||
File diff suppressed because one or more lines are too long
Reference in New Issue
Block a user