Add branch and Azure repository support for ingestion

api/rag.py

@@ -343,7 +343,7 @@ def _validate_and_filter_embeddings(self, documents: List) -> List:
 
     def prepare_retriever(self, repo_url_or_path: str, type: str = "github", access_token: str = None,
                       excluded_dirs: List[str] = None, excluded_files: List[str] = None,
-                      included_dirs: List[str] = None, included_files: List[str] = None):
+                      included_dirs: List[str] = None, included_files: List[str] = None, branch: str = None):
         """
         Prepare the retriever for a repository.
         Will load database from local storage if available.
@@ -355,6 +355,7 @@ def prepare_retriever(self, repo_url_or_path: str, type: str = "github", access_
             excluded_files: Optional list of file patterns to exclude from processing
             included_dirs: Optional list of directories to include exclusively
             included_files: Optional list of file patterns to include exclusively
+            branch: Optional branch to use for cloning and indexing
         """
         self.initialize_db_manager()
         self.repo_url_or_path = repo_url_or_path
@@ -366,7 +367,8 @@ def prepare_retriever(self, repo_url_or_path: str, type: str = "github", access_
             excluded_dirs=excluded_dirs,
             excluded_files=excluded_files,
             included_dirs=included_dirs,
-            included_files=included_files
+            included_files=included_files,
+            branch=branch
         )
         logger.info(f"Loaded {len(self.transformed_docs)} documents for retrieval")
 

api/websocket_wiki.py

@@ -38,6 +38,7 @@ class ChatCompletionRequest(BaseModel):
     filePath: Optional[str] = Field(None, description="Optional path to a file in the repository to include in the prompt")
     token: Optional[str] = Field(None, description="Personal access token for private repositories")
     type: Optional[str] = Field("github", description="Type of repository (e.g., 'github', 'gitlab', 'bitbucket')")
+    branch: Optional[str] = Field(None, description="Branch name to use for retrieval and file content")
 
     # model parameters
     provider: str = Field("google", description="Model provider (google, openai, openrouter, ollama, azure)")
@@ -95,7 +96,7 @@ async def handle_websocket_chat(websocket: WebSocket):
                 included_files = [unquote(file_pattern) for file_pattern in request.included_files.split('\n') if file_pattern.strip()]
                 logger.info(f"Using custom included files: {included_files}")
 
-            request_rag.prepare_retriever(request.repo_url, request.type, request.token, excluded_dirs, excluded_files, included_dirs, included_files)
+            request_rag.prepare_retriever(request.repo_url, request.type, request.token, excluded_dirs, excluded_files, included_dirs, included_files, branch=request.branch)
             logger.info(f"Retriever prepared for {request.repo_url}")
         except ValueError as e:
             if "No valid documents with embeddings found" in str(e):
@@ -391,8 +392,8 @@ async def handle_websocket_chat(websocket: WebSocket):
         file_content = ""
         if request.filePath:
             try:
-                file_content = get_file_content(request.repo_url, request.filePath, request.type, request.token)
-                logger.info(f"Successfully retrieved content for file: {request.filePath}")
+                file_content = get_file_content(request.repo_url, request.filePath, request.type, request.token, request.branch)
+                logger.info(f"Successfully retrieved content for file: {request.filePath} (branch: {request.branch or 'default'})")
             except Exception as e:
                 logger.error(f"Error retrieving file content: {str(e)}")
                 # Continue without file content if there's an error

src/app/[owner]/[repo]/page.tsx

@@ -106,7 +106,8 @@ const addTokensToRequestBody = (
   excludedDirs?: string,
   excludedFiles?: string,
   includedDirs?: string,
-  includedFiles?: string
+  includedFiles?: string,
+  branch?: string
 ): void => {
   if (token !== '') {
     requestBody.token = token;
@@ -134,7 +135,9 @@ const addTokensToRequestBody = (
   if (includedFiles) {
     requestBody.included_files = includedFiles;
   }
-
+  if (branch) {
+    requestBody.branch = branch;
+  }
 };
 
 const createGithubHeaders = (githubToken: string): HeadersInit => {
@@ -199,6 +202,7 @@ export default function RepoWikiPage() {
       : repoUrl?.includes('github.com')
         ? 'github'
         : searchParams.get('type') || 'github';
+  const branch = searchParams.get('branch') || '';
 
   // Import language context for translations
   const { messages } = useLanguage();
@@ -210,8 +214,9 @@ export default function RepoWikiPage() {
     type: repoType,
     token: token || null,
     localPath: localPath || null,
-    repoUrl: repoUrl || null
-  }), [owner, repo, repoType, localPath, repoUrl, token]);
+    repoUrl: repoUrl || null,
+    branch: branch || null,
+  }), [owner, repo, repoType, localPath, repoUrl, token, branch]);
 
   // State variables
   const [isLoading, setIsLoading] = useState(true);
@@ -511,7 +516,7 @@ Remember:
         };
 
         // Add tokens if available
-        addTokensToRequestBody(requestBody, currentToken, effectiveRepoInfo.type, selectedProviderState, selectedModelState, isCustomSelectedModelState, customSelectedModelState, language, modelExcludedDirs, modelExcludedFiles, modelIncludedDirs, modelIncludedFiles);
+        addTokensToRequestBody(requestBody, currentToken, effectiveRepoInfo.type, selectedProviderState, selectedModelState, isCustomSelectedModelState, customSelectedModelState, language, modelExcludedDirs, modelExcludedFiles, modelIncludedDirs, modelIncludedFiles, branch);
 
         // Use WebSocket for communication
         let content = '';
@@ -808,7 +813,7 @@ IMPORTANT:
       };
 
       // Add tokens if available
-      addTokensToRequestBody(requestBody, currentToken, effectiveRepoInfo.type, selectedProviderState, selectedModelState, isCustomSelectedModelState, customSelectedModelState, language, modelExcludedDirs, modelExcludedFiles, modelIncludedDirs, modelIncludedFiles);
+      addTokensToRequestBody(requestBody, currentToken, effectiveRepoInfo.type, selectedProviderState, selectedModelState, isCustomSelectedModelState, customSelectedModelState, language, modelExcludedDirs, modelExcludedFiles, modelIncludedDirs, modelIncludedFiles, branch);
 
       // Use WebSocket for communication
       let responseText = '';
@@ -1211,22 +1216,31 @@ IMPORTANT:
         };
 
         const githubApiBaseUrl = getGithubApiUrl(effectiveRepoInfo.repoUrl);
-        // First, try to get the default branch from the repository info
-        let defaultBranchLocal = null;
-        try {
-          const repoInfoResponse = await fetch(`${githubApiBaseUrl}/repos/${owner}/${repo}`, {
-            headers: createGithubHeaders(currentToken)
-          });
-
-          if (repoInfoResponse.ok) {
-            const repoData = await repoInfoResponse.json();
-            defaultBranchLocal = repoData.default_branch;
-            console.log(`Found default branch: ${defaultBranchLocal}`);
-            // Store the default branch in state
-            setDefaultBranch(defaultBranchLocal || 'main');
+        // First, try to get the default branch from effectiveRepoInfo or the repository info
+        let defaultBranchLocal = effectiveRepoInfo?.branch && `${effectiveRepoInfo.branch}`.trim() !== ''
+          ? `${effectiveRepoInfo.branch}`.trim()
+          : null;
+
+        // If branch not provided, fetch repository info to determine default branch
+        if (!defaultBranchLocal) {
+          try {
+            const repoInfoResponse = await fetch(`${githubApiBaseUrl}/repos/${owner}/${repo}`, {
+              headers: createGithubHeaders(currentToken)
+            });
+
+            if (repoInfoResponse.ok) {
+              const repoData = await repoInfoResponse.json();
+              defaultBranchLocal = repoData.default_branch;
+              console.log(`Found default branch: ${defaultBranchLocal}`);
+              // Store the default branch in state
+              setDefaultBranch(defaultBranchLocal || 'main');
+            }
+          } catch (err) {
+            console.warn('Could not fetch repository info for default branch:', err);
           }
-        } catch (err) {
-          console.warn('Could not fetch repository info for default branch:', err);
+        } else {
+          // If we already have a branch from effectiveRepoInfo, store it
+          setDefaultBranch(defaultBranchLocal);
         }
 
         // Create list of branches to try, prioritizing the actual default branch
@@ -2243,7 +2257,7 @@ IMPORTANT:
         onApply={confirmRefresh}
         showWikiType={true}
         showTokenInput={effectiveRepoInfo.type !== 'local' && !currentToken} // Show token input if not local and no current token
-        repositoryType={effectiveRepoInfo.type as 'github' | 'gitlab' | 'bitbucket'}
+        repositoryType={effectiveRepoInfo.type as 'github' | 'gitlab' | 'bitbucket' | 'azure'}
         authRequired={authRequired}
         authCode={authCode}
         setAuthCode={setAuthCode}

src/app/page.tsx

@@ -76,6 +76,7 @@ export default function Home() {
   };
 
   const [repositoryInput, setRepositoryInput] = useState('https://github.com/AsyncFuncAI/deepwiki-open');
+  const [branch, setBranch] = useState<string>('');
 
   const REPO_CONFIG_CACHE_KEY = 'deepwikiRepoConfigCache';
 
@@ -98,6 +99,7 @@ export default function Home() {
           setExcludedFiles(config.excludedFiles || '');
           setIncludedDirs(config.includedDirs || '');
           setIncludedFiles(config.includedFiles || '');
+          setBranch(config.branch || '');
         }
       }
     } catch (error) {
@@ -134,7 +136,7 @@ export default function Home() {
   const [excludedFiles, setExcludedFiles] = useState('');
   const [includedDirs, setIncludedDirs] = useState('');
   const [includedFiles, setIncludedFiles] = useState('');
-  const [selectedPlatform, setSelectedPlatform] = useState<'github' | 'gitlab' | 'bitbucket'>('github');
+  const [selectedPlatform, setSelectedPlatform] = useState<'github' | 'gitlab' | 'bitbucket' | 'azure'>('github');
   const [accessToken, setAccessToken] = useState('');
   const [error, setError] = useState<string | null>(null);
   const [isSubmitting, setIsSubmitting] = useState(false);
@@ -212,6 +214,8 @@ export default function Home() {
         type = 'gitlab';
       } else if (domain?.includes('bitbucket.org') || domain?.includes('bitbucket.')) {
         type = 'bitbucket';
+      } else if (domain?.includes('azure.com') || domain?.includes('dev.azure.com')) {
+        type = 'azure';
       } else {
         type = 'web'; // fallback for other git hosting services
       }
@@ -221,6 +225,13 @@ export default function Home() {
       if (parts.length >= 2) {
         repo = parts[parts.length - 1] || '';
         owner = parts[parts.length - 2] || '';
+        if (type === 'azure') {
+          const userInfoMatch = input.match(/^(?:https?:\/\/)?([^@\/=\s]+)@/);
+          const candidateOwner = userInfoMatch?.[1] || parts[0] || '';
+          owner = candidateOwner === '_git' ? (parts[0] || '') : candidateOwner;
+        } else {
+          owner = parts[parts.length - 2] || '';
+        }
       }
     }
     // Unsupported URL formats
@@ -322,6 +333,7 @@ export default function Home() {
           excludedFiles,
           includedDirs,
           includedFiles,
+          branch,
         };
         existingConfigs[currentRepoUrl] = configToSave;
         localStorage.setItem(REPO_CONFIG_CACHE_KEY, JSON.stringify(existingConfigs));
@@ -376,6 +388,11 @@ export default function Home() {
       params.append('included_files', includedFiles);
     }
 
+    // Add branch parameter if provided
+    if (branch && branch.trim() !== '') {
+      params.append('branch', branch.trim());
+    }
+
     // Add language parameter
     params.append('language', selectedLanguage);
 
@@ -445,6 +462,8 @@ export default function Home() {
             isOpen={isConfigModalOpen}
             onClose={() => setIsConfigModalOpen(false)}
             repositoryInput={repositoryInput}
+            branch={branch}
+            setBranch={setBranch}
             selectedLanguage={selectedLanguage}
             setSelectedLanguage={setSelectedLanguage}
             supportedLanguages={supportedLanguages}

src/components/ConfigurationModal.tsx

@@ -12,6 +12,10 @@ interface ConfigurationModalProps {
   // Repository input
   repositoryInput: string;
 
+  // Branch input
+  branch: string;
+  setBranch: (value: string) => void;
+
   // Language selection
   selectedLanguage: string;
   setSelectedLanguage: (value: string) => void;
@@ -32,8 +36,8 @@ interface ConfigurationModalProps {
   setCustomModel: (value: string) => void;
 
   // Platform selection
-  selectedPlatform: 'github' | 'gitlab' | 'bitbucket';
-  setSelectedPlatform: (value: 'github' | 'gitlab' | 'bitbucket') => void;
+  selectedPlatform: 'github' | 'gitlab' | 'bitbucket' | 'azure';
+  setSelectedPlatform: (value: 'github' | 'gitlab' | 'bitbucket' | 'azure') => void;
 
   // Access token
   accessToken: string;
@@ -64,6 +68,8 @@ export default function ConfigurationModal({
   isOpen,
   onClose,
   repositoryInput,
+  branch,
+  setBranch,
   selectedLanguage,
   setSelectedLanguage,
   supportedLanguages,
@@ -135,6 +141,29 @@ export default function ConfigurationModal({
               </div>
             </div>
 
+            {/* Branch input */}
+            <div className="mb-4">
+              <label htmlFor="branch-input" className="block text-sm font-medium text-[var(--foreground)] mb-2">
+                {t.form?.branchLabel || 'Branch (Optional)'}
+              </label>
+              <input
+                type="text"
+                id="branch-input"
+                value={branch}
+                onChange={(e) => setBranch(e.target.value)}
+                placeholder="e.g., main, master, develop"
+                className="input-japanese block w-full px-3 py-2 text-sm rounded-md bg-transparent text-[var(--foreground)] focus:outline-none focus:border-[var(--accent-primary)]"
+              />
+              <div className="flex items-start mt-2 text-xs text-[var(--muted)]">
+                <svg xmlns="http://www.w3.org/2000/svg" className="h-4 w-4 mr-1 text-[var(--muted)]" fill="none" viewBox="0 0 24 24" stroke="currentColor">
+                  <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M13 16h-1v-4h-1m1-4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z" />
+                </svg>
+                <span>
+                  {t.form?.branchHelp || "If left empty, the repository's default branch will be used."}
+                </span>
+              </div>
+            </div>
+
             {/* Language selection */}
             <div className="mb-4">
               <label htmlFor="language-select" className="block text-sm font-medium text-[var(--foreground)] mb-2">

src/components/ModelSelectionModal.tsx

@@ -37,7 +37,7 @@ interface ModelSelectionModalProps {
 
   // Token input for refresh
   showTokenInput?: boolean;
-  repositoryType?: 'github' | 'gitlab' | 'bitbucket';
+  repositoryType?: 'github' | 'gitlab' | 'bitbucket' | 'azure';
   // Authentication
   authRequired?: boolean;
   authCode?: string;
@@ -91,7 +91,7 @@ export default function ModelSelectionModal({
 
   // Token input state
   const [localAccessToken, setLocalAccessToken] = useState('');
-  const [localSelectedPlatform, setLocalSelectedPlatform] = useState<'github' | 'gitlab' | 'bitbucket'>(repositoryType);
+  const [localSelectedPlatform, setLocalSelectedPlatform] = useState<'github' | 'gitlab' | 'bitbucket' | 'azure'>(repositoryType);
   const [showTokenSection, setShowTokenSection] = useState(showTokenInput);
 
   // Reset local state when modal is opened

src/components/TokenInput.tsx

@@ -4,8 +4,8 @@ import React from 'react';
 import { useLanguage } from '@/contexts/LanguageContext';
 
 interface TokenInputProps {
-  selectedPlatform: 'github' | 'gitlab' | 'bitbucket';
-  setSelectedPlatform: (value: 'github' | 'gitlab' | 'bitbucket') => void;
+  selectedPlatform: 'github' | 'gitlab' | 'bitbucket' | 'azure';
+  setSelectedPlatform: (value: 'github' | 'gitlab' | 'bitbucket' | 'azure') => void;
   accessToken: string;
   setAccessToken: (value: string) => void;
   showTokenSection?: boolean;
@@ -76,6 +76,16 @@ export default function TokenInput({
                 >
                   <span className="text-sm">Bitbucket</span>
                 </button>
+                <button
+                  type="button"
+                  onClick={() => setSelectedPlatform('azure')}
+                  className={`flex-1 flex items-center justify-center gap-2 px-3 py-2 rounded-md border transition-all ${selectedPlatform === 'azure'
+                    ? 'bg-[var(--accent-primary)]/10 border-[var(--accent-primary)] text-[var(--accent-primary)] shadow-sm'
+                    : 'border-[var(--border-color)] text-[var(--foreground)] hover:bg-[var(--background)]'
+                    }`}
+                >
+                  <span className="text-sm">Azure</span>
+                </button>
               </div>
             </div>
           )}

src/messages/en.json

@@ -33,6 +33,8 @@
     "configureWiki": "Configure Wiki",
     "repoPlaceholder": "owner/repo or GitHub/GitLab/Bitbucket URL",
     "wikiLanguage": "Wiki Language",
+    "branchLabel": "Branch (Optional)",
+    "branchHelp": "If left empty, the repository's default branch will be used.",
     "modelOptions": "Model Options",
     "modelProvider": "Model Provider",
     "modelSelection": "Model Selection",

src/messages/es.json

@@ -33,6 +33,8 @@
     "configureWiki": "Configurar Wiki",
     "repoPlaceholder": "propietario/repositorio o URL de GitHub/GitLab/Bitbucket",
     "wikiLanguage": "Idioma del Wiki",
+    "branchLabel": "Rama (Opcional)",
+    "branchHelp": "Si no ingresas nada, se usará la rama por defecto del repositorio.",
     "modelOptions": "Opciones de Modelo",
     "modelProvider": "Proveedor de Modelo",
     "modelSelection": "Selección de Modelo",

src/messages/fr.json

@@ -33,6 +33,8 @@
       "configureWiki": "Configurer le Wiki",
       "repoPlaceholder": "propriétaire/dépôt ou URL GitHub/GitLab/Bitbucket",
       "wikiLanguage": "Langue du Wiki",
+      "branchLabel": "Branche (facultatif)",
+      "branchHelp": "Si ce champ est vide, la branche par défaut du dépôt sera utilisée.",
       "modelOptions": "Options du Modèle",
       "modelProvider": "Fournisseur du Modèle",
       "modelSelection": "Sélection du Modèle",

src/messages/ja.json

@@ -33,6 +33,8 @@
     "configureWiki": "Wiki設定",
     "repoPlaceholder": "所有者/リポジトリまたはGitHub/GitLab/BitbucketのURL",
     "wikiLanguage": "Wiki言語",
+    "branchLabel": "ブランチ（任意）",
+    "branchHelp": "空の場合、リポジトリのデフォルトブランチが使用されます。",
     "modelOptions": "モデルオプション",
     "modelProvider": "モデルプロバイダー",
     "modelSelection": "モデル選択",

src/messages/kr.json

@@ -33,6 +33,8 @@
       "configureWiki": "위키 구성",
       "repoPlaceholder": "owner/repo 또는 GitHub/GitLab/Bitbucket URL",
       "wikiLanguage": "위키 언어",
+      "branchLabel": "브랜치(선택 사항)",
+      "branchHelp": "비워두면 저장소의 기본 브랜치가 사용됩니다.",
       "modelOptions": "모델 옵션",
       "modelProvider": "모델 제공자",
       "modelSelection": "모델 선택",