From b43f29a5fe6e3056c76a50deb9607172a3317624 Mon Sep 17 00:00:00 2001 From: Ahmed Khaleel Date: Fri, 27 Dec 2024 07:48:29 -0500 Subject: [PATCH] allow users to use personal api key for large repos --- backend/app/routers/generate.py | 16 ++- backend/app/services/claude_service.py | 12 +- package.json | 1 + pnpm-lock.yaml | 188 +++++++++++++++++++++++++ src/app/[username]/[repo]/page.tsx | 25 +++- src/components/api-key-button.tsx | 18 +++ src/components/api-key-dialog.tsx | 69 +++++++++ src/components/ui/dialog.tsx | 122 ++++++++++++++++ src/hooks/useDiagram.ts | 40 ++++++ src/lib/fetch-backend.ts | 6 +- 10 files changed, 488 insertions(+), 9 deletions(-) create mode 100644 src/components/api-key-button.tsx create mode 100644 src/components/api-key-dialog.tsx create mode 100644 src/components/ui/dialog.tsx diff --git a/backend/app/routers/generate.py b/backend/app/routers/generate.py index a378307..8bac027 100644 --- a/backend/app/routers/generate.py +++ b/backend/app/routers/generate.py @@ -40,6 +40,7 @@ class ApiRequest(BaseModel): username: str repo: str instructions: str + api_key: str | None = None @router.post("") @@ -64,9 +65,17 @@ async def generate(request: Request, body: ApiRequest): # Check combined token count combined_content = f"{file_tree}\n{readme}" token_count = claude_service.count_tokens(combined_content) - if token_count > 50000: + + # Modified token limit check + if 50000 < token_count < 190000 and not body.api_key: + return { + "error": f"File tree and README combined exceeds token limit (50,000). Current size: {token_count} tokens. This GitHub repository is too large for my wallet, but you can continue by providing your own Anthropic API key.", + "token_count": token_count, + "requires_api_key": True + } + elif token_count > 200000: return { - "error": f"File tree and README combined exceeds token limit (50,000). Current size: {token_count} tokens. This GitHub repository is too large for my wallet, but if you still want the diagram and it's under 200k tokens, you can run GitDiagram locally with your own Anthropic API key." + "error": f"Repository is too large (>200k tokens) for analysis. Claude 3.5 Sonnet's max context length is 200k tokens. Current size: {token_count} tokens." } # Prepare system prompts with instructions if provided @@ -85,7 +94,8 @@ async def generate(request: Request, body: ApiRequest): "file_tree": file_tree, "readme": readme, "instructions": body.instructions - } + }, + api_key=body.api_key ) # Check for BAD_INSTRUCTIONS response diff --git a/backend/app/services/claude_service.py b/backend/app/services/claude_service.py index 1d22eb2..9cc39af 100644 --- a/backend/app/services/claude_service.py +++ b/backend/app/services/claude_service.py @@ -6,15 +6,16 @@ class ClaudeService: def __init__(self): - self.client = Anthropic() + self.default_client = Anthropic() - def call_claude_api(self, system_prompt: str, data: dict) -> str: + def call_claude_api(self, system_prompt: str, data: dict, api_key: str | None = None) -> str: """ Makes an API call to Claude and returns the response. Args: system_prompt (str): The instruction/system prompt data (dict): Dictionary of variables to format into the user message + api_key (str | None): Optional custom API key Returns: str: Claude's response text @@ -22,7 +23,10 @@ def call_claude_api(self, system_prompt: str, data: dict) -> str: # Create the user message with the data user_message = self._format_user_message(data) - message = self.client.messages.create( + # Use custom client if API key provided, otherwise use default + client = Anthropic(api_key=api_key) if api_key else self.default_client + + message = client.messages.create( model="claude-3-5-sonnet-latest", max_tokens=4096, temperature=0, @@ -73,7 +77,7 @@ def count_tokens(self, prompt: str) -> int: Returns: int: Number of input tokens """ - response = self.client.messages.count_tokens( + response = self.default_client.messages.count_tokens( model="claude-3-5-sonnet-latest", messages=[{ "role": "user", diff --git a/package.json b/package.json index 45a60a8..442baff 100644 --- a/package.json +++ b/package.json @@ -21,6 +21,7 @@ }, "dependencies": { "@neondatabase/serverless": "^0.10.4", + "@radix-ui/react-dialog": "^1.1.4", "@radix-ui/react-progress": "^1.1.1", "@radix-ui/react-slot": "^1.1.1", "@radix-ui/react-tooltip": "^1.1.6", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index cbcd0c4..cc80cea 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -11,6 +11,9 @@ importers: '@neondatabase/serverless': specifier: ^0.10.4 version: 0.10.4 + '@radix-ui/react-dialog': + specifier: ^1.1.4 + version: 1.1.4(@types/react-dom@18.3.5(@types/react@18.3.16))(@types/react@18.3.16)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) '@radix-ui/react-progress': specifier: ^1.1.1 version: 1.1.1(@types/react-dom@18.3.5(@types/react@18.3.16))(@types/react@18.3.16)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) @@ -728,6 +731,19 @@ packages: '@types/react': optional: true + '@radix-ui/react-dialog@1.1.4': + resolution: {integrity: sha512-Ur7EV1IwQGCyaAuyDRiOLA5JIUZxELJljF+MbM/2NC0BYwfuRrbpS30BiQBJrVruscgUkieKkqXYDOoByaxIoA==} + peerDependencies: + '@types/react': '*' + '@types/react-dom': '*' + react: ^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc + react-dom: ^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc + peerDependenciesMeta: + '@types/react': + optional: true + '@types/react-dom': + optional: true + '@radix-ui/react-dismissable-layer@1.1.3': resolution: {integrity: sha512-onrWn/72lQoEucDmJnr8uczSNTujT0vJnA/X5+3AkChVPowr8n1yvIKIabhWyMQeMvvmdpsvcyDqx3X1LEXCPg==} peerDependencies: @@ -741,6 +757,28 @@ packages: '@types/react-dom': optional: true + '@radix-ui/react-focus-guards@1.1.1': + resolution: {integrity: sha512-pSIwfrT1a6sIoDASCSpFwOasEwKTZWDw/iBdtnqKO7v6FeOzYJ7U53cPzYFVR3geGGXgVHaH+CdngrrAzqUGxg==} + peerDependencies: + '@types/react': '*' + react: ^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc + peerDependenciesMeta: + '@types/react': + optional: true + + '@radix-ui/react-focus-scope@1.1.1': + resolution: {integrity: sha512-01omzJAYRxXdG2/he/+xy+c8a8gCydoQ1yOxnWNcRhrrBW5W+RQJ22EK1SaO8tb3WoUsuEw7mJjBozPzihDFjA==} + peerDependencies: + '@types/react': '*' + '@types/react-dom': '*' + react: ^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc + react-dom: ^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc + peerDependenciesMeta: + '@types/react': + optional: true + '@types/react-dom': + optional: true + '@radix-ui/react-id@1.1.0': resolution: {integrity: sha512-EJUrI8yYh7WOjNOqpoJaf1jlFIH2LvtgAl+YcFqNCa+4hj64ZXmPkAKOFs/ukjz3byN6bdb/AVUqHkI8/uWWMA==} peerDependencies: @@ -1157,6 +1195,10 @@ packages: argparse@2.0.1: resolution: {integrity: sha512-8+9WqebbFzpX9OR+Wa6O29asIogeRMzcGtAINdpMHHyAg10f05aSFVBbcEqGf/PXw1EjAZ+q2/bEBg3DvurK3Q==} + aria-hidden@1.2.4: + resolution: {integrity: sha512-y+CcFFwelSXpLZk/7fMB2mUbGtX9lKycf1MWJ7CaTIERyitVlyQx6C+sxcROU2BAJ24OiZyK+8wj2i8AlBoS3A==} + engines: {node: '>=10'} + aria-query@5.3.2: resolution: {integrity: sha512-COROpnaoap1E2F000S62r6A60uHZnmlvomhfyT2DlTcrY1OrBKn2UhH7qn5wTC9zMvD0AY7csdPSNwKP+7WiQw==} engines: {node: '>= 0.4'} @@ -1543,6 +1585,9 @@ packages: resolution: {integrity: sha512-bwy0MGW55bG41VqxxypOsdSdGqLwXPI/focwgTYCFMbdUiBAxLg9CFzG08sz2aqzknwiX7Hkl0bQENjg8iLByw==} engines: {node: '>=8'} + detect-node-es@1.1.0: + resolution: {integrity: sha512-ypdmJU/TbBby2Dxibuv7ZLW3Bs1QEmM7nHjEANfohJLvE0XVujisn1qPJcZxg+qDucsr+bP6fLD1rPS3AhJ7EQ==} + didyoumean@1.2.2: resolution: {integrity: sha512-gxtyfqMg7GKyhQmb056K7M3xszy/myH8w+B4RT+QXBQsvAOdc3XymqDDPHx1BgPgsdAA5SIifona89YtRATDzw==} @@ -1917,6 +1962,10 @@ packages: resolution: {integrity: sha512-qxsEs+9A+u85HhllWJJFicJfPDhRmjzoYdl64aMWW9yRIJmSyxdn8IEkuIM530/7T+lv0TIHd8L6Q/ra0tEoeA==} engines: {node: '>= 0.4'} + get-nonce@1.0.1: + resolution: {integrity: sha512-FJhYRoDaiatfEkUK8HKlicmu/3SGFD51q3itKDGoSTysQJBnfOcxU5GxnhE1E6soB76MbT0MBtnKJuXyAx+96Q==} + engines: {node: '>=6'} + get-symbol-description@1.0.2: resolution: {integrity: sha512-g0QYk1dZBxGwk+Ngc+ltRH2IBp2f7zBkBMBJZCDerh6EhlhSR6+9irMCuT/09zD6qkarHUSn529sK/yL4S27mg==} engines: {node: '>= 0.4'} @@ -2623,6 +2672,36 @@ packages: react-is@16.13.1: resolution: {integrity: sha512-24e6ynE2H+OKt4kqsOvNd8kBpV65zoxbA4BVsEOB3ARVWQki/DHzaUoC5KuON/BiccDaCCTZBuOcfZs70kR8bQ==} + react-remove-scroll-bar@2.3.8: + resolution: {integrity: sha512-9r+yi9+mgU33AKcj6IbT9oRCO78WriSj6t/cF8DWBZJ9aOGPOTEDvdUDz1FwKim7QXWwmHqtdHnRJfhAxEG46Q==} + engines: {node: '>=10'} + peerDependencies: + '@types/react': '*' + react: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0 + peerDependenciesMeta: + '@types/react': + optional: true + + react-remove-scroll@2.6.2: + resolution: {integrity: sha512-KmONPx5fnlXYJQqC62Q+lwIeAk64ws/cUw6omIumRzMRPqgnYqhSSti99nbj0Ry13bv7dF+BKn7NB+OqkdZGTw==} + engines: {node: '>=10'} + peerDependencies: + '@types/react': '*' + react: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0 || ^19.0.0-rc + peerDependenciesMeta: + '@types/react': + optional: true + + react-style-singleton@2.2.3: + resolution: {integrity: sha512-b6jSvxvVnyptAiLjbkWLE/lOnR4lfTtDAl+eUC7RZy+QQWc6wRzIV2CE6xBuMmDxc2qIihtDCZD5NPOFl7fRBQ==} + engines: {node: '>=10'} + peerDependencies: + '@types/react': '*' + react: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0 || ^19.0.0-rc + peerDependenciesMeta: + '@types/react': + optional: true + react@18.3.1: resolution: {integrity: sha512-wS+hAgJShR0KhEvPJArfuPVN1+Hz1t0Y6n5jLrGQbkb4urgPE/0Rve+1kMB1v/oWgHgm4WIcV+i7F2pTVj+2iQ==} engines: {node: '>=0.10.0'} @@ -2932,6 +3011,26 @@ packages: uri-js@4.4.1: resolution: {integrity: sha512-7rKUyy33Q1yc98pQ1DAmLtwX109F7TIfWlW1Ydo8Wl1ii1SeHieeh0HHfPeL2fMXK6z0s8ecKs9frCuLJvndBg==} + use-callback-ref@1.3.3: + resolution: {integrity: sha512-jQL3lRnocaFtu3V00JToYz/4QkNWswxijDaCVNZRiRTO3HQDLsdu1ZtmIUvV4yPp+rvWm5j0y0TG/S61cuijTg==} + engines: {node: '>=10'} + peerDependencies: + '@types/react': '*' + react: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0 || ^19.0.0-rc + peerDependenciesMeta: + '@types/react': + optional: true + + use-sidecar@1.1.3: + resolution: {integrity: sha512-Fedw0aZvkhynoPYlA5WXrMCAMm+nSWdZt6lzJQ7Ok8S6Q+VsHmHpRWndVRJ8Be0ZbkfPc5LRYH+5XrzXcEeLRQ==} + engines: {node: '>=10'} + peerDependencies: + '@types/react': '*' + react: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0 || ^19.0.0-rc + peerDependenciesMeta: + '@types/react': + optional: true + util-deprecate@1.0.2: resolution: {integrity: sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==} @@ -3438,6 +3537,28 @@ snapshots: optionalDependencies: '@types/react': 18.3.16 + '@radix-ui/react-dialog@1.1.4(@types/react-dom@18.3.5(@types/react@18.3.16))(@types/react@18.3.16)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)': + dependencies: + '@radix-ui/primitive': 1.1.1 + '@radix-ui/react-compose-refs': 1.1.1(@types/react@18.3.16)(react@18.3.1) + '@radix-ui/react-context': 1.1.1(@types/react@18.3.16)(react@18.3.1) + '@radix-ui/react-dismissable-layer': 1.1.3(@types/react-dom@18.3.5(@types/react@18.3.16))(@types/react@18.3.16)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) + '@radix-ui/react-focus-guards': 1.1.1(@types/react@18.3.16)(react@18.3.1) + '@radix-ui/react-focus-scope': 1.1.1(@types/react-dom@18.3.5(@types/react@18.3.16))(@types/react@18.3.16)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) + '@radix-ui/react-id': 1.1.0(@types/react@18.3.16)(react@18.3.1) + '@radix-ui/react-portal': 1.1.3(@types/react-dom@18.3.5(@types/react@18.3.16))(@types/react@18.3.16)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) + '@radix-ui/react-presence': 1.1.2(@types/react-dom@18.3.5(@types/react@18.3.16))(@types/react@18.3.16)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) + '@radix-ui/react-primitive': 2.0.1(@types/react-dom@18.3.5(@types/react@18.3.16))(@types/react@18.3.16)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) + '@radix-ui/react-slot': 1.1.1(@types/react@18.3.16)(react@18.3.1) + '@radix-ui/react-use-controllable-state': 1.1.0(@types/react@18.3.16)(react@18.3.1) + aria-hidden: 1.2.4 + react: 18.3.1 + react-dom: 18.3.1(react@18.3.1) + react-remove-scroll: 2.6.2(@types/react@18.3.16)(react@18.3.1) + optionalDependencies: + '@types/react': 18.3.16 + '@types/react-dom': 18.3.5(@types/react@18.3.16) + '@radix-ui/react-dismissable-layer@1.1.3(@types/react-dom@18.3.5(@types/react@18.3.16))(@types/react@18.3.16)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)': dependencies: '@radix-ui/primitive': 1.1.1 @@ -3451,6 +3572,23 @@ snapshots: '@types/react': 18.3.16 '@types/react-dom': 18.3.5(@types/react@18.3.16) + '@radix-ui/react-focus-guards@1.1.1(@types/react@18.3.16)(react@18.3.1)': + dependencies: + react: 18.3.1 + optionalDependencies: + '@types/react': 18.3.16 + + '@radix-ui/react-focus-scope@1.1.1(@types/react-dom@18.3.5(@types/react@18.3.16))(@types/react@18.3.16)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)': + dependencies: + '@radix-ui/react-compose-refs': 1.1.1(@types/react@18.3.16)(react@18.3.1) + '@radix-ui/react-primitive': 2.0.1(@types/react-dom@18.3.5(@types/react@18.3.16))(@types/react@18.3.16)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) + '@radix-ui/react-use-callback-ref': 1.1.0(@types/react@18.3.16)(react@18.3.1) + react: 18.3.1 + react-dom: 18.3.1(react@18.3.1) + optionalDependencies: + '@types/react': 18.3.16 + '@types/react-dom': 18.3.5(@types/react@18.3.16) + '@radix-ui/react-id@1.1.0(@types/react@18.3.16)(react@18.3.1)': dependencies: '@radix-ui/react-use-layout-effect': 1.1.0(@types/react@18.3.16)(react@18.3.1) @@ -3883,6 +4021,10 @@ snapshots: argparse@2.0.1: {} + aria-hidden@1.2.4: + dependencies: + tslib: 2.8.1 + aria-query@5.3.2: {} array-buffer-byte-length@1.0.1: @@ -4328,6 +4470,8 @@ snapshots: detect-libc@2.0.3: optional: true + detect-node-es@1.1.0: {} + didyoumean@1.2.2: {} dlv@1.1.3: {} @@ -4825,6 +4969,8 @@ snapshots: hasown: 2.0.2 math-intrinsics: 1.0.0 + get-nonce@1.0.1: {} + get-symbol-description@1.0.2: dependencies: call-bind: 1.0.8 @@ -5471,6 +5617,33 @@ snapshots: react-is@16.13.1: {} + react-remove-scroll-bar@2.3.8(@types/react@18.3.16)(react@18.3.1): + dependencies: + react: 18.3.1 + react-style-singleton: 2.2.3(@types/react@18.3.16)(react@18.3.1) + tslib: 2.8.1 + optionalDependencies: + '@types/react': 18.3.16 + + react-remove-scroll@2.6.2(@types/react@18.3.16)(react@18.3.1): + dependencies: + react: 18.3.1 + react-remove-scroll-bar: 2.3.8(@types/react@18.3.16)(react@18.3.1) + react-style-singleton: 2.2.3(@types/react@18.3.16)(react@18.3.1) + tslib: 2.8.1 + use-callback-ref: 1.3.3(@types/react@18.3.16)(react@18.3.1) + use-sidecar: 1.1.3(@types/react@18.3.16)(react@18.3.1) + optionalDependencies: + '@types/react': 18.3.16 + + react-style-singleton@2.2.3(@types/react@18.3.16)(react@18.3.1): + dependencies: + get-nonce: 1.0.1 + react: 18.3.1 + tslib: 2.8.1 + optionalDependencies: + '@types/react': 18.3.16 + react@18.3.1: dependencies: loose-envify: 1.4.0 @@ -5879,6 +6052,21 @@ snapshots: dependencies: punycode: 2.3.1 + use-callback-ref@1.3.3(@types/react@18.3.16)(react@18.3.1): + dependencies: + react: 18.3.1 + tslib: 2.8.1 + optionalDependencies: + '@types/react': 18.3.16 + + use-sidecar@1.1.3(@types/react@18.3.16)(react@18.3.1): + dependencies: + detect-node-es: 1.1.0 + react: 18.3.1 + tslib: 2.8.1 + optionalDependencies: + '@types/react': 18.3.16 + util-deprecate@1.0.2: {} uuid@9.0.1: {} diff --git a/src/app/[username]/[repo]/page.tsx b/src/app/[username]/[repo]/page.tsx index fda1561..f8091ad 100644 --- a/src/app/[username]/[repo]/page.tsx +++ b/src/app/[username]/[repo]/page.tsx @@ -5,6 +5,9 @@ import MainCard from "~/components/main-card"; import Loading from "~/components/loading"; import MermaidChart from "~/components/mermaid-diagram"; import { useDiagram } from "~/hooks/useDiagram"; +import { ApiKeyDialog } from "~/components/api-key-dialog"; +import { Button } from "~/components/ui/button"; +import { ApiKeyButton } from "~/components/api-key-button"; export default function Repo() { const params = useParams<{ username: string; repo: string }>(); @@ -15,9 +18,14 @@ export default function Repo() { lastGenerated, cost, isRegenerating, + showApiKeyDialog, + tokenCount, handleModify, handleRegenerate, handleCopy, + handleApiKeySubmit, + handleCloseApiKeyDialog, + handleOpenApiKeyDialog, } = useDiagram(params.username, params.repo); return ( @@ -41,12 +49,20 @@ export default function Repo() { ) : error ? (
-

{error}

+

+ {error} +

{error.includes("Rate limit") && (

Rate limits: 1 request per minute, 5 requests per day

)} + {error.includes("token limit") && ( +
+ +

Your key will not be stored

+
+ )}
) : (
@@ -54,6 +70,13 @@ export default function Repo() {
)} + + ); } diff --git a/src/components/api-key-button.tsx b/src/components/api-key-button.tsx new file mode 100644 index 0000000..863be4f --- /dev/null +++ b/src/components/api-key-button.tsx @@ -0,0 +1,18 @@ +import { Key } from "lucide-react"; +import { Button } from "./ui/button"; + +interface ApiKeyButtonProps { + onClick: () => void; +} + +export function ApiKeyButton({ onClick }: ApiKeyButtonProps) { + return ( + + ); +} diff --git a/src/components/api-key-dialog.tsx b/src/components/api-key-dialog.tsx new file mode 100644 index 0000000..72cc9b1 --- /dev/null +++ b/src/components/api-key-dialog.tsx @@ -0,0 +1,69 @@ +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "./ui/dialog"; +import { Input } from "./ui/input"; +import { Button } from "./ui/button"; +import { useState } from "react"; + +interface ApiKeyDialogProps { + isOpen: boolean; + onClose: () => void; + onSubmit: (apiKey: string) => void; + tokenCount: number; +} + +export function ApiKeyDialog({ + isOpen, + onClose, + onSubmit, + tokenCount, +}: ApiKeyDialogProps) { + const [apiKey, setApiKey] = useState(""); + + const handleSubmit = (e: React.FormEvent) => { + e.preventDefault(); + onSubmit(apiKey); + setApiKey(""); + }; + + return ( + + + + + Enter Anthropic API Key + + +
+
+ This repository requires {tokenCount.toLocaleString()} tokens to + analyze. Enter your Anthropic API key to continue. The key will not + be stored. +
+ setApiKey(e.target.value)} + className="flex-1 rounded-md border-[3px] border-black px-3 py-2 text-base font-bold shadow-[4px_4px_0_0_#000000] placeholder:text-base placeholder:font-normal placeholder:text-gray-700" + required + /> +
+ + +
+
+ + + ); +} diff --git a/src/components/ui/dialog.tsx b/src/components/ui/dialog.tsx new file mode 100644 index 0000000..7ce4535 --- /dev/null +++ b/src/components/ui/dialog.tsx @@ -0,0 +1,122 @@ +"use client"; + +import * as React from "react"; +import * as DialogPrimitive from "@radix-ui/react-dialog"; +import { X } from "lucide-react"; + +import { cn } from "~/lib/utils"; + +const Dialog = DialogPrimitive.Root; + +const DialogTrigger = DialogPrimitive.Trigger; + +const DialogPortal = DialogPrimitive.Portal; + +const DialogClose = DialogPrimitive.Close; + +const DialogOverlay = React.forwardRef< + React.ElementRef, + React.ComponentPropsWithoutRef +>(({ className, ...props }, ref) => ( + +)); +DialogOverlay.displayName = DialogPrimitive.Overlay.displayName; + +const DialogContent = React.forwardRef< + React.ElementRef, + React.ComponentPropsWithoutRef +>(({ className, children, ...props }, ref) => ( + + + + {children} + + + Close + + + +)); +DialogContent.displayName = DialogPrimitive.Content.displayName; + +const DialogHeader = ({ + className, + ...props +}: React.HTMLAttributes) => ( +
+); +DialogHeader.displayName = "DialogHeader"; + +const DialogFooter = ({ + className, + ...props +}: React.HTMLAttributes) => ( +
+); +DialogFooter.displayName = "DialogFooter"; + +const DialogTitle = React.forwardRef< + React.ElementRef, + React.ComponentPropsWithoutRef +>(({ className, ...props }, ref) => ( + +)); +DialogTitle.displayName = DialogPrimitive.Title.displayName; + +const DialogDescription = React.forwardRef< + React.ElementRef, + React.ComponentPropsWithoutRef +>(({ className, ...props }, ref) => ( + +)); +DialogDescription.displayName = DialogPrimitive.Description.displayName; + +export { + Dialog, + DialogPortal, + DialogOverlay, + DialogClose, + DialogTrigger, + DialogContent, + DialogHeader, + DialogFooter, + DialogTitle, + DialogDescription, +}; diff --git a/src/hooks/useDiagram.ts b/src/hooks/useDiagram.ts index e903134..85cf6e2 100644 --- a/src/hooks/useDiagram.ts +++ b/src/hooks/useDiagram.ts @@ -15,6 +15,8 @@ export function useDiagram(username: string, repo: string) { const [isRegenerating, setIsRegenerating] = useState(false); const [lastGenerated, setLastGenerated] = useState(); const [cost, setCost] = useState(""); + const [showApiKeyDialog, setShowApiKeyDialog] = useState(false); + const [tokenCount, setTokenCount] = useState(0); const getDiagram = useCallback(async () => { setLoading(true); @@ -42,6 +44,9 @@ export function useDiagram(username: string, repo: string) { if (result.error) { console.error("Diagram generation failed:", result.error); + if (result.requires_api_key) { + setTokenCount(result.token_count ?? 0); + } setError(result.error); } else if (result.diagram) { setDiagram(result.diagram); @@ -143,6 +148,36 @@ export function useDiagram(username: string, repo: string) { } }; + const handleApiKeySubmit = async (apiKey: string) => { + setShowApiKeyDialog(false); + setLoading(true); + setError(""); + + try { + const result = await generateAndCacheDiagram(username, repo, "", apiKey); + if (result.error) { + setError(result.error); + } else if (result.diagram) { + setDiagram(result.diagram); + const date = await getLastGeneratedDate(username, repo); + setLastGenerated(date ?? undefined); + } + } catch (error) { + console.error("Error generating with API key:", error); + setError("Failed to generate diagram with provided API key."); + } finally { + setLoading(false); + } + }; + + const handleCloseApiKeyDialog = () => { + setShowApiKeyDialog(false); + }; + + const handleOpenApiKeyDialog = () => { + setShowApiKeyDialog(true); + }; + return { diagram, error, @@ -153,5 +188,10 @@ export function useDiagram(username: string, repo: string) { handleModify, handleRegenerate, handleCopy, + showApiKeyDialog, + tokenCount, + handleApiKeySubmit, + handleCloseApiKeyDialog, + handleOpenApiKeyDialog, }; } diff --git a/src/lib/fetch-backend.ts b/src/lib/fetch-backend.ts index 5b8a898..20abcb3 100644 --- a/src/lib/fetch-backend.ts +++ b/src/lib/fetch-backend.ts @@ -8,6 +8,8 @@ interface GenerateApiResponse { error?: string; diagram?: string; explanation?: string; + token_count?: number; + requires_api_key?: boolean; } interface ModifyApiResponse { @@ -24,6 +26,7 @@ export async function generateAndCacheDiagram( username: string, repo: string, instructions?: string, + api_key?: string, ): Promise { try { const baseUrl = @@ -39,6 +42,7 @@ export async function generateAndCacheDiagram( username, repo, instructions: instructions ?? "", + api_key: api_key, }), }); @@ -49,7 +53,7 @@ export async function generateAndCacheDiagram( const data = (await response.json()) as GenerateApiResponse; if (data.error) { - return { error: data.error }; + return data; // pass the whole thing for multiple data fields } // Call the server action to cache the diagram