ziya 0.1.41__py3-none-any.whl → 0.1.42__py3-none-any.whl

app/agents/agent.py

@@ -55,7 +55,14 @@ def get_combined_docs_from_files(files) -> str:
         try:
             full_file_path = os.path.join(user_codebase_dir, file_path)
             docs = TextLoader(full_file_path).load()
-            for doc in docs:
+            for doc_index, doc in enumerate(docs):
+                # Add line numbers to the content
+                lines = doc.page_content.split('\n')
+                numbered_lines = [f"{i + 1:4d} | {line}" for i, line in enumerate(lines)]
+                numbered_content = '\n'.join(numbered_lines)
+                
+                # Use the numbered content instead of the original
+                doc.page_content = numbered_content
                 combined_contents += f"File: {file_path}\n{doc.page_content}\n\n"
 
         except Exception as e:

app/agents/prompts.py

@@ -6,6 +6,10 @@ template = """
 
 You are an excellent coder. Help the user with their coding tasks. You are given the codebase of the user in your context.
 
+IMPORTANT: The file contents provided to you have line numbers added at the beginning of each line in the format "   
+1 | <line content>". These line numbers start at 1, matching Git's conventions. Use these line numbers when generating 
+diff hunks to ensure accurate line numbering.
+
 IMPORTANT: When recommending code changes, format your response as a standard Git diff format unless the user specifies otherwise. 
 Follow these strict guidelines for diff formatting:
 
@@ -31,11 +35,26 @@ Follow these strict guidelines for diff formatting:
    
 5. End each diff block with ``` on a new line
 
+CRITICAL: After generating each hunk diff, carefully review and verify the following:
+1. Ensure all Hunk Headers (e.g., @@ -4,7 +4,2 @@) are accurate. The numbers should correctly reflect the lines being changed, added, or removed.
+2. Verify that the line numbers in the Hunk Headers match the actual content changes in the diff.
+3. Check that the diff can be applied cleanly using `git apply`. This means:
+   - The context lines (unchanged lines) in the diff should match the original file content.
+   - The line numbers and content should be consistent throughout the diff.
+   - There should be no conflicts or inconsistencies in the changes.
+4. If you find any errors or inconsistencies, correct them before finalizing the diff.
+5. For each hunk in the diff, please make sure it starts or ends with a line containing content instead of empty line, if possible.
+6. When creating a new file, ensure the line `new file mode 100644` is included to specify file permissions.
+5. When deleting a file, include `deleted file mode` to indicate that the file has been removed. Each line in the 
+deleted file should be prefixed with `-` to indicate the content removal.
+
 Do not include any explanatory text within the diff blocks. If you need to provide explanations or comments, do so outside the diff blocks.
 
 The codebase is provided at the end of this prompt in a specific format. 
 The code that the user has given to you for context is in the format like below where first line has the File path and then the content follows.
 Each file starts with "File: <filepath>" followed by its content on subsequent lines. 
+Remember, each line of the file content now starts with a line number, beginning from 1. Use these numbers to generate 
+accurate hunk diffs, but do not include the line numbers in the actual diff content.
 
 File: <filepath>
 <Content of the file>. 
@@ -63,7 +82,6 @@ conversational_prompt = ChatPromptTemplate.from_messages(
     ]
 )
 
-
 def parse_output(message):
     text = message.content
     return AgentFinish(return_values={"output": text}, log=text)

app/main.py

@@ -23,6 +23,8 @@ def parse_arguments():
                         help="Port number to run Ziya frontend on (e.g., --port 8080)")
     parser.add_argument("--version", action="store_true",
                         help="Prints the version of Ziya")
+    parser.add_argument("--enable-code-apply", action="store_true",
+                        help="Enable the feature to show the button to apply code diffs (unstable)")
     parser.add_argument("--max-depth", type=int, default=15,
                         help="Maximum depth for folder structure traversal (e.g., --max-depth 20)")
     return parser.parse_args()
@@ -38,6 +40,8 @@ def setup_environment(args):
         os.environ["ZIYA_AWS_PROFILE"] = args.profile
     if args.model:
         os.environ["ZIYA_AWS_MODEL"] = args.model
+    if args.enable_code_apply:
+        os.environ["ZIYA_ENABLE_CODE_APPLY"] = "true"
 
     os.environ["ZIYA_MAX_DEPTH"] = str(args.max_depth)
 

app/server.py

@@ -2,18 +2,21 @@ import os
 from typing import Dict, Any, List, Tuple
 
 import tiktoken
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 from langserve import add_routes
 from app.agents.agent import agent_executor
 from fastapi.responses import FileResponse
+from pydantic import BaseModel
 
 # import pydevd_pycharm
 import uvicorn
 
+from app.utils.code_util import us_git_to_apply_code_diff, correct_git_diff
 from app.utils.directory_util import get_ignored_patterns
+from app.utils.logging_utils import logger
 from app.utils.gitignore_parser import parse_gitignore_patterns
 
 app = FastAPI()
@@ -35,7 +38,13 @@ add_routes(app, agent_executor, disabled_endpoints=["playground"], path="/ziya")
 
 @app.get("/")
 async def root(request: Request):
-    return templates.TemplateResponse("index.html", {"request": request})
+    enable_code_apply = os.environ.get("ZIYA_ENABLE_CODE_APPLY", "false")
+    enable_code_apply_bool = enable_code_apply.lower() == "true"
+    logger.info(f"enable_code_apply_bool: {enable_code_apply_bool}")
+    return templates.TemplateResponse("index.html", {
+        "request": request,
+        "enable_code_apply": str(enable_code_apply_bool).lower()
+    })
 
 
 @app.get("/favicon.ico", include_in_schema=False)
@@ -97,5 +106,27 @@ def get_default_included_folders():
     return {'defaultIncludedFolders': []}
 
 
+class ApplyChangesRequest(BaseModel):
+    diff: str
+    filePath: str
+
+@app.post('/api/apply-changes')
+async def apply_changes(request: ApplyChangesRequest):
+    try:
+        logger.info(f"Received request to apply changes to file: {request.filePath}")
+        logger.info(f"Diff content: \n{request.diff}")
+
+        user_codebase_dir = os.environ.get("ZIYA_USER_CODEBASE_DIR")
+        if not user_codebase_dir:
+            raise ValueError("ZIYA_USER_CODEBASE_DIR environment variable is not set")
+
+        corrected_diff = correct_git_diff(request.diff)
+        us_git_to_apply_code_diff(corrected_diff)
+        return {'message': 'Changes applied successfully'}
+    except Exception as e:
+        logger.error(f"Error applying changes: {str(e)}", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)

app/utils/code_util.py

@@ -0,0 +1,234 @@
+import os
+import subprocess
+from app.utils.logging_utils import logger
+import time
+import re
+
+HUNK_HEADER_REGEX = re.compile(r'^@@ -(\d+)(?:,(\d+))? \+(\d+)(?:,(\d+))? @@')
+
+def us_git_to_apply_code_diff(git_diff: str):
+    """
+    Apply a git diff to the user's codebase.
+
+    This function takes a git diff as a string and applies it to the user's codebase
+    specified by the ZIYA_USER_CODEBASE_DIR environment variable. It creates a temporary
+    file with the diff content and uses the 'git apply' command to apply the changes.
+
+    Args:
+        git_diff (str): A string containing the git diff to be applied.
+
+    Note:
+        This function assumes that the git command-line tool is available in the system path.
+        It uses the --ignore-whitespace and --ignore-space-change options with 'git apply'
+        to handle potential whitespace issues.
+    """
+
+    user_codebase_dir = os.environ.get("ZIYA_USER_CODEBASE_DIR")
+    if not user_codebase_dir:
+        raise ValueError("ZIYA_USER_CODEBASE_DIR environment variable is not set")
+
+    # Create a temporary file with the diff content and a timestamp
+    timestamp = int(time.time() * 1000)  # Get current timestamp in milliseconds
+    temp_file = os.path.join(user_codebase_dir, f'temp_{timestamp}.diff')
+
+    with open(temp_file, 'w', newline='\n') as f:
+        f.write(git_diff)
+        # need to add a newline at the end of the file of git apply will fail
+        f.write("\n")
+    logger.info(f"Created temporary diff file: {temp_file}")
+
+    try:
+        # Apply the changes using git apply
+        cmd = ['git', 'apply', '--verbose', '--ignore-whitespace', '--ignore-space-change', temp_file]
+        logger.info(f"Executing command: {' '.join(cmd)}")
+        result = subprocess.run(cmd, capture_output=True, text=True, cwd=user_codebase_dir, check=True)
+
+        logger.info(f"Git apply stdout: {result.stdout}")
+        logger.error(f"Git apply stderr: {result.stderr}")
+        logger.info(f"Git apply return code: {result.returncode}")
+
+    except Exception as e:
+        logger.error(f"Error applying changes: {str(e)}", exc_info=True)
+        raise
+    finally:
+        os.remove(temp_file)
+
+def correct_git_diff(git_diff: str) -> str:
+    """
+    Corrects the hunk headers in a git diff string by recalculating the line counts and
+    adjusting the starting line numbers in the new file, considering the cumulative effect of
+    previous hunks.
+
+    The function assumes that the `start_line_old` values in the hunk headers are correct,
+    but the line counts and `start_line_new` may be incorrect due to manual edits or errors.
+
+    Parameters:
+        git_diff (str): The git diff string to be corrected. It may contain multiple hunks
+                        and incorrect hunk headers.
+
+    Returns:
+        str: The git diff string with corrected hunk headers.
+
+    The function works by parsing the diff, recalculating the line counts for each hunk,
+    adjusting the `start_line_new` values based on the cumulative changes from previous hunks,
+    and reconstructing the diff with corrected hunk headers.
+
+    Notes:
+        - The diff is expected to be in unified diff format.
+        - Lines starting with '+++', '---', or other file headers are not treated as additions or deletions.
+        - Context lines are lines that do not start with '+', '-', or '@@'.
+        - All lines in the hunk (including empty lines) are considered in the counts.
+
+    """
+    # Split the diff into lines
+    lines = git_diff.split('\n')
+
+    corrected_lines = []
+    line_index = 0
+
+    # Keep track of the cumulative line number adjustments
+    cumulative_line_offset = 0
+
+
+    while line_index < len(lines):
+        line = lines[line_index]
+        hunk_match = HUNK_HEADER_REGEX.match(line)
+
+        if hunk_match:
+            # Process the hunk
+            corrected_hunk_header, hunk_lines, line_index, line_offset = _process_hunk(
+                lines, line_index, cumulative_line_offset
+            )
+
+            # Update cumulative_line_offset
+            cumulative_line_offset += line_offset
+
+            # Append corrected hunk header and hunk lines
+            corrected_lines.append(corrected_hunk_header)
+            corrected_lines.extend(hunk_lines)
+        else:
+            # For non-hunk header lines, just append them
+            corrected_lines.append(line)
+            line_index += 1
+
+    # Join the corrected lines back into a string
+    corrected_diff = '\n'.join(corrected_lines)
+    return corrected_diff
+
+def _process_hunk(lines: list, start_index: int, cumulative_line_offset: int):
+    """
+    Processes a single hunk starting at start_index in lines, recalculates the line counts,
+    and returns the corrected hunk header, hunk lines, and the updated index after the hunk.
+
+    Parameters:
+        lines (list): The list of lines from the diff.
+        start_index (int): The index in lines where the hunk header is located.
+        cumulative_line_offset (int): The cumulative line offset from previous hunks.
+
+    Returns:
+        tuple:
+            - corrected_hunk_header (str): The corrected hunk header.
+            - hunk_lines (list): The list of lines in the hunk (excluding the hunk header).
+            - end_index (int): The index in lines after the hunk.
+            - line_offset (int): The line offset caused by this hunk (to adjust future hunks).
+
+    The function reads the hunk lines, counts the number of lines in the original and new files,
+    and adjusts the starting line number in the new file based on cumulative changes.
+    """
+
+    line_index = start_index
+    hunk_header_line = lines[line_index]
+    hunk_match = HUNK_HEADER_REGEX.match(hunk_header_line)
+
+    # Extract the starting line numbers from the hunk header
+    start_line_old = int(hunk_match.group(1))
+
+    # Initialize counts for recalculation
+    actual_count_old = 0
+    actual_count_new = 0
+
+    # Move to the next line after the hunk header
+    line_index += 1
+
+    hunk_lines = []
+
+    # Collect hunk lines until the next hunk header or end of diff
+    while line_index < len(lines):
+        hunk_line = lines[line_index]
+        if HUNK_HEADER_REGEX.match(hunk_line):
+            break
+        else:
+            hunk_lines.append(hunk_line)
+            line_index += 1
+
+    # Now process hunk_lines to calculate counts
+    for hunk_line in hunk_lines:
+        if hunk_line.startswith('+') and not hunk_line.startswith('+++'):
+            actual_count_new += 1
+        elif hunk_line.startswith('-') and not hunk_line.startswith('---'):
+            actual_count_old += 1
+        else:
+            # Context line (unchanged line)
+            actual_count_old += 1
+            actual_count_new += 1
+
+    # Adjust start_line_new considering previous line offsets
+    corrected_start_line_new = start_line_old + cumulative_line_offset
+
+    # Calculate line offset for subsequent hunks
+    line_offset = actual_count_new - actual_count_old
+
+    # Reconstruct the corrected hunk header
+    corrected_hunk_header = _format_hunk_header(
+        start_line_old, actual_count_old, corrected_start_line_new, actual_count_new
+    )
+
+    return corrected_hunk_header, hunk_lines, line_index, line_offset
+
+def _format_hunk_header(start_old: int, count_old: int, start_new: int, count_new: int) -> str:
+    """
+    Formats the hunk header according to git diff syntax, omitting counts when they are 1.
+
+    Parameters:
+        start_old (int): Starting line number in the original file.
+        count_old (int): Number of lines in the hunk in the original file.
+        start_new (int): Starting line number in the new file.
+        count_new (int): Number of lines in the hunk in the new file.
+
+    Returns:
+        str: The formatted hunk header.
+
+    The hunk header format is:
+        @@ -start_old[,count_old] +start_new[,count_new] @@
+
+    If count_old or count_new is 1, the count is omitted.
+    """
+    # Omit counts when they are equal to 1
+    old_part = f'-{start_old}'
+    if count_old != 1:
+        old_part += f',{count_old}'
+    new_part = f'+{start_new}'
+    if count_new != 1:
+        new_part += f',{count_new}'
+    return f'@@ {old_part} {new_part} @@'
+
+
+
+if __name__ == '__main__':
+    # TODO: Create unit test and move these code to unit test
+    diff = """\
+diff --git a/file.txt b/file.txt
+index e69de29..4b825dc 100644
+--- a/file.txt
++++ b/file.txt
+@@ -1,5 +1,5 @@
+ Line one
++Line two
+ Line three
+ Line four
+ Line five
+@@ -10,2 +10,2 @@
+ Line ten
+-Line eleven
++Line eleven modified"""
+    print(correct_git_diff(diff))

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ziya"
-version = "0.1.41"
+version = "0.1.42"
 description = ""
 authors = ["Vishnu Krishnaprasad <vishnukool@gmail.com>"]
 readme = "README.md"
@@ -18,6 +18,7 @@ tiktoken = "^0.8.0"
 boto3 = "^1.34.88"
 langchain-aws = ">=0.2,<0.3"
 langchain = ">=0.3,<0.4"
+langgraph = ">=0.2,<0.3"
 langchainhub = ">=0.1.15"
 langchain-anthropic = ">=0.2,<0.3"
 langchain-cli = ">=0.0.15"

templates/asset-manifest.json

@@ -1,7 +1,7 @@
 {
   "files": {
     "main.css": "/static/css/main.8af23da0.css",
-    "main.js": "/static/js/main.69f6f0d1.js",
+    "main.js": "/static/js/main.fee8aad7.js",
     "static/media/fa-solid-900.ttf": "/static/media/fa-solid-900.bacd5de623fb563b961a.ttf",
     "static/media/fa-brands-400.ttf": "/static/media/fa-brands-400.60127e352b7a11f7f1bc.ttf",
     "static/media/fa-solid-900.woff2": "/static/media/fa-solid-900.4d986b00ff9ca3828fbd.woff2",
@@ -12,10 +12,10 @@
     "static/media/fa-v4compatibility.woff2": "/static/media/fa-v4compatibility.cf7f5903d06b79ad60f1.woff2",
     "index.html": "/index.html",
     "main.8af23da0.css.map": "/static/css/main.8af23da0.css.map",
-    "main.69f6f0d1.js.map": "/static/js/main.69f6f0d1.js.map"
+    "main.fee8aad7.js.map": "/static/js/main.fee8aad7.js.map"
   },
   "entrypoints": [
     "static/css/main.8af23da0.css",
-    "static/js/main.69f6f0d1.js"
+    "static/js/main.fee8aad7.js"
   ]
 }

templates/index.html

@@ -1 +1 @@
-<!doctype html><html lang="en"><head><meta charset="UTF-8"><meta name="viewport" content="width=device-width,initial-scale=1"><title>Ziya - Code Assistant</title><link rel="icon" href="/favicon.ico" type="image/x-icon"><script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script><script defer="defer" src="/static/js/main.69f6f0d1.js"></script><link href="/static/css/main.8af23da0.css" rel="stylesheet"></head><body><div id="root"></div></body></html>
+<!doctype html><html lang="en"><head><meta charset="UTF-8"><meta name="viewport" content="width=device-width,initial-scale=1"><title>Ziya - Code Assistant</title><link rel="icon" href="/favicon.ico" type="image/x-icon"><script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script><script>window.enableCodeApply="{{ enable_code_apply | lower }}"</script><script defer="defer" src="/static/js/main.fee8aad7.js"></script><link href="/static/css/main.8af23da0.css" rel="stylesheet"></head><body><div id="root"></div></body></html>