iframe-proxy | Sunbelt Computer Software

549 lines (467 loc) · 22.5 KB
import os, json, subprocess, tempfile, re
from google.genai import types
from vimini import util, context
# Global data store keyed by buffer number to exchange data between python calls.
_BUFFER_DATA_STORE = {}
# Separator line to distinguish between thoughts/summary and the actual diff
_DIFF_SEPARATOR = "========== VIMINI DIFF START =========="
def code(prompt, verbose=False, temperature=None):
    Uploads all open files, sends them to the Gemini API with a prompt
    to generate code. Runs asynchronously using a thread and Vim timer.
    util.log_info(f"code({prompt}, verbose={verbose}, temperature={temperature})")
    # --- 1. Initialization and Setup ---
        client = util.get_client()
        if not client:
            return
        project_root = util.get_git_repo_root()
        if not project_root:
            project_root = vim.eval('getcwd()')
        # Gather context files manually and resolve them relative to project_root
        file_paths_to_include = []
        for b in vim.buffers:
            if b.name and os.path.exists(b.name):
                file_paths_to_include.append(os.path.abspath(b.name))
        try:
            context_files_list = vim.eval("get(g:, 'context_files', [])")
            if isinstance(context_files_list, list):
                for f in context_files_list:
                    if os.path.isabs(f):
                        file_paths_to_include.append(os.path.abspath(f))
                        file_paths_to_include.append(os.path.abspath(os.path.join(project_root, f)))
        except Exception:
            pass
        # Upload context files using the helper function.
        uploaded_files = context.upload_context_files(client, file_paths_to_include=file_paths_to_include)
        if uploaded_files is None:
            return  # The helper function has already displayed an error.
    except Exception as e:
        util.display_message(f"Error during initialization: {e}", error=True)
        return
    # --- 2. Define Schema and Prompt ---
    file_object_schema = types.Schema(
        type=types.Type.OBJECT,
        properties={
            'file_path': types.Schema(type=types.Type.STRING, description="The full path of the file relative to the project directory."),
            'file_type': types.Schema(type=types.Type.STRING, description="The content type. Use 'text/plain' for the full file content or 'text/x-diff' for a patch in the unified diff format."),
            'file_content': types.Schema(type=types.Type.STRING, description="The new, complete source code for the file, or a patch in the unified diff format, corresponding to the file_type.")
        required=['file_path', 'file_type', 'file_content']
    multi_file_output_schema = types.Schema(
        type=types.Type.OBJECT,
        properties={
            'files': types.Schema(
                type=types.Type.ARRAY,
                items=file_object_schema
        required=['files']
    original_buffer = vim.current.buffer
    if original_buffer.name:
        main_file_name = os.path.relpath(original_buffer.name, project_root) if os.path.isabs(original_buffer.name) else original_buffer.name
        task_instruction = f"Your primary task is to modify the file named '{main_file_name}'."
        task_instruction = "Your primary task is to address the concern in the active buffer (if any).\n"
        buffer_content = "\n".join(original_buffer[:])
        if buffer_content.strip():
            task_instruction += f"\n\nAdditional context from the current active buffer:\n{buffer_content}\n"
    context_file_names = sorted([f.display_name for f in uploaded_files])
    file_list_str = "\n".join(f"- {name}" for name in context_file_names)
    context_files_section = ""
    if file_list_str:
        context_files_section = (
            "The following files have been uploaded for context:\n"
            f"{file_list_str}\n\n"
    full_prompt = [
            f"{prompt}\n\n"
            "Based on the user's request, please generate the code. "
            "Your identity is Vimini, and you are integrated into the vimini project."
            f"{task_instruction}\n\n"
            f"{context_files_section}"
            "IMPORTANT:\n"
            "1. Your response must be a single JSON object with a 'files' key.\n"
            "2. The value of 'files' must be an array of file objects.\n"
            "3. Each file object must have three string keys: 'file_path', 'file_type', and 'file_content'.\n"
            "4. 'file_path' must be the full path of the file relative to the project directory. When modifying a file from the context, you MUST use its original file path for the 'file_path' property.\n"
            "5. 'file_type' must be either 'text/plain' for the full file content or 'text/x-diff' for a patch in the unified diff format.\n"
            "6. 'file_content' must contain either the new, complete source code or the diff patch, corresponding to the 'file_type'.\n"
            "7. Diffs ('text/x-diff') can be returned only if explicitly mentioned as an acceptable output in the prompt or if the files are really difficult or too large to process. For small files, returning the entire modified file ('text/plain') is the most preferred option.\n"
            "8. You can modify existing files or create new files as needed to fulfill the request."
        *uploaded_files
    job_name = f"Code: {prompt}"
    job_id = util.reserve_next_job_id(job_name)
    # --- 3. Create Code Buffer ---
    # Create the buffer immediately to store thoughts or request summary.
    util.new_split()
    base_buffer_name = f"[{job_id}] Vimini Code"
    safe_name = f"{base_buffer_name} [->G?]".replace(" ", "\\ ")
    vim.command(f"file {safe_name}")
    vim.command("setlocal buftype=nofile")
    vim.command("setlocal bufhidden=wipe")
    vim.command("setlocal noswapfile")
    vim.command("setlocal filetype=markdown")
    code_buffer = vim.current.buffer
    code_buffer_num = code_buffer.number
    # Store buffer-local variables
    vim.command(f"let b:vimini_project_root = '{project_root}'")
    vim.command(f"let b:vimini_job_id = '{job_id}'")
    util.append_job_summary(code_buffer_num, job_id, prompt, context_file_names)
    util.display_message("Processing... (Async)")
    # State for the closure
    json_aggregator = ""
    started_receiving = False
    def update_status_receiving():
        nonlocal started_receiving
        if not started_receiving:
            started_receiving = True
            try:
                code_buffer.name = f"{base_buffer_name} [<-G]"
            except Exception:
    def on_chunk(text):
        nonlocal json_aggregator
        update_status_receiving()
        json_aggregator += text
    def on_thought(text):
        update_status_receiving()
        if verbose:
            util.append_to_buffer(code_buffer_num, text)
    def on_finish():
        try:
            code_buffer.name = base_buffer_name
        except Exception:
            pass
        return _finalize_code_generation(json_aggregator, project_root, job_id, code_buffer_num)
    def on_error(msg):
        util.append_to_buffer(code_buffer_num, f"\nError: {msg}")
        return f"Error: {msg}"
    kwargs = util.create_generation_kwargs(
        contents=full_prompt,
        temperature=temperature,
        verbose=verbose,
        response_mime_type="application/json",
        response_schema=multi_file_output_schema
    util.start_async_job(client, kwargs, {
        "on_chunk": on_chunk,
        "on_thought": on_thought,
        "on_finish": on_finish,
        "on_error": on_error
    }, job_id=job_id)
def _finalize_code_generation(json_aggregator, project_root, job_id, buffer_num):
    """Parses accumulated JSON and generates diff."""
    global _BUFFER_DATA_STORE
    # --- 5. Parse JSON Response ---
        parsed_json = json.loads(json_aggregator)
        files_to_process = parsed_json.get("files", [])
        if not isinstance(files_to_process, list):
            raise ValueError("'files' key is not a list.")
    except (json.JSONDecodeError, ValueError) as e:
        util.append_to_buffer(buffer_num, f"\nError parsing JSON: {e}\n\nRaw JSON:\n{json_aggregator}")
        return f"AI did not return valid JSON for files: {e}"
    if not files_to_process:
        util.append_to_buffer(buffer_num, "\nAI returned no file changes.")
        return "AI returned no file changes."
    # Update Data Store
    _BUFFER_DATA_STORE[buffer_num] = {
        "files_to_apply": files_to_process,
        "project_root": project_root,
        "job_id": job_id
    # --- 6. Generate and Display Diff ---
        combined_diff_output = []
        for file_op in files_to_process:
            api_path = file_op["file_path"]
            ai_generated_code = file_op["file_content"]
            # Ensure file ends with newline
            if ai_generated_code and not ai_generated_code.endswith('\n'):
                ai_generated_code += '\n'
            file_type = file_op.get("file_type", "text/plain")
            absolute_path = util.get_absolute_path_from_api_path(api_path)
            file_exists = os.path.exists(absolute_path)
            relative_path = os.path.relpath(absolute_path, project_root)
            if file_type == "text/x-diff":
                # Do not strip trailing empty lines indiscriminately, but clean up the string ends.
                # .strip('\n') removes leading/trailing newlines but preserves context spaces.
                lines = ai_generated_code.strip('\n').split("\n")
                if lines:
                    # Validate that the first line is a diff command
                    if not lines[0].startswith("diff"):
                        lines.insert(0, f"diff --git a/{relative_path} b/{relative_path}")
                    # Validate that individual --- +++ lines are suitable for patch -p1
                    fixed_lines = []
                    for line in lines:
                        if line.startswith("--- "):
                            path = line[4:].strip()
                            # Trust /dev/null if explicitly mentioned
                            if path == "/dev/null":
                                fixed_lines.append(line)
                                # Otherwise, rewrite based on metadata to ensure correctness
                                if not file_exists:
                                    fixed_lines.append("--- /dev/null")
                                    fixed_lines.append(f"--- a/{relative_path}")
                        elif line.startswith("+++ "):
                            path = line[4:].strip()
                            if path == "/dev/null":
                                fixed_lines.append(line)
                                fixed_lines.append(f"+++ b/{relative_path}")
                            fixed_lines.append(line)
                    combined_diff_output.extend(fixed_lines)
            else: # 'text/plain' or unspecified
                original_content = ""
                if file_exists:
                        with open(absolute_path, "r", encoding="utf-8") as f:
                            original_content = f.read()
                    except Exception as e:
                        continue
                with tempfile.NamedTemporaryFile(mode="w+", delete=False, encoding="utf-8") as f_orig, \
                     tempfile.NamedTemporaryFile(mode="w+", delete=False, encoding="utf-8") as f_ai:
                    f_orig.write(original_content)
                    f_ai.write(ai_generated_code)
                    orig_filepath = f_orig.name
                    ai_filepath = f_ai.name
                    source_path = orig_filepath if file_exists else "/dev/null"
                    cmd = ["diff", "-u", source_path, ai_filepath]
                    result = subprocess.run(cmd, capture_output=True, text=True, check=False)
                    if result.returncode > 1:
                        continue # diff error
                    diff_lines = result.stdout.split("\n")
                    if len(diff_lines) <= 2 and not original_content and not ai_generated_code:
                        continue # Empty diff
                    combined_diff_output.append(f"diff --git a/{relative_path} b/{relative_path}")
                    if not file_exists:
                        combined_diff_output.append("new file mode 100644")
                        combined_diff_output.append(f"--- /dev/null")
                        combined_diff_output.append(f"+++ b/{relative_path}")
                        combined_diff_output.append(f"--- a/{relative_path}")
                        combined_diff_output.append(f"+++ b/{relative_path}")
                    combined_diff_output.extend(diff_lines[2:])
                finally:
                    if os.path.exists(orig_filepath): os.remove(orig_filepath)
                    if os.path.exists(ai_filepath): os.remove(ai_filepath)
        if not combined_diff_output:
            util.append_to_buffer(buffer_num, "\nAI content is identical to the original files or returned empty diff.")
            return "AI content is identical to the original files or returned empty diff."
        # Append Separator and Diff
        separator_block = f"\n{_DIFF_SEPARATOR}\n"
        diff_text = "\n".join(combined_diff_output)
        util.append_to_buffer(buffer_num, separator_block + diff_text)
        # Switch filetype to diff for syntax highlighting
        # We use setbufvar to avoid switching windows
        vim.command(f"call setbufvar({buffer_num}, '&filetype', 'diff')")
        return "Diff generated."
    except Exception as e:
        util.append_to_buffer(buffer_num, f"\nError generating diff: {e}")
        return f"Error generating or displaying diff: {e}"
def show_diff():
    Shows the current git modifications in a new buffer.
    util.log_info("show_diff()")
        repo_path = util.get_git_repo_root()
        if not repo_path:
            return # Error message handled by helper
        # Command to get the diff.
        # -C ensures git runs in the correct directory.
        cmd = ['git', '-C', repo_path, 'diff', '--color=never']
        # Execute the command.
        util.display_message("Running git diff...")
        result = subprocess.run(cmd, capture_output=True, text=True, check=False)
        util.display_message("") # Clear message
        # Handle git errors (e.g., not a git repository).
        if result.returncode != 0 and not result.stdout.strip():
            error_message = result.stderr.strip()
            util.display_message(f"Git error: {error_message}", error=True)
            return
        # Handle case with no modifications.
        if not result.stdout.strip():
            util.display_message("No modifications found.", history=True)
            return
        # Display the diff in a new split window.
        util.new_split()
        vim.command("file Git Diff")
        # Setting filetype to 'diff' helps with syntax highlighting
        vim.command("setlocal buftype=nofile filetype=diff noswapfile")
        # The output from git contains ANSI escape codes for color.
        # We place this raw output into the buffer.
        vim.current.buffer[:] = result.stdout.split("\n")
    except FileNotFoundError:
        util.display_message("Error: `git` command not found. Is it in your PATH?", error=True)
    except Exception as e:
        util.display_message(f"Error: {e}", error=True)
def apply_patch(diff_content, project_root=None, silent=False):
    Applies a unified diff patch and reloads affected buffers.
    Returns (True, message) if successful, (False, error_message) otherwise.
    if not diff_content:
        msg = "Diff is empty. Nothing to apply."
        if not silent: util.display_message(msg, history=True)
        return False, msg
    if not project_root:
        project_root = util.get_git_repo_root() or vim.eval("getcwd()")
        # Use patch command to apply the diff
        # -p1 strips the first path component (a/ and b/)
        # -N ignores patches that seem already applied
        # -r - rejects to stdout (avoids .rej files)
        result = subprocess.run(
            ["patch", "-p1", "-N", "-r", "-"],
            input=diff_content, text=True, check=False,
            capture_output=True, cwd=project_root
        if result.returncode != 0:
            err_msg = f"Patch command failed. Please review the output and the diff.\nSTDOUT: {result.stdout}\nSTDERR: {result.stderr}"
            if not silent: util.display_message(err_msg, error=True)
            return False, err_msg
        # Success
        util.display_message("Successfully applied modified diff.", history=True)
        # Parse diff to identify modified files for reloading
        modified_files = set()
        for line in diff_content.split("\n"):
            if line.startswith("--- a/"):
                path = line[len("--- a/"):].strip()
                if path != "/dev/null":
                    modified_files.add(path)
            elif line.startswith("+++ b/"):
                path = line[len("+++ b/"):].strip()
                if path != "/dev/null":
                    modified_files.add(path)
        for relative_path in modified_files:
            absolute_path = os.path.join(project_root, relative_path)
            # Clean up .orig files potentially created by patch
            orig_path = absolute_path + ".orig"
            if os.path.exists(orig_path):
                    os.remove(orig_path)
                except Exception:
            normalized_target_path = os.path.abspath(absolute_path)
            for buf in vim.buffers:
                if buf.name and os.path.abspath(buf.name) == normalized_target_path:
                    # Reload buffer if visible
                    win_nr = vim.eval(f"bufwinnr({buf.number})")
                    if int(win_nr) > 0:
                        vim.command(f"{win_nr}wincmd w")
                        vim.command("e!")
                        vim.command("wincmd p")
                        # Mark buffer to be reloaded when entered
                        vim.command(f"checktime {buf.number}")
        return True
    except FileNotFoundError:
        util.display_message("Error: `patch` command not found. Is it in your PATH?", error=True)
        return False
    except Exception as e:
        util.display_message(f"Error applying diff: {e}", error=True)
        return False
def apply_code(job_id=None):
    Finds the 'Vimini Code' buffer, writes all specified file changes to
    disk, and reloads any affected open buffers. If an error occurs, the
    diff buffer is preserved for manual editing and re-application.
    global _BUFFER_DATA_STORE
    util.log_info(f"apply_code(job_id={job_id})")
    diff_buffer = None
    # 1. Find all potential Vimini Code buffers
    candidates = []
    for buf in vim.buffers:
        if buf.name and 'Vimini Code' in os.path.basename(buf.name):
            candidates.append(buf)
    if not candidates:
        util.display_message("`Vimini Code` buffer not found. Was :ViminiCode run?", error=True)
        return
    # 2. Filter by job_id if provided, or handle selection logic
    if job_id is not None:
        target_candidates = []
        for buf in candidates:
            # Try matching by internal buffer variable
            try:
                bid = vim.eval(f"getbufvar({buf.number}, 'vimini_job_id', '')")
                if bid and int(bid) == job_id:
                    target_candidates.append(buf)
                    continue
            except:
            # Try matching by filename pattern "[{job_id}] Vimini Code"
            basename = os.path.basename(buf.name)
            if f"[{job_id}]" in basename:
                 target_candidates.append(buf)
        if not target_candidates:
            util.display_message(f"No Vimini Code buffer found for Job ID {job_id}.", error=True)
            return
        # If for some reason multiple buffers match the same ID, take the last one
        diff_buffer = target_candidates[-1]
        # No job ID provided.
        if vim.current.buffer in candidates:
            diff_buffer = vim.current.buffer
        elif len(candidates) > 1:
            # Multiple buffers exist: Error and list them
            msg = "Multiple Vimini Code buffers found. Please specify which job to apply using -j <job_id>.\nAvailable Jobs:\n"
            for buf in candidates:
                bid = "Unknown"
                    bid = vim.eval(f"getbufvar({buf.number}, 'vimini_job_id', '')")
                except: pass
                if not bid or bid == "Unknown":
                     m = re.search(r'\[(\d+)\]', os.path.basename(buf.name))
                     if m: bid = m.group(1)
                msg += f"- Job {bid} (Buffer {buf.number})\n"
            util.display_message(msg.strip(), error=True)
            return
        else:
            # Only one buffer exists
            diff_buffer = candidates[0]
    # 4. Extract diff content using separator
    project_root = vim.eval(f"getbufvar({diff_buffer.number}, 'vimini_project_root', '')")
    if not project_root:
        project_root = util.get_git_repo_root() or vim.eval("getcwd()")
    separator_index = -1
    for i, line in enumerate(diff_buffer):
        if _DIFF_SEPARATOR in line:
            separator_index = i
            break
    if separator_index != -1:
        diff_content = "\n".join(diff_buffer[separator_index + 1:])
        # Ensure the patch content ends with a newline
        if diff_content and not diff_content.endswith('\n'):
            diff_content += '\n'
        err_msg = "DIFF section not found, did you remove the separator?"
        util.display_message(err_msg, error=True)
        return  # Preserve buffer
    if not diff_content:
        util.display_message("Diff is empty. Nothing to apply.", history=True)
        vim.command(f"bdelete! {diff_buffer.number}")
        if diff_buffer.number in _BUFFER_DATA_STORE:
            del _BUFFER_DATA_STORE[diff_buffer.number]
        return
    if apply_patch(diff_content, project_root):
        # Remove from data store
        if diff_buffer.number in _BUFFER_DATA_STORE:
            del _BUFFER_DATA_STORE[diff_buffer.number]
        # Cleanup
        vim.command(f"bdelete! {diff_buffer.number}")
Sunbelt Computer Software

PL/B Language Development and Support

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Sunbelt Computer Software

PL/B Language Development and Support

FilesExpand file tree

code.py

Latest commit

History

code.py

File metadata and controls
