feat: batch index/clone example repos (#4628)

Author: sahil485

servers/fai-lambda/fai-code-indexing/src/handler.py

@@ -14,7 +14,7 @@
 
 from .operations import (
     run_code_search_tool_call,
-    setup_repo_for_domain,
+    setup_repos_for_domain,
 )
 
 logger = logging.getLogger()
@@ -32,9 +32,10 @@ def handler(event: dict[str, Any], context: Any) -> dict[str, Any]:
         event_type: Literal["codeSearch", "indexRepo"] = validate_body_param_or_throw(body, "eventType")
 
         if event_type == "indexRepo":
-            repo_url = validate_body_param_or_throw(body, "repoUrl")
-            asyncio.run(setup_repo_for_domain(domain=domain, repo_url=repo_url))
-            message = "Repository indexed successfully"
+            repo_urls = validate_body_param_or_throw(body, "repoUrls", list[str])
+            asyncio.run(setup_repos_for_domain(domain=domain, repo_urls=repo_urls))
+
+            message = f"Successfully indexed {len(repo_urls)} repositories"
         elif event_type == "codeSearch":
             asyncio.run(run_code_search_tool_call(domain=domain))
             message = "Code search completed successfully"

servers/fai-lambda/fai-code-indexing/src/operations/__init__.py

@@ -1,5 +1,5 @@
 from .analysis import analyze_repositories_for_domain
-from .indexing import setup_repo_for_domain
+from .indexing import setup_repos_for_domain
 from .search import run_code_search_tool_call
 
-__all__ = ["analyze_repositories_for_domain", "setup_repo_for_domain", "run_code_search_tool_call"]
+__all__ = ["analyze_repositories_for_domain", "setup_repos_for_domain", "run_code_search_tool_call"]

servers/fai-lambda/fai-code-indexing/src/operations/indexing.py

@@ -7,20 +7,20 @@
 logger = logging.getLogger()
 
 
-async def setup_repo_for_domain(domain: str, repo_url: str) -> SetupRepoResult:
-    """Set up a repository for a domain by cloning and indexing it.
+async def setup_repos_for_domain(domain: str, repo_urls: list[str]) -> SetupRepoResult:
+    """Set up repositories for a domain by cloning and indexing them.
 
     Args:
-        domain: The domain to associate the repository with
-        repo_url: The GitHub repository URL to clone and index
+        domain: The domain to associate the repositories with
+        repo_urls: List of GitHub repository URLs to clone and index
 
     Returns:
         Dictionary with setup results
     """
-    logger.info(f"Setting up repository for domain: {domain}, repo: {repo_url}")
+    logger.info(f"Setting up repositories for domain: {domain}, repos: {repo_urls}")
 
-    repo_path = clone_repo_to_domain(domain=domain, repo_url=repo_url)
-    logger.info(f"Repository cloned to: {repo_path}")
+    for repo_url in repo_urls:
+        await clone_repo_to_domain(domain=domain, repo_url=repo_url)
 
     analysis_result = await analyze_repositories_for_domain(domain=domain)
     logger.info(f"Analysis completed for domain {domain} repositories")

servers/fai-lambda/fai-code-indexing/src/utils/git.py

@@ -1,12 +1,12 @@
+import asyncio
 import logging
 import os
-import subprocess
 from pathlib import Path
 
 logger = logging.getLogger()
 
 
-def clone_repo_to_domain(domain: str, repo_url: str) -> str:
+async def clone_repo_to_domain(domain: str, repo_url: str) -> str:
     """Clone a GitHub repository into EFS under a domain folder.
 
     Args:
@@ -33,40 +33,69 @@ def clone_repo_to_domain(domain: str, repo_url: str) -> str:
     if repo_path.exists():
         logger.info(f"Repository already exists at {repo_path}, pulling latest changes")
         try:
-            subprocess.run(
-                ["git", "config", "--global", "--add", "safe.directory", str(repo_path)],
-                capture_output=True,
-                text=True,
+            config_process = await asyncio.create_subprocess_exec(
+                "git",
+                "config",
+                "--global",
+                "--add",
+                "safe.directory",
+                str(repo_path),
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
             )
-            subprocess.run(
-                ["git", "-C", str(repo_path), "fetch", "origin"],
-                check=True,
-                capture_output=True,
-                text=True,
+            _, _ = await config_process.communicate()
+
+            process = await asyncio.create_subprocess_exec(
+                "git",
+                "-C",
+                str(repo_path),
+                "fetch",
+                "origin",
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
             )
-            subprocess.run(
-                ["git", "-C", str(repo_path), "pull", "origin"],
-                check=True,
-                capture_output=True,
-                text=True,
+            _, stderr = await process.communicate()
+            if process.returncode != 0:
+                raise RuntimeError(f"Failed to fetch: {stderr.decode()}")
+
+            process = await asyncio.create_subprocess_exec(
+                "git",
+                "-C",
+                str(repo_path),
+                "pull",
+                "origin",
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
             )
+            _, stderr = await process.communicate()
+            if process.returncode != 0:
+                raise RuntimeError(f"Failed to pull: {stderr.decode()}")
+
             logger.info(f"Successfully pulled latest changes at {repo_path}")
-        except subprocess.CalledProcessError as e:
-            logger.error(f"Failed to pull repository: {e.stderr}")
-            raise RuntimeError(f"Failed to pull latest changes: {e.stderr}")
+        except Exception as e:
+            logger.error(f"Failed to pull repository: {e}")
+            raise RuntimeError(f"Failed to pull latest changes: {e}")
     else:
         clone_url = f"https://x-access-token:{github_token}@github.com/{repo_identifier}.git"
 
         logger.info(f"Cloning {repo_identifier} into {repo_path} (shallow)")
         try:
-            subprocess.run(
-                ["git", "clone", "--depth", "1", clone_url, str(repo_path)],
-                check=True,
-                capture_output=True,
-                text=True,
+            process = await asyncio.create_subprocess_exec(
+                "git",
+                "clone",
+                "--depth",
+                "1",
+                clone_url,
+                str(repo_path),
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
             )
-        except subprocess.CalledProcessError as e:
-            logger.error(f"Failed to clone repository: {e.stderr}")
-            raise RuntimeError(f"Failed to clone {repo_identifier}: {e.stderr}")
+            _, stderr = await process.communicate()
+            if process.returncode != 0:
+                raise RuntimeError(f"Failed to clone {repo_identifier}: {stderr.decode()}")
+            logger.info(f"Repository cloned to: {repo_path}")
+        except Exception as e:
+            logger.error(f"Failed to clone repository: {e}")
+            raise RuntimeError(f"Failed to clone {repo_identifier}: {e}")
 
     return str(repo_path)

servers/fai-lambda/shared/utils/validation.py

@@ -1,8 +1,20 @@
-from typing import Any
+from typing import Any, Type, TypeVar, get_origin
 
+T = TypeVar("T")
 
-def validate_body_param_or_throw(body: dict[str, Any], param_name: str) -> str:
+
+def validate_body_param_or_throw(body: dict[str, Any], param_name: str, expected_type: Type[T] = str) -> T:
     value = body.get(param_name)
-    if not isinstance(value, str) or not value:
+
+    if value is None:
         raise ValueError(f"Missing required request body field: {param_name}")
+
+    origin = get_origin(expected_type)
+    if origin is not None:
+        if not isinstance(value, origin):
+            raise ValueError(f"Field '{param_name}' must be of type {expected_type}, got {type(value).__name__}")
+    else:
+        if not isinstance(value, expected_type):
+            raise ValueError(f"Field '{param_name}' must be of type {expected_type.__name__}, got {type(value).__name__}")
+
     return value

servers/fai/src/fai/models/api/github_source_api.py

@@ -7,12 +7,12 @@
 
 
 class IndexGithubRequest(BaseModel):
-    repo_url: str = Field(description="GitHub repository URL to index")
+    repo_urls: list[str] = Field(description="GitHub repository URLs to index")
 
 
 class IndexGithubResponse(BaseModel):
     job_id: str = Field(description="Job ID for tracking indexing progress")
-    repo_url: str = Field(description="GitHub repository URL being indexed")
+    repo_urls: list[str] = Field(description="GitHub repository URLs being indexed")
 
 
 class GithubIndexStatusResponse(BaseModel):

servers/fai/src/fai/routes/sources/__init__.py

@@ -0,0 +1,3 @@
+from . import github
+
+__all__ = ["github"]

servers/fai/src/fai/routes/sources/github.py

@@ -1,5 +1,17 @@
-from fastapi import Depends
-from fastapi.responses import JSONResponse
+import json
+import logging
+import uuid
+from datetime import (
+    UTC,
+    datetime,
+)
+
+import aioboto3
+from botocore.exceptions import ClientError
+from fastapi import (
+    Depends,
+    HTTPException,
+)
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from fai.app import fai_app
@@ -9,12 +21,18 @@
     verify_token,
 )
 from fai.models.api.github_source_api import (
-    GithubIndexStatusResponse,
     IndexGithubRequest,
     IndexGithubResponse,
-    ReindexGithubRequest,
-    ReindexGithubResponse,
 )
+from fai.models.db.index_source_db import (
+    IndexSourceDb,
+    IndexSourceStatus,
+    SourceType,
+)
+
+logger = logging.getLogger(__name__)
+
+LAMBDA_FUNCTION_NAME = "fai-code-indexing-dev2"
 
 
 @fai_app.post(
@@ -23,60 +41,65 @@
     dependencies=[Depends(verify_token)],
     openapi_extra={"x-fern-audiences": ["internal"]},
 )
-async def index_github_source(
+async def index_github_source_repos(
     domain: str,
     request: IndexGithubRequest,
     db: AsyncSession = Depends(get_db),
-) -> JSONResponse:
+) -> IndexGithubResponse:
     """Start indexing a GitHub repository for a domain."""
-    strip_domain(domain)
+    stripped_domain = strip_domain(domain)
 
-    # TODO: Implement GitHub indexing logic
-    # - Create IndexSourceDb record
-    # - Start background indexing job
-    # - Return job_id and repo_url
+    job_id = str(uuid.uuid4())
 
-    raise NotImplementedError("GitHub indexing not yet implemented")
+    now = datetime.now(UTC)
+    for repo_url in request.repo_urls:
+        index_source = IndexSourceDb(
+            id=str(uuid.uuid4()),
+            domain=stripped_domain,
+            source_type=SourceType.GITHUB,
+            source_identifier=repo_url,
+            config={},
+            job_id=job_id,
+            status=IndexSourceStatus.INDEXING,
+            metrics={},
+            created_at=now,
+            updated_at=now,
+        )
+        db.add(index_source)
 
+    await db.commit()
 
-@fai_app.get(
-    "/sources/github/{domain}/status",
-    response_model=GithubIndexStatusResponse,
-    dependencies=[Depends(verify_token)],
-    openapi_extra={"x-fern-audiences": ["internal"]},
-)
-async def get_github_index_status(
-    domain: str,
-    db: AsyncSession = Depends(get_db),
-) -> JSONResponse:
-    """Get the indexing status for a GitHub repository."""
-    strip_domain(domain)
-
-    # TODO: Implement status check logic
-    # - Query IndexSourceDb for the domain
-    # - Return current status and metrics
-
-    raise NotImplementedError("GitHub status check not yet implemented")
+    try:
+        session = aioboto3.Session()
+        async with session.client("lambda") as lambda_client:
+            payload = {
+                "domain": stripped_domain,
+                "eventType": "indexRepo",
+                "repoUrls": request.repo_urls,
+            }
 
+            response = await lambda_client.invoke(
+                FunctionName=LAMBDA_FUNCTION_NAME,
+                InvocationType="Event",
+                Payload=json.dumps(payload),
+            )
 
-@fai_app.post(
-    "/sources/github/{domain}/reindex",
-    response_model=ReindexGithubResponse,
-    dependencies=[Depends(verify_token)],
-    openapi_extra={"x-fern-audiences": ["internal"]},
-)
-async def reindex_github_source(
-    domain: str,
-    request: ReindexGithubRequest,
-    db: AsyncSession = Depends(get_db),
-) -> JSONResponse:
-    """Delete existing index and start a new indexing job for a GitHub repository."""
-    strip_domain(domain)
+            logger.info(
+                f"Successfully invoked code indexing Lambda. "
+                f"StatusCode: {response.get('StatusCode')}, "
+                f"Domain: {stripped_domain}, "
+                f"RepoUrls: {request.repo_urls}, "
+                f"JobId: {job_id}"
+            )
 
-    # TODO: Implement reindexing logic
-    # - Delete old IndexSourceDb records and indexed content
-    # - Create new IndexSourceDb record
-    # - Start new background indexing job
-    # - Return new job_id
+    except ClientError as e:
+        logger.error(
+            f"Failed to invoke code indexing Lambda: {e.response['Error']['Code']} - {e.response['Error']['Message']}",
+            exc_info=True,
+        )
+        raise HTTPException(status_code=500, detail="Failed to start indexing job")
+    except Exception as e:
+        logger.error(f"Unexpected error invoking code indexing Lambda: {str(e)}", exc_info=True)
+        raise HTTPException(status_code=500, detail="Failed to start indexing job")
 
-    raise NotImplementedError("GitHub reindexing not yet implemented")
+    return IndexGithubResponse(job_id=job_id, repo_urls=request.repo_urls)