Merge pull request #1 from setkyar/feature/async

setkyar · web-flow · commit f8eb995db68f · 2025-02-13T22:50:50.000+08:00
Add pytest-asyncio dependency and refactor review process to support …
diff --git a/aireview/ai_reviewer.py b/aireview/ai_reviewer.py
@@ -1,7 +1,8 @@
 """Module for handling AI review generation."""
 import click
+import asyncio
 from dataclasses import dataclass
-from openai import OpenAI
+from openai import AsyncOpenAI
 from typing import List, Optional
 from .git_handler import FileChange
 
@@ -13,44 +14,74 @@ class Review:
 
 class AIReviewer:
     def __init__(self, model: str, api_key: str, base_url: Optional[str] = None):
-        self.client = OpenAI(
+        self.client = AsyncOpenAI(
             api_key=api_key,
             base_url=base_url if base_url else None
         )
         self.model = model
     
-    def review_changes(self, changes: List[FileChange], 
-                      project_context: str, prompt_template: str) -> List[Review]:
-        """Generate AI reviews for all file changes."""
-        reviews = []
+    async def review_changes(self, changes: List[FileChange], 
+                           project_context: str, prompt_template: str) -> List[Review]:
+        """Generate AI reviews for all file changes in parallel."""
+        click.echo(f"Generating reviews for {len(changes)} files...")
+        
+        # Create tasks for all reviews
+        tasks = []
         for change in changes:
-            click.echo(f"Generating review for {change.filename}...")
+            click.echo(f"Starting review for {change.filename}...")
             
-            review_content = self._get_review(
-                change.content,
-                change.filename,
-                project_context,
-                prompt_template
+            # Create prompt with filename included
+            prompt = self._create_prompt(
+                changes=change.content,
+                filename=change.filename,
+                file_content=change.file_content,
+                project_context=project_context,
+                prompt_template=prompt_template
             )
-            reviews.append(Review(filename=change.filename, content=review_content))
+            
+            # Create task for this review
+            tasks.append(self._get_review(prompt, change.filename))
+        
+        # Run all reviews concurrently
+        review_contents = await asyncio.gather(*tasks)
+        
+        # Create Review objects from results
+        reviews = [
+            Review(filename=changes[i].filename, content=content)
+            for i, content in enumerate(review_contents)
+        ]
+        
         return reviews
     
-    def _get_review(self, changes: str, filename: str, 
-                    project_context: str, prompt_template: str) -> str:
-        """Get AI review for a single file's changes."""
-        prompt = f"""{project_context}
+    def _create_prompt(self, changes: str, filename: str,
+                      file_content: Optional[str], project_context: str,
+                      prompt_template: str) -> str:
+        """Create the prompt for a single file review."""
+        # Include file content in the prompt if available
+        file_content_section = ""
+        if file_content:
+            file_content_section = f"""
+            Current file content:
+            ```
+            {file_content}
+            ```
+            """
+
+        return f"""{project_context}
 
         {prompt_template}
         
         Review the following changes in {filename}:
         ```
         {changes}
         ```
-        
+        {file_content_section}
         Please focus your review on these specific changes."""
-
+    
+    async def _get_review(self, prompt: str, filename: str) -> str:
+        """Get AI review for the provided prompt."""
         try:
-            completion = self.client.chat.completions.create(
+            completion = await self.client.chat.completions.create(
                 model=self.model,
                 n=1,
                 messages=[
@@ -59,6 +90,7 @@ def _get_review(self, changes: str, filename: str,
                 ]
             )
             
+            click.echo(f"Completed review for {filename}")
             return f"## Review for changes in {filename}\n\n{completion.choices[0].message.content}"
         except Exception as e:
             raise RuntimeError(f"OpenAI API error for {filename}: {str(e)}")
diff --git a/aireview/git_handler.py b/aireview/git_handler.py
@@ -1,18 +1,20 @@
 """Module for handling Git operations."""
 import subprocess
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import List, Optional, Dict
+import os
 
 @dataclass
 class FileChange:
     """Represents changes in a single file."""
     filename: str
     content: str
+    file_content: Optional[str] = None
 
 class GitHandler:
     @staticmethod
     def get_file_changes() -> List[FileChange]:
-        """Retrieves only staged changes from Git (after git add)."""
+        """Retrieves staged changes from Git and their corresponding file content efficiently."""
         try:
             # Get staged changes
             staged_cmd = subprocess.run(
@@ -22,12 +24,115 @@ def get_file_changes() -> List[FileChange]:
             
             if not staged_cmd.stdout:
                 return []
-                
-            return GitHandler._parse_diff_output(staged_cmd.stdout)
+            
+            # Parse the diff output first
+            changes = GitHandler._parse_diff_output(staged_cmd.stdout)
+            
+            # Get the list of files we need content for
+            files_to_fetch = [change.filename for change in changes]
+            
+            # Batch fetch file contents
+            file_contents = GitHandler._batch_get_file_contents(files_to_fetch)
+            
+            # Update FileChange objects with their content
+            for change in changes:
+                change.file_content = file_contents.get(change.filename)
+            
+            return changes
             
         except subprocess.CalledProcessError as e:
             raise RuntimeError(f"Git command failed: {e.stderr}")
     
+    @staticmethod
+    def _batch_get_file_contents(filenames: List[str]) -> Dict[str, Optional[str]]:
+        """
+        Efficiently get contents of multiple files using git cat-file --batch.
+        Returns a dictionary mapping filenames to their content.
+        """
+        if not filenames:
+            return {}
+        
+        try:
+            # Get object IDs for staged versions of files
+            file_revs = {}
+            for filename in filenames:
+                try:
+                    rev_cmd = subprocess.run(
+                        ['git', 'rev-parse', f':"{filename}"'],
+                        capture_output=True, text=True, check=True
+                    )
+                    file_revs[filename] = rev_cmd.stdout.strip()
+                except subprocess.CalledProcessError:
+                    # File might be new/deleted
+                    file_revs[filename] = None
+            
+            # Prepare batch input
+            valid_revs = {f: rev for f, rev in file_revs.items() if rev is not None}
+            if not valid_revs:
+                return {f: None for f in filenames}
+            
+            # Start git cat-file --batch process
+            process = subprocess.Popen(
+                ['git', 'cat-file', '--batch'],
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE
+            )
+            
+            # Write object IDs to git cat-file
+            input_data = '\n'.join(valid_revs.values()) + '\n'
+            stdout, stderr = process.communicate(input_data.encode())
+            
+            if process.returncode != 0:
+                raise subprocess.CalledProcessError(
+                    process.returncode, 'git cat-file', stderr
+                )
+            
+            # Parse the output
+            contents = {}
+            current_content = []
+            current_file = None
+            rev_to_file = {rev: f for f, rev in valid_revs.items()}
+            
+            for line in stdout.decode().split('\n'):
+                if line.strip() and ' blob ' in line:
+                    # New blob header - save previous content if any
+                    if current_file and current_content:
+                        contents[current_file] = ''.join(current_content)
+                        current_content = []
+                    
+                    # Get filename for this blob
+                    obj_id = line.split()[0]
+                    current_file = rev_to_file.get(obj_id)
+                else:
+                    current_content.append(line + '\n')
+            
+            # Save last file's content
+            if current_file and current_content:
+                contents[current_file] = ''.join(current_content)
+            
+            # Include None for files that weren't found
+            return {f: contents.get(f) for f in filenames}
+            
+        except Exception as e:
+            # If batch operation fails, fall back to individual git show commands
+            return GitHandler._fallback_get_file_contents(filenames)
+    
+    @staticmethod
+    def _fallback_get_file_contents(filenames: List[str]) -> Dict[str, Optional[str]]:
+        """Fallback method to get file contents using git show."""
+        contents = {}
+        for filename in filenames:
+            try:
+                show_cmd = subprocess.run(
+                    ['git', 'show', f':{filename}'],
+                    capture_output=True, text=True, check=True
+                )
+                contents[filename] = show_cmd.stdout
+            except subprocess.CalledProcessError:
+                contents[filename] = None
+        return contents
+    
     @staticmethod
     def _parse_diff_output(diff_output: str) -> List[FileChange]:
         """Parse git diff output into FileChange objects."""
diff --git a/aireview/main.py b/aireview/main.py
@@ -1,6 +1,7 @@
 """Main module for the AI code review tool."""
 import click
 import logging
+import asyncio
 from typing import List
 from .config import ConfigLoader
 from .git_handler import GitHandler
@@ -47,11 +48,12 @@ def main(config: str):
             base_url=ai_config.base_url
         )
         
-        reviews = reviewer.review_changes(
+        # Run the async review process
+        reviews = asyncio.run(reviewer.review_changes(
             file_changes,
             review_config.project_context,
             review_config.prompt_template
-        )
+        ))
         
         # Write output
         write_reviews(reviews, review_config.output_file)
@@ -61,4 +63,7 @@ def main(config: str):
     except Exception as e:
         click.echo(f"Error: {str(e)}", err=True)
         logging.error(f"Error: {str(e)}")
-        return
+        return
+
+if __name__ == '__main__':
+    main()
diff --git a/setup.py b/setup.py
@@ -84,6 +84,7 @@ def get_version_from_git():
             'pytest>=7.0.0,<8.0.0',
             'pytest-cov>=4.1.0,<5.0.0',
             'pytest-mock>=3.10.0',
+            'pytest-asyncio>=0.23.0',
         ],
     },
 )
diff --git a/tests/test_ai_reviewer.py b/tests/test_ai_reviewer.py
@@ -1,45 +1,87 @@
 import pytest
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, patch, AsyncMock
 from aireview.ai_reviewer import AIReviewer
 from aireview.git_handler import FileChange
 
 @pytest.fixture
 def mock_openai():
-    """Mock OpenAI client responses."""
-    with patch('aireview.ai_reviewer.OpenAI') as mock:
+    """Mock AsyncOpenAI client responses."""
+    with patch('aireview.ai_reviewer.AsyncOpenAI') as mock:
         mock_client = Mock()
-        mock_client.chat.completions.create.return_value = Mock(
-            choices=[Mock(message=Mock(content="Test review content"))]
+        # Use AsyncMock for async methods
+        mock_client.chat = Mock()
+        mock_client.chat.completions = Mock()
+        mock_client.chat.completions.create = AsyncMock(
+            return_value=Mock(
+                choices=[Mock(message=Mock(content="Mock review content"))]
+            )
         )
         mock.return_value = mock_client
         yield mock
 
-def test_review_changes(mock_openai):
-    """Test AI review generation."""
+@pytest.mark.asyncio
+async def test_review_changes_with_file_content(mock_openai):
+    """Test AI review generation with file content."""
     reviewer = AIReviewer("test-model", "test-key")
     changes = [
         FileChange(
             filename="test.py",
-            content="Added: print('hello world')"
+            content="Added: print('hello world')",
+            file_content="print('hello')\nprint('hello world')"
         )
     ]
     
-    reviews = reviewer.review_changes(
+    reviews = await reviewer.review_changes(
         changes,
         project_context="Test context",
         prompt_template="Test template"
     )
     
+    # Verify the review was created
     assert len(reviews) == 1
     assert reviews[0].filename == "test.py"
-    assert "Test review content" in reviews[0].content
+    
+    # Verify the prompt sent to the API includes file content
+    api_call_args = mock_openai.return_value.chat.completions.create.call_args
+    prompt_sent = api_call_args[1]['messages'][1]['content']
+    assert "Current file content:" in prompt_sent
+    assert "print('hello')" in prompt_sent
+
+@pytest.mark.asyncio
+async def test_review_changes_without_file_content(mock_openai):
+    """Test AI review generation without file content."""
+    reviewer = AIReviewer("test-model", "test-key")
+    changes = [
+        FileChange(
+            filename="test.py",
+            content="Added: print('hello world')",
+            file_content=None  # No file content
+        )
+    ]
+    
+    reviews = await reviewer.review_changes(
+        changes,
+        project_context="Test context",
+        prompt_template="Test template"
+    )
+    
+    # Verify the review was created
+    assert len(reviews) == 1
+    assert reviews[0].filename == "test.py"
+    
+    # Verify the prompt sent to the API doesn't include file content section
+    api_call_args = mock_openai.return_value.chat.completions.create.call_args
+    prompt_sent = api_call_args[1]['messages'][1]['content']
+    assert "Current file content:" not in prompt_sent
+    assert changes[0].content in prompt_sent
 
-def test_review_changes_api_error(mock_openai):
+@pytest.mark.asyncio
+async def test_review_changes_api_error(mock_openai):
     """Test handling of API errors during review."""
     mock_openai.return_value.chat.completions.create.side_effect = Exception("API Error")
     
     reviewer = AIReviewer("test-model", "test-key")
     changes = [FileChange(filename="test.py", content="test content")]
     
     with pytest.raises(RuntimeError, match="OpenAI API error"):
-        reviewer.review_changes(changes, "", "")
+        await reviewer.review_changes(changes, "", "")
diff --git a/tests/test_git_handler.py b/tests/test_git_handler.py
diff --git a/tests/test_main.py b/tests/test_main.py

Original file line number	Diff line number	Diff line change
`@@ -84,6 +84,7 @@ def get_version_from_git():`
`84`	`84`	`'pytest>=7.0.0,<8.0.0',`
`85`	`85`	`'pytest-cov>=4.1.0,<5.0.0',`
`86`	`86`	`'pytest-mock>=3.10.0',`
	`87`	`+ 'pytest-asyncio>=0.23.0',`
`87`	`88`	`],`
`88`	`89`	`},`
`89`	`90`	`)`