Fix #16: Add support for private repository analysis

ahmedkhaleel2004 · ahmedkhaleel2004 · commit 702d2519eb5d · 2025-02-12T19:41:57.000-05:00
diff --git a/backend/app/routers/generate.py b/backend/app/routers/generate.py
@@ -21,29 +21,32 @@
 router = APIRouter(prefix="/generate", tags=["Claude"])
 
 # Initialize services
-github_service = GitHubService()
 # claude_service = ClaudeService()
 o3_service = OpenRouterO3Service()
 
 
-# cache github data for 5 minutes to avoid double API calls from cost and generate
+# cache github data to avoid double API calls from cost and generate
 @lru_cache(maxsize=100)
-def get_cached_github_data(username: str, repo: str):
-    default_branch = github_service.get_default_branch(username, repo)
+def get_cached_github_data(username: str, repo: str, github_pat: str | None = None):
+    # Create a new service instance for each call with the appropriate PAT
+    current_github_service = GitHubService(pat=github_pat)
+
+    default_branch = current_github_service.get_default_branch(username, repo)
     if not default_branch:
         default_branch = "main"  # fallback value
 
-    file_tree = github_service.get_github_file_paths_as_list(username, repo)
-    readme = github_service.get_github_readme(username, repo)
+    file_tree = current_github_service.get_github_file_paths_as_list(username, repo)
+    readme = current_github_service.get_github_readme(username, repo)
 
     return {"default_branch": default_branch, "file_tree": file_tree, "readme": readme}
 
 
 class ApiRequest(BaseModel):
     username: str
     repo: str
-    instructions: str
+    instructions: str = ""
     api_key: str | None = None
+    github_pat: str | None = None
 
 
 @router.post("")
@@ -63,8 +66,8 @@ async def generate(request: Request, body: ApiRequest):
         ]:
             return {"error": "Example repos cannot be regenerated"}
 
-        # Get cached github data
-        github_data = get_cached_github_data(body.username, body.repo)
+        # Get cached github data with PAT if provided
+        github_data = get_cached_github_data(body.username, body.repo, body.github_pat)
 
         # Get default branch first
         default_branch = github_data["default_branch"]
@@ -205,7 +208,7 @@ async def generate(request: Request, body: ApiRequest):
 async def get_generation_cost(request: Request, body: ApiRequest):
     try:
         # Get file tree and README content
-        github_data = get_cached_github_data(body.username, body.repo)
+        github_data = get_cached_github_data(body.username, body.repo, body.github_pat)
         file_tree = github_data["file_tree"]
         readme = github_data["readme"]
 
diff --git a/backend/app/services/github_service.py b/backend/app/services/github_service.py
@@ -9,14 +9,14 @@
 
 
 class GitHubService:
-    def __init__(self):
+    def __init__(self, pat: str | None = None):
         # Try app authentication first
         self.client_id = os.getenv("GITHUB_CLIENT_ID")
         self.private_key = os.getenv("GITHUB_PRIVATE_KEY")
         self.installation_id = os.getenv("GITHUB_INSTALLATION_ID")
 
-        # Fallback to PAT if app credentials not found
-        self.github_token = os.getenv("GITHUB_PAT")
+        # Use provided PAT if available, otherwise fallback to env PAT
+        self.github_token = pat or os.getenv("GITHUB_PAT")
 
         # If no credentials are provided, warn about rate limits
         if (
diff --git a/src/components/header.tsx b/src/components/header.tsx
@@ -1,19 +1,34 @@
-import React from "react";
+"use client";
+
+import React, { useState, useEffect } from "react";
 import Link from "next/link";
 import { FaGithub } from "react-icons/fa";
 import { getStarCount } from "~/app/_actions/github";
+import { PrivateReposDialog } from "./private-repos-dialog";
+
+export function Header() {
+  const [isPrivateReposDialogOpen, setIsPrivateReposDialogOpen] =
+    useState(false);
+  const [starCount, setStarCount] = useState<number | null>(null);
 
-export async function Header() {
-  const starCount = await getStarCount();
+  useEffect(() => {
+    void getStarCount().then(setStarCount);
+  }, []);
 
   const formatStarCount = (count: number | null) => {
-    if (!count) return "0";
+    if (!count) return "2.0k"; // Default to 2.0k if count is null (it can only go up from here)
     if (count >= 1000) {
       return `${(count / 1000).toFixed(1)}k`;
     }
     return count.toString();
   };
 
+  const handlePrivateReposSubmit = (pat: string) => {
+    // Store the PAT in localStorage
+    localStorage.setItem("github_pat", pat);
+    setIsPrivateReposDialogOpen(false);
+  };
+
   return (
     <header className="border-b-[3px] border-black">
       <div className="mx-auto flex h-16 max-w-4xl items-center justify-between px-8">
@@ -34,6 +49,12 @@ export async function Header() {
           >
             API
           </Link>
+          <span
+            onClick={() => setIsPrivateReposDialogOpen(true)}
+            className="cursor-pointer text-sm font-medium text-black transition-transform hover:translate-y-[-2px] hover:text-purple-600"
+          >
+            Private Repos
+          </span>
           <Link
             href="https://github.com/ahmedkhaleel2004/gitdiagram"
             className="flex items-center gap-2 text-sm font-medium text-black transition-transform hover:translate-y-[-2px] hover:text-purple-600"
@@ -46,6 +67,12 @@ export async function Header() {
             {formatStarCount(starCount)}
           </span>
         </nav>
+
+        <PrivateReposDialog
+          isOpen={isPrivateReposDialogOpen}
+          onClose={() => setIsPrivateReposDialogOpen(false)}
+          onSubmit={handlePrivateReposSubmit}
+        />
       </div>
     </header>
   );
diff --git a/src/components/private-repos-dialog.tsx b/src/components/private-repos-dialog.tsx
@@ -0,0 +1,100 @@
+import { Dialog, DialogContent, DialogHeader, DialogTitle } from "./ui/dialog";
+import { Input } from "./ui/input";
+import { Button } from "./ui/button";
+import { useState } from "react";
+import Link from "next/link";
+
+interface PrivateReposDialogProps {
+  isOpen: boolean;
+  onClose: () => void;
+  onSubmit: (pat: string) => void;
+}
+
+export function PrivateReposDialog({
+  isOpen,
+  onClose,
+  onSubmit,
+}: PrivateReposDialogProps) {
+  const [pat, setPat] = useState(() => {
+    // Initialize from localStorage if available
+    if (typeof window !== "undefined") {
+      return localStorage.getItem("github_pat") ?? "";
+    }
+    return "";
+  });
+
+  const handleSubmit = (e: React.FormEvent) => {
+    e.preventDefault();
+    onSubmit(pat);
+    setPat("");
+  };
+
+  return (
+    <Dialog open={isOpen} onOpenChange={onClose}>
+      <DialogContent className="border-[3px] border-black bg-purple-200 p-6 shadow-[8px_8px_0_0_#000000] sm:max-w-md">
+        <DialogHeader>
+          <DialogTitle className="text-xl font-bold text-black">
+            Enter GitHub Personal Access Token
+          </DialogTitle>
+        </DialogHeader>
+        <form onSubmit={handleSubmit} className="space-y-4">
+          <div className="text-sm">
+            To enable private repositories, you&apos;ll need to provide a GitHub
+            Personal Access Token with repo scope. The token will be stored
+            locally in your browser. Find out how{" "}
+            <Link
+              href="https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens"
+              className="text-purple-600 transition-colors duration-200 hover:text-purple-500"
+            >
+              here
+            </Link>
+            .
+          </div>
+          <details className="group text-sm [&>summary:focus-visible]:outline-none">
+            <summary className="cursor-pointer font-medium text-purple-700 hover:text-purple-600">
+              Data storage disclaimer
+            </summary>
+            <div className="animate-accordion-down mt-2 space-y-2 overflow-hidden pl-2">
+              <p>
+                Take note that the diagram data will be stored in my database
+                (not that I would use it for anything anyways). You can also
+                self-host this app by following the instructions in the{" "}
+                <Link
+                  href="https://github.com/ahmedkhaleel2004/gitdiagram"
+                  className="text-purple-600 transition-colors duration-200 hover:text-purple-500"
+                >
+                  README
+                </Link>
+                .
+              </p>
+            </div>
+          </details>
+          <Input
+            type="password"
+            placeholder="ghp_..."
+            value={pat}
+            onChange={(e) => setPat(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
+          />
+          <div className="flex justify-end gap-3">
+            <Button
+              type="button"
+              onClick={onClose}
+              className="border-[3px] border-black bg-gray-200 px-4 py-2 text-black shadow-[4px_4px_0_0_#000000] transition-transform hover:-translate-x-0.5 hover:-translate-y-0.5 hover:bg-gray-300"
+            >
+              Cancel
+            </Button>
+            <Button
+              type="submit"
+              disabled={!pat.startsWith("ghp_")}
+              className="border-[3px] border-black bg-purple-400 px-4 py-2 text-black shadow-[4px_4px_0_0_#000000] transition-transform hover:-translate-x-0.5 hover:-translate-y-0.5 hover:bg-purple-300 disabled:opacity-50"
+            >
+              Save Token
+            </Button>
+          </div>
+        </form>
+      </DialogContent>
+    </Dialog>
+  );
+}
diff --git a/src/hooks/useDiagram.ts b/src/hooks/useDiagram.ts
@@ -25,13 +25,19 @@ export function useDiagram(username: string, repo: string) {
     setIsRegenerating(true);
     try {
       const cached = await getCachedDiagram(username, repo);
+      const github_pat = localStorage.getItem("github_pat");
 
       if (cached) {
         setDiagram(cached);
         const date = await getLastGeneratedDate(username, repo);
         setLastGenerated(date ?? undefined);
       } else {
-        const costEstimate = await getCostOfGeneration(username, repo, ""); // empty instructions so lru cache is used
+        const costEstimate = await getCostOfGeneration(
+          username,
+          repo,
+          "",
+          github_pat ?? undefined,
+        ); // empty instructions so lru cache is used
 
         if (costEstimate.error) {
           console.error("Cost estimation failed:", costEstimate.error);
@@ -40,7 +46,11 @@ export function useDiagram(username: string, repo: string) {
 
         setCost(costEstimate.cost ?? "");
 
-        const result = await generateAndCacheDiagram(username, repo);
+        const result = await generateAndCacheDiagram(
+          username,
+          repo,
+          github_pat ?? undefined,
+        );
 
         if (result.error) {
           console.error("Diagram generation failed:", result.error);
@@ -110,6 +120,7 @@ export function useDiagram(username: string, repo: string) {
     setCost("");
     setIsRegenerating(true);
     try {
+      const github_pat = localStorage.getItem("github_pat");
       const costEstimate = await getCostOfGeneration(username, repo, "");
 
       if (costEstimate.error) {
@@ -122,6 +133,7 @@ export function useDiagram(username: string, repo: string) {
       const result = await generateAndCacheDiagram(
         username,
         repo,
+        github_pat ?? undefined,
         instructions,
       );
       if (result.error) {
@@ -198,9 +210,15 @@ export function useDiagram(username: string, repo: string) {
     setShowApiKeyDialog(false);
     setLoading(true);
     setError("");
-
+    const github_pat = localStorage.getItem("github_pat");
     try {
-      const result = await generateAndCacheDiagram(username, repo, "", apiKey);
+      const result = await generateAndCacheDiagram(
+        username,
+        repo,
+        github_pat ?? undefined,
+        "",
+        apiKey,
+      );
       if (result.error) {
         setError(result.error);
       } else if (result.diagram) {
diff --git a/src/lib/fetch-backend.ts b/src/lib/fetch-backend.ts
@@ -25,6 +25,7 @@ interface CostApiResponse {
 export async function generateAndCacheDiagram(
   username: string,
   repo: string,
+  github_pat?: string,
   instructions?: string,
   api_key?: string,
 ): Promise<GenerateApiResponse> {
@@ -43,6 +44,7 @@ export async function generateAndCacheDiagram(
         repo,
         instructions: instructions ?? "",
         api_key: api_key,
+        github_pat: github_pat,
       }),
     });
 
@@ -130,6 +132,7 @@ export async function getCostOfGeneration(
   username: string,
   repo: string,
   instructions: string,
+  github_pat?: string,
 ): Promise<CostApiResponse> {
   try {
     const baseUrl =
@@ -144,6 +147,7 @@ export async function getCostOfGeneration(
       body: JSON.stringify({
         username,
         repo,
+        github_pat: github_pat,
         instructions: instructions ?? "",
       }),
     });
