cognify-code 0.2.0__py3-none-any.whl

ai_code_assistant/generator/code_gen.py

@@ -0,0 +1,265 @@
+"""Code generator for creating code from specifications."""
+
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import List, Literal, Optional
+
+from ai_code_assistant.config import Config
+from ai_code_assistant.llm import LLMManager
+from ai_code_assistant.generator.prompts import GENERATION_PROMPTS
+
+
+GenerationMode = Literal["function", "class", "script", "test", "generic"]
+
+
+@dataclass
+class GenerationResult:
+    """Result of code generation."""
+    code: str
+    language: str
+    mode: GenerationMode
+    description: str
+    raw_response: str = ""
+    error: Optional[str] = None
+
+    @property
+    def success(self) -> bool:
+        return self.error is None and bool(self.code.strip())
+
+
+class CodeGenerator:
+    """Generates code from natural language descriptions."""
+
+    def __init__(self, config: Config, llm_manager: LLMManager):
+        self.config = config
+        self.llm = llm_manager
+
+    def generate_function(
+        self,
+        description: str,
+        name: str,
+        language: str = "python",
+        parameters: str = "",
+        return_type: str = "None",
+    ) -> GenerationResult:
+        """Generate a function from specification."""
+        prompt = GENERATION_PROMPTS["function"]
+        
+        try:
+            response = self.llm.invoke_with_template(
+                prompt,
+                language=language,
+                description=description,
+                name=name,
+                parameters=parameters or "None specified",
+                return_type=return_type,
+                include_type_hints=self.config.generation.include_type_hints,
+                include_docstrings=self.config.generation.include_docstrings,
+            )
+            code = self._extract_code(response, language)
+            return GenerationResult(
+                code=code,
+                language=language,
+                mode="function",
+                description=description,
+                raw_response=response,
+            )
+        except Exception as e:
+            return GenerationResult(
+                code="",
+                language=language,
+                mode="function",
+                description=description,
+                error=str(e),
+            )
+
+    def generate_class(
+        self,
+        description: str,
+        name: str,
+        language: str = "python",
+        attributes: str = "",
+        methods: str = "",
+    ) -> GenerationResult:
+        """Generate a class from specification."""
+        prompt = GENERATION_PROMPTS["class"]
+        
+        try:
+            response = self.llm.invoke_with_template(
+                prompt,
+                language=language,
+                description=description,
+                name=name,
+                attributes=attributes or "None specified",
+                methods=methods or "None specified",
+                include_type_hints=self.config.generation.include_type_hints,
+                include_docstrings=self.config.generation.include_docstrings,
+            )
+            code = self._extract_code(response, language)
+            return GenerationResult(
+                code=code,
+                language=language,
+                mode="class",
+                description=description,
+                raw_response=response,
+            )
+        except Exception as e:
+            return GenerationResult(
+                code="",
+                language=language,
+                mode="class",
+                description=description,
+                error=str(e),
+            )
+
+    def generate_script(
+        self,
+        description: str,
+        requirements: List[str],
+        language: str = "python",
+    ) -> GenerationResult:
+        """Generate a complete script or module."""
+        prompt = GENERATION_PROMPTS["script"]
+        
+        try:
+            response = self.llm.invoke_with_template(
+                prompt,
+                language=language,
+                description=description,
+                requirements="\n".join(f"- {r}" for r in requirements),
+                include_type_hints=self.config.generation.include_type_hints,
+                include_docstrings=self.config.generation.include_docstrings,
+            )
+            code = self._extract_code(response, language)
+            return GenerationResult(
+                code=code,
+                language=language,
+                mode="script",
+                description=description,
+                raw_response=response,
+            )
+        except Exception as e:
+            return GenerationResult(
+                code="",
+                language=language,
+                mode="script",
+                description=description,
+                error=str(e),
+            )
+
+    def generate_tests(
+        self,
+        source_code: str,
+        language: str = "python",
+        test_framework: str = "pytest",
+    ) -> GenerationResult:
+        """Generate tests for existing code."""
+        prompt = GENERATION_PROMPTS["test"]
+        
+        try:
+            response = self.llm.invoke_with_template(
+                prompt,
+                language=language,
+                source_code=source_code,
+                test_framework=test_framework,
+            )
+            code = self._extract_code(response, language)
+            return GenerationResult(
+                code=code,
+                language=language,
+                mode="test",
+                description=f"Tests for provided {language} code",
+                raw_response=response,
+            )
+        except Exception as e:
+            return GenerationResult(
+                code="",
+                language=language,
+                mode="test",
+                description="Test generation failed",
+                error=str(e),
+            )
+
+    def generate(
+        self,
+        description: str,
+        language: str = "python",
+    ) -> GenerationResult:
+        """Generate code from a generic description."""
+        prompt = GENERATION_PROMPTS["generic"]
+        
+        try:
+            response = self.llm.invoke_with_template(
+                prompt,
+                language=language,
+                description=description,
+                include_type_hints=self.config.generation.include_type_hints,
+                include_docstrings=self.config.generation.include_docstrings,
+            )
+            code = self._extract_code(response, language)
+            return GenerationResult(
+                code=code,
+                language=language,
+                mode="generic",
+                description=description,
+                raw_response=response,
+            )
+        except Exception as e:
+            return GenerationResult(
+                code="",
+                language=language,
+                mode="generic",
+                description=description,
+                error=str(e),
+            )
+
+    def _extract_code(self, response: str, language: str) -> str:
+        """Extract code block from LLM response."""
+        # Try to find language-specific code block
+        pattern = rf"```{language}\s*\n(.*?)```"
+        match = re.search(pattern, response, re.DOTALL | re.IGNORECASE)
+        if match:
+            return match.group(1).strip()
+
+        # Try generic code block
+        pattern = r"```\s*\n(.*?)```"
+        match = re.search(pattern, response, re.DOTALL)
+        if match:
+            return match.group(1).strip()
+
+        # Try any code block
+        pattern = r"```\w*\s*\n(.*?)```"
+        match = re.search(pattern, response, re.DOTALL)
+        if match:
+            return match.group(1).strip()
+
+        # Look for code that starts with common patterns
+        lines = response.strip().split('\n')
+        code_lines = []
+        in_code = False
+
+        for line in lines:
+            # Detect start of code
+            if not in_code:
+                if line.strip().startswith(('def ', 'class ', 'import ', 'from ', '#!', 'function ', 'const ', 'let ', 'var ')):
+                    in_code = True
+
+            if in_code:
+                code_lines.append(line)
+
+        if code_lines:
+            return '\n'.join(code_lines).strip()
+
+        # Return raw response if no code block found
+        return response.strip()
+
+    def save_to_file(self, result: GenerationResult, output_path: Path) -> bool:
+        """Save generated code to a file."""
+        try:
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            output_path.write_text(result.code)
+            return True
+        except Exception:
+            return False
+

ai_code_assistant/generator/prompts.py

@@ -0,0 +1,114 @@
+"""Prompt templates for code generation."""
+
+from langchain_core.prompts import ChatPromptTemplate
+
+# System prompt for code generation
+GENERATION_SYSTEM_PROMPT = """You are an expert software developer who writes clean, efficient, and well-documented code.
+
+When generating code:
+1. Follow best practices and design patterns for the language
+2. Include comprehensive docstrings and comments
+3. Add type hints where applicable
+4. Handle edge cases and errors appropriately
+5. Write code that is readable and maintainable
+
+Output ONLY the code, wrapped in appropriate markdown code blocks.
+Do not include explanations unless specifically requested."""
+
+# Function generation prompt
+FUNCTION_PROMPT = ChatPromptTemplate.from_messages([
+    ("system", GENERATION_SYSTEM_PROMPT),
+    ("human", """Generate a {language} function with the following specification:
+
+**Description:** {description}
+**Function name:** {name}
+**Parameters:** {parameters}
+**Return type:** {return_type}
+**Additional requirements:**
+- Include type hints: {include_type_hints}
+- Include docstring: {include_docstrings}
+
+Generate only the function code, no usage examples.""")
+])
+
+# Class generation prompt
+CLASS_PROMPT = ChatPromptTemplate.from_messages([
+    ("system", GENERATION_SYSTEM_PROMPT),
+    ("human", """Generate a {language} class with the following specification:
+
+**Description:** {description}
+**Class name:** {name}
+**Attributes:** {attributes}
+**Methods:** {methods}
+**Additional requirements:**
+- Include type hints: {include_type_hints}
+- Include docstrings: {include_docstrings}
+- Follow {language} best practices and conventions
+
+Generate the complete class implementation.""")
+])
+
+# Script/module generation prompt
+SCRIPT_PROMPT = ChatPromptTemplate.from_messages([
+    ("system", GENERATION_SYSTEM_PROMPT),
+    ("human", """Generate a complete {language} script/module:
+
+**Description:** {description}
+**Features/functionality required:**
+{requirements}
+
+**Additional requirements:**
+- Include proper imports
+- Include type hints: {include_type_hints}
+- Include docstrings: {include_docstrings}
+- Include a main entry point if applicable
+- Handle errors appropriately
+
+Generate the complete script.""")
+])
+
+# Test file generation prompt
+TEST_PROMPT = ChatPromptTemplate.from_messages([
+    ("system", """You are an expert at writing comprehensive test suites.
+Write tests that cover happy paths, edge cases, and error conditions.
+Use appropriate testing frameworks and assertions."""),
+    ("human", """Generate {language} tests for the following code:
+
+```{language}
+{source_code}
+```
+
+**Test framework:** {test_framework}
+**Coverage requirements:**
+- Test all public functions/methods
+- Include edge case tests
+- Include error handling tests
+- Use descriptive test names
+
+Generate a complete test file.""")
+])
+
+# Generic code generation prompt
+GENERIC_PROMPT = ChatPromptTemplate.from_messages([
+    ("system", GENERATION_SYSTEM_PROMPT),
+    ("human", """Generate {language} code for the following request:
+
+{description}
+
+Requirements:
+- Include type hints: {include_type_hints}
+- Include docstrings: {include_docstrings}
+- Follow best practices
+
+Generate clean, production-ready code.""")
+])
+
+# Collect all prompts
+GENERATION_PROMPTS = {
+    "function": FUNCTION_PROMPT,
+    "class": CLASS_PROMPT,
+    "script": SCRIPT_PROMPT,
+    "test": TEST_PROMPT,
+    "generic": GENERIC_PROMPT,
+}
+

ai_code_assistant/git/__init__.py

@@ -0,0 +1,6 @@
+"""Git integration module for AI Code Assistant."""
+
+from ai_code_assistant.git.manager import GitManager
+from ai_code_assistant.git.commit_generator import CommitMessageGenerator
+
+__all__ = ["GitManager", "CommitMessageGenerator"]

ai_code_assistant/git/commit_generator.py

@@ -0,0 +1,130 @@
+"""AI-powered commit message generator."""
+
+from typing import Optional
+
+from ai_code_assistant.git.manager import GitManager, GitDiff, GitStatus
+
+
+COMMIT_MESSAGE_PROMPT = """Analyze the following git diff and generate a concise, descriptive commit message.
+
+Follow these conventions:
+1. Start with a type prefix: feat, fix, docs, style, refactor, test, chore
+2. Use imperative mood ("Add feature" not "Added feature")
+3. Keep the first line under 72 characters
+4. If needed, add a blank line and bullet points for details
+
+Git Status:
+- Branch: {branch}
+- Files changed: {files_changed}
+- Insertions: {insertions}
+- Deletions: {deletions}
+
+Changed files:
+{changed_files}
+
+Diff (truncated to 3000 chars):
+```
+{diff}
+```
+
+Generate ONLY the commit message, nothing else. No explanations, no markdown formatting around it.
+"""
+
+
+class CommitMessageGenerator:
+    """Generates commit messages using AI."""
+    
+    def __init__(self, llm_manager):
+        self.llm = llm_manager
+    
+    def generate(self, git_manager: GitManager, 
+                 staged_only: bool = True,
+                 style: str = "conventional") -> str:
+        """Generate a commit message based on the current changes.
+        
+        Args:
+            git_manager: GitManager instance
+            staged_only: If True, only consider staged changes
+            style: Commit message style (conventional, simple)
+            
+        Returns:
+            Generated commit message
+        """
+        status = git_manager.get_status()
+        diff = git_manager.get_diff(staged=staged_only)
+        
+        if not diff.diff_text and not status.untracked:
+            return ""
+        
+        # Build list of changed files
+        changed_files = []
+        if staged_only:
+            changed_files = status.staged
+        else:
+            changed_files = status.staged + status.modified + status.untracked
+        
+        changed_files_str = "\n".join(f"  - {f}" for f in changed_files[:20])
+        if len(changed_files) > 20:
+            changed_files_str += f"\n  ... and {len(changed_files) - 20} more files"
+        
+        # Truncate diff to avoid token limits
+        diff_text = diff.diff_text[:3000]
+        if len(diff.diff_text) > 3000:
+            diff_text += "\n... (truncated)"
+        
+        prompt = COMMIT_MESSAGE_PROMPT.format(
+            branch=status.branch,
+            files_changed=diff.files_changed or len(changed_files),
+            insertions=diff.insertions,
+            deletions=diff.deletions,
+            changed_files=changed_files_str,
+            diff=diff_text,
+        )
+        
+        # Generate message
+        response = self.llm.invoke(prompt)
+        
+        # Clean up response
+        message = response.strip()
+        # Remove any markdown code blocks if present
+        if message.startswith("```"):
+            lines = message.split("\n")
+            message = "\n".join(lines[1:-1] if lines[-1] == "```" else lines[1:])
+        
+        return message.strip()
+    
+    def generate_from_description(self, description: str, 
+                                   git_manager: Optional[GitManager] = None) -> str:
+        """Generate a commit message from a description.
+        
+        Args:
+            description: User's description of changes
+            git_manager: Optional GitManager for context
+            
+        Returns:
+            Generated commit message
+        """
+        context = ""
+        if git_manager:
+            status = git_manager.get_status()
+            diff = git_manager.get_diff(staged=True)
+            context = f"""
+Context:
+- Branch: {status.branch}
+- Files changed: {diff.files_changed}
+- Staged files: {', '.join(status.staged[:5])}
+"""
+        
+        prompt = f"""Convert this description into a proper git commit message.
+
+Follow conventional commits format:
+- Start with type: feat, fix, docs, style, refactor, test, chore
+- Use imperative mood
+- Keep first line under 72 characters
+{context}
+Description: {description}
+
+Generate ONLY the commit message, nothing else:"""
+        
+        response = self.llm.invoke(prompt)
+        return response.strip()