Spaces:
Running
Running
| # server/tasks.py | |
| import os | |
| import shutil | |
| import tempfile | |
| import zipfile | |
| import json | |
| import asyncio | |
| from pathlib import Path | |
| from typing import AsyncGenerator, List | |
| from git import Repo, GitCommandError | |
| from codescribe.config import load_config | |
| from codescribe.llm_handler import LLMHandler | |
| from codescribe.orchestrator import DocstringOrchestrator | |
| from codescribe.readme_generator import ReadmeGenerator | |
| async def process_project( | |
| project_path: Path, | |
| description: str, | |
| readme_note: str, | |
| is_temp: bool, | |
| exclude_list: List[str], | |
| new_branch_name: str = None, | |
| repo_full_name: str = None, | |
| github_token: str = None | |
| ) -> AsyncGenerator[str, None]: | |
| loop = asyncio.get_running_loop() | |
| queue = asyncio.Queue() | |
| def emit_event(event: str, data: dict): | |
| """A thread-safe way to send events from the worker thread to the async generator.""" | |
| loop.call_soon_threadsafe(queue.put_nowait, {"event": event, "data": data}) | |
| def _blocking_process(): | |
| """The main, synchronous processing logic that runs in a separate thread.""" | |
| try: | |
| # --- Setup --- | |
| config = load_config() | |
| llm_handler = LLMHandler(config.api_keys, progress_callback=lambda msg: emit_event("log", {"message": msg})) | |
| # --- 1. Docstrings Phase (using the orchestrator) --- | |
| doc_orchestrator = DocstringOrchestrator( | |
| path_or_url=str(project_path), | |
| description=description, | |
| exclude=exclude_list, | |
| llm_handler=llm_handler, | |
| progress_callback=emit_event, | |
| repo_full_name=repo_full_name # <<< THIS IS THE FIX | |
| ) | |
| doc_orchestrator.project_path = project_path | |
| doc_orchestrator.is_temp_dir = False | |
| doc_orchestrator.run() | |
| # --- 2. READMEs Phase (using the generator) --- | |
| readme_gen = ReadmeGenerator( | |
| path_or_url=str(project_path), | |
| description=description, | |
| exclude=exclude_list, | |
| llm_handler=llm_handler, | |
| user_note=readme_note, | |
| progress_callback=emit_event, | |
| repo_full_name=repo_full_name | |
| ) | |
| readme_gen.project_path = project_path | |
| readme_gen.is_temp_dir = False | |
| readme_gen.run() | |
| # --- 3. Output Phase --- | |
| emit_event("phase", {"id": "output", "status": "in-progress"}) | |
| if new_branch_name: | |
| # Git logic remains the same | |
| emit_event("subtask", {"parentId": "output", "listId": "output-step-list", "id": "git-check", "name": "Checking for changes...", "status": "in-progress"}) | |
| repo = Repo(project_path) | |
| if repo.untracked_files: | |
| repo.git.add(repo.untracked_files) | |
| if not repo.is_dirty(untracked_files=True): | |
| emit_event("subtask", {"parentId": "output", "id": "git-check", "status": "success"}) | |
| emit_event("log", {"message": "No changes were generated by the AI."}) | |
| emit_event("done", {"type": "github", "url": f"https://github.com/{repo_full_name}", "message": "✅ No changes needed. Your repository is up to date."}) | |
| return | |
| emit_event("subtask", {"parentId": "output", "id": "git-check", "status": "success"}) | |
| emit_event("subtask", {"parentId": "output", "listId": "output-step-list", "id": "git-push", "name": f"Pushing to branch '{new_branch_name}'...", "status": "in-progress"}) | |
| try: | |
| if new_branch_name in repo.heads: | |
| repo.heads[new_branch_name].checkout() | |
| else: | |
| repo.create_head(new_branch_name).checkout() | |
| repo.git.add(A=True) | |
| repo.index.commit("docs: Add AI-generated documentation by CodeScribe") | |
| origin = repo.remote(name='origin') | |
| origin.push(new_branch_name, force=True) | |
| pr_url = f"https://github.com/{repo_full_name}/pull/new/{new_branch_name}" | |
| emit_event("subtask", {"parentId": "output", "id": "git-push", "status": "success"}) | |
| emit_event("done", {"type": "github", "url": pr_url, "message": "✅ Successfully pushed changes!"}) | |
| except GitCommandError as e: | |
| emit_event("log", {"message": f"Git Error: {e}"}) | |
| raise RuntimeError(f"Failed to push to GitHub: {e}") | |
| else: | |
| # ZIP logic remains the same | |
| emit_event("subtask", {"parentId": "output", "listId": "output-step-list", "id": "zip-create", "name": "Creating downloadable ZIP file...", "status": "in-progress"}) | |
| temp_dir = tempfile.gettempdir() | |
| try: | |
| contained_dir = next(project_path.iterdir()) | |
| project_name = contained_dir.name | |
| except StopIteration: | |
| project_name = "documented-project" | |
| zip_filename_base = f"codescribe-docs-{project_name}" | |
| zip_path_base = Path(temp_dir) / zip_filename_base | |
| zip_full_path = shutil.make_archive(str(zip_path_base), 'zip', project_path) | |
| zip_file_name = Path(zip_full_path).name | |
| emit_event("subtask", {"parentId": "output", "id": "zip-create", "status": "success"}) | |
| emit_event("done", {"type": "zip", "download_path": zip_file_name, "message": "✅ Your documented project is ready for download."}) | |
| emit_event("phase", {"id": "output", "status": "success"}) | |
| except Exception as e: | |
| emit_event("error", str(e)) | |
| finally: | |
| loop.call_soon_threadsafe(queue.put_nowait, None) | |
| # Async runner remains the same | |
| main_task = loop.run_in_executor(None, _blocking_process) | |
| while True: | |
| message = await queue.get() | |
| if message is None: | |
| break | |
| json_line = json.dumps({"type": message["event"], "payload": message["data"]}) | |
| yield f"{json_line}\n" | |
| await main_task | |
| if is_temp and project_path and project_path.exists(): | |
| shutil.rmtree(project_path, ignore_errors=True) |