akitallm 0.1.1__py3-none-any.whl → 1.1.0__py3-none-any.whl

akita/core/providers.py

@@ -0,0 +1,181 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+import requests
+from pydantic import BaseModel
+
+class ModelInfo(BaseModel):
+    id: str
+    name: Optional[str] = None
+
+class BaseProvider(ABC):
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        pass
+
+    @abstractmethod
+    def validate_key(self, api_key: str) -> bool:
+        pass
+
+    @abstractmethod
+    def list_models(self, api_key: str) -> List[ModelInfo]:
+        pass
+
+class OpenAIProvider(BaseProvider):
+    @property
+    def name(self) -> str:
+        return "openai"
+
+    def validate_key(self, api_key: str) -> bool:
+        if not api_key.startswith("sk-"):
+            return False
+        # Simple validation request
+        try:
+            response = requests.get(
+                "https://api.openai.com/v1/models",
+                headers={"Authorization": f"Bearer {api_key}"},
+                timeout=5
+            )
+            return response.status_code == 200
+        except Exception:
+            return False
+
+    def list_models(self, api_key: str) -> List[ModelInfo]:
+        response = requests.get(
+            "https://api.openai.com/v1/models",
+            headers={"Authorization": f"Bearer {api_key}"},
+            timeout=10
+        )
+        response.raise_for_status()
+        data = response.json()
+        exclude_keywords = ["vision", "instruct", "audio", "realtime", "tts", "dall-e", "embedding", "moderation", "davinci", "babbage", "curie", "ada"]
+        
+        models = []
+        for m in data["data"]:
+            model_id = m["id"]
+            if not any(kw in model_id.lower() for kw in exclude_keywords):
+                if model_id.startswith("gpt-") or model_id.startswith("o1") or model_id.startswith("o3"):
+                    models.append(ModelInfo(id=model_id))
+        return sorted(models, key=lambda x: x.id)
+
+class AnthropicProvider(BaseProvider):
+    @property
+    def name(self) -> str:
+        return "anthropic"
+
+    def validate_key(self, api_key: str) -> bool:
+        if not api_key.startswith("sk-ant-"):
+            return False
+        # Anthropic validation usually requires a full request, but we'll check prefix for now
+        # or do a no-op call if possible.
+        return True
+
+    def list_models(self, api_key: str) -> List[ModelInfo]:
+        # Anthropic doesn't have a public models list API like OpenAI
+        return [
+            ModelInfo(id="claude-3-5-sonnet-latest", name="Claude 3.5 Sonnet (Latest)"),
+            ModelInfo(id="claude-3-5-haiku-latest", name="Claude 3.5 Haiku (Latest)"),
+            ModelInfo(id="claude-3-opus-20240229", name="Claude 3 Opus"),
+            ModelInfo(id="claude-3-sonnet-20240229", name="Claude 3 Sonnet"),
+            ModelInfo(id="claude-3-haiku-20240307", name="Claude 3 Haiku"),
+        ]
+
+class OllamaProvider(BaseProvider):
+    @property
+    def name(self) -> str:
+        return "ollama"
+
+    def validate_key(self, api_key: str) -> bool:
+        # Ollama doesn't use keys by default, we just check if it's reachable
+        try:
+            response = requests.get("http://localhost:11434/api/tags", timeout=2)
+            return response.status_code == 200
+        except Exception:
+            return False
+
+    def list_models(self, api_key: str) -> List[ModelInfo]:
+        response = requests.get("http://localhost:11434/api/tags", timeout=5)
+        response.raise_for_status()
+        data = response.json()
+        return [ModelInfo(id=m["name"]) for m in data["models"]]
+
+class GeminiProvider(BaseProvider):
+    @property
+    def name(self) -> str:
+        return "gemini"
+
+    def validate_key(self, api_key: str) -> bool:
+        if not api_key.startswith("AIza"):
+            return False
+        return True
+
+    def list_models(self, api_key: str) -> List[ModelInfo]:
+        # Gemini API URL for listing models
+        url = f"https://generativelanguage.googleapis.com/v1beta/models?key={api_key}"
+        response = requests.get(url, timeout=10)
+        response.raise_for_status()
+        data = response.json()
+        
+        exclude_keywords = ["nano", "banana", "vision", "embedding", "aqa", "learnlm"]
+        
+        models = []
+        for m in data["models"]:
+            model_id = m["name"].split("/")[-1]
+            display_name = m["displayName"]
+            
+            # Check if it supports generation and doesn't have excluded keywords
+            if "generateContent" in m["supportedGenerationMethods"]:
+                if not any(kw in model_id.lower() or kw in display_name.lower() for kw in exclude_keywords):
+                    models.append(ModelInfo(id=model_id, name=display_name))
+        
+        return models
+
+class GroqProvider(BaseProvider):
+    @property
+    def name(self) -> str:
+        return "groq"
+
+    def validate_key(self, api_key: str) -> bool:
+        if not api_key.startswith("gsk_"):
+            return False
+        return True
+
+    def list_models(self, api_key: str) -> List[ModelInfo]:
+        # Groq uses OpenAI-compatible models endpoint
+        response = requests.get(
+            "https://api.groq.com/openai/v1/models",
+            headers={"Authorization": f"Bearer {api_key}"},
+            timeout=10
+        )
+        response.raise_for_status()
+        data = response.json()
+        
+        # Filter for text models
+        exclude_keywords = ["vision", "audio"]
+        models = []
+        for m in data["data"]:
+            model_id = m["id"]
+            if not any(kw in model_id.lower() for kw in exclude_keywords):
+                models.append(ModelInfo(id=model_id))
+        return sorted(models, key=lambda x: x.id)
+
+def detect_provider(api_key: str) -> Optional[BaseProvider]:
+    """
+    Attempts to detect the provider based on the API key or environment.
+    """
+    if api_key.lower() == "ollama":
+        return OllamaProvider()
+    
+    if api_key.startswith("sk-ant-"):
+        return AnthropicProvider()
+    
+    if api_key.startswith("gsk_"):
+        return GroqProvider()
+    
+    if api_key.startswith("sk-"):
+        return OpenAIProvider()
+    
+    if api_key.startswith("AIza"):
+        return GeminiProvider()
+        
+    return None

akita/core/trace.py

@@ -0,0 +1,18 @@
+from typing import List, Dict, Any
+from datetime import datetime
+from pydantic import BaseModel, Field
+
+class TraceStep(BaseModel):
+    timestamp: datetime = Field(default_factory=datetime.now)
+    action: str
+    details: str
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+class ReasoningTrace(BaseModel):
+    steps: List[TraceStep] = Field(default_factory=list)
+
+    def add_step(self, action: str, details: str, metadata: Dict[str, Any] = None):
+        self.steps.append(TraceStep(action=action, details=details, metadata=metadata or {}))
+
+    def __str__(self):
+        return "\n".join([f"[{s.timestamp.strftime('%H:%M:%S')}] {s.action}: {s.details}" for s in self.steps])

akita/models/base.py

@@ -35,14 +35,19 @@ def get_model(model_name: Optional[str] = None) -> AIModel:
     """
     Get an AIModel instance based on config or provided name.
     """
+    provider = get_config_value("model", "provider", "openai")
+    api_key = get_config_value("model", "api_key")
+    
     if model_name is None:
         model_name = get_config_value("model", "name", "gpt-4o-mini")
     
-    provider = get_config_value("model", "provider", "openai")
-    
-    # LiteLLM usually wants "provider/model_name" for some providers 
-    # but for OpenAI it handles "gpt-3.5-turbo" directly.
-    # If it's a custom provider, we might need to prepend it.
-    full_model_name = f"{provider}/{model_name}" if provider != "openai" else model_name
+    # LiteLLM wants "provider/model_name" for non-OpenAI providers
+    if provider == "openai":
+        full_model_name = model_name
+    elif provider == "gemini":
+        full_model_name = f"gemini/{model_name}"
+    else:
+        full_model_name = f"{provider}/{model_name}"
     
-    return AIModel(model_name=full_model_name)
+    # For Ollama, we might need a base_url, but for now we assume default
+    return AIModel(model_name=full_model_name, api_key=api_key)

akita/plugins/__init__.py

@@ -0,0 +1 @@
+# Official AkitaLLM Plugins

akita/plugins/files.py

@@ -0,0 +1,34 @@
+from akita.core.plugins import AkitaPlugin
+from akita.tools.base import FileSystemTools
+from typing import List, Dict, Any
+
+class FilesPlugin(AkitaPlugin):
+    @property
+    def name(self) -> str:
+        return "files"
+
+    @property
+    def description(self) -> str:
+        return "Standard filesystem operations (read, write, list)."
+
+    def get_tools(self) -> List[Dict[str, Any]]:
+        return [
+            {
+                "name": "read_file",
+                "description": "Read content from a file.",
+                "parameters": {"path": "string"},
+                "func": FileSystemTools.read_file
+            },
+            {
+                "name": "write_file",
+                "description": "Write content to a file.",
+                "parameters": {"path": "string", "content": "string"},
+                "func": FileSystemTools.write_file
+            },
+            {
+                "name": "list_dir",
+                "description": "List files in a directory.",
+                "parameters": {"path": "string"},
+                "func": FileSystemTools.list_dir
+            }
+        ]

akita/reasoning/engine.py

@@ -1,16 +1,23 @@
 from typing import List, Dict, Any, Optional
 from akita.models.base import AIModel, get_model
-from akita.tools.base import ShellTools, FileSystemTools
+from akita.tools.base import ShellTools
+from akita.core.plugins import PluginManager
 from akita.tools.context import ContextBuilder
 from akita.schemas.review import ReviewResult
+from akita.core.trace import ReasoningTrace
+from akita.reasoning.session import ConversationSession
 import json
 from rich.console import Console
 
 console = Console()
-
+ 
 class ReasoningEngine:
     def __init__(self, model: AIModel):
         self.model = model
+        self.plugin_manager = PluginManager()
+        self.plugin_manager.discover_all()
+        self.trace = ReasoningTrace()
+        self.session: Optional[ConversationSession] = None
 
     def run_review(self, path: str) -> ReviewResult:
         """
@@ -91,27 +98,46 @@ class ReasoningEngine:
         ])
         return response.content
 
-    def run_solve(self, query: str, path: str = ".") -> str:
+    def run_solve(self, query: str, path: str = ".", session: Optional[ConversationSession] = None) -> str:
         """
         Generates a Unified Diff solution for the given query.
+        Supports iterative refinement if a session is provided.
         """
-        console.print(f"🔍 [bold]Building context for solution...[/]")
-        builder = ContextBuilder(path)
-        snapshot = builder.build()
-        
-        files_str = "\n---\n".join([f"FILE: {f.path}\nCONTENT:\n{f.content}" for f in snapshot.files[:10]]) # Limit for solve
+        self.trace.add_step("Solve", f"Starting solve for query: {query}")
         
-        system_prompt = (
-            "You are an Expert Programmer. Solve the requested task by providing code changes in Unified Diff format. "
-            "Respond ONLY with the Diff block. Use +++ and --- with file paths relative to project root."
-        )
-        user_prompt = f"Task: {query}\n\nContext:\n{files_str}\n\nGenerate the Unified Diff."
+        if not session:
+            self.trace.add_step("Context", f"Building context for {path}")
+            builder = ContextBuilder(path)
+            snapshot = builder.build(query=query)
+            
+            files_str = "\n---\n".join([f"FILE: {f.path}\nCONTENT:\n{f.content}" for f in snapshot.files[:10]])
+            
+            rag_str = ""
+            if snapshot.rag_snippets:
+                rag_str = "\n\nRELEVANT SNIPPETS (RAG):\n" + "\n".join([
+                    f"- {s['path']} ({s['name']}):\n{s['content']}" for s in snapshot.rag_snippets
+                ])
+
+            tools_info = "\n".join([f"- {t['name']}: {t['description']}" for t in self.plugin_manager.get_all_tools()])
+            
+            system_prompt = (
+                "You are an Expert Programmer. Solve the requested task by providing code changes in Unified Diff format. "
+                "Respond ONLY with the Diff block. Use +++ and --- with file paths relative to project root.\n\n"
+                f"Available Tools:\n{tools_info}"
+            )
+            
+            session = ConversationSession()
+            session.add_message("system", system_prompt)
+            session.add_message("user", f"Task: {query}\n\nContext:\n{files_str}{rag_str}")
+            self.session = session
+        else:
+            session.add_message("user", query)
+
+        console.print("🤖 [bold green]Thinking...[/]")
+        response = self.model.chat(session.get_messages_dict())
+        session.add_message("assistant", response.content)
         
-        console.print("🤖 [bold green]Generating solution...[/]")
-        response = self.model.chat([
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_prompt}
-        ])
+        self.trace.add_step("LLM Response", "Received solution from model")
         return response.content
 
     def run_pipeline(self, task: str):

akita/reasoning/session.py

@@ -0,0 +1,15 @@
+from typing import List, Dict, Any
+from pydantic import BaseModel, Field
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ConversationSession(BaseModel):
+    messages: List[ChatMessage] = Field(default_factory=list)
+
+    def add_message(self, role: str, content: str):
+        self.messages.append(ChatMessage(role=role, content=content))
+
+    def get_messages_dict(self) -> List[Dict[str, str]]:
+        return [m.model_dump() for m in self.messages]

akita/tools/base.py

@@ -15,7 +15,12 @@ class FileSystemTools:
             return f.read()
 
     @staticmethod
-    def list_files(path: str) -> List[str]:
+    def write_file(path: str, content: str):
+        with open(path, 'w', encoding='utf-8') as f:
+            f.write(content)
+
+    @staticmethod
+    def list_dir(path: str) -> List[str]:
         return os.listdir(path)
 
 class ShellTools:

akita/tools/context.py

@@ -1,16 +1,18 @@
 import os
 from pathlib import Path
-from typing import List, Dict, Optional
+from typing import Any, List, Dict, Optional
 from pydantic import BaseModel
 
 class FileContext(BaseModel):
     path: str
     content: str
     extension: str
+    summary: Optional[str] = None # New field for semantic summary
 
 class ContextSnapshot(BaseModel):
     files: List[FileContext]
     project_structure: List[str]
+    rag_snippets: Optional[List[Dict[str, Any]]] = None
 
 class ContextBuilder:
     def __init__(
@@ -19,26 +21,49 @@ class ContextBuilder:
         extensions: Optional[List[str]] = None,
         exclude_dirs: Optional[List[str]] = None,
         max_file_size_kb: int = 50,
-        max_files: int = 50
+        max_files: int = 50,
+        use_semantical_context: bool = True
     ):
         self.base_path = Path(base_path)
         self.extensions = extensions or [".py", ".js", ".ts", ".cpp", ".h", ".toml", ".md", ".json"]
         self.exclude_dirs = exclude_dirs or [".git", ".venv", "node_modules", "__pycache__", "dist", "build"]
         self.max_file_size_kb = max_file_size_kb
         self.max_files = max_files
+        self.use_semantical_context = use_semantical_context
+        
+        if self.use_semantical_context:
+            try:
+                from akita.core.ast_utils import ASTParser
+                from akita.core.indexing import CodeIndexer
+                self.ast_parser = ASTParser()
+                self.indexer = CodeIndexer(str(self.base_path))
+            except ImportError:
+                self.ast_parser = None
+                self.indexer = None
 
-    def build(self) -> ContextSnapshot:
-        """Scan the path and build a context snapshot."""
+    def build(self, query: Optional[str] = None) -> ContextSnapshot:
+        """
+        Scan the path and build a context snapshot.
+        If a query is provided and indexer is available, it includes RAG snippets.
+        """
         files_context = []
         project_structure = []
+        rag_snippets = None
         
+        if query and self.indexer:
+            try:
+                # Ensure index exists (lazy indexing for now)
+                # In production, we'd have a separate command or check timestamps
+                rag_snippets = self.indexer.search(query, n_results=10)
+            except Exception:
+                pass
+
         if self.base_path.is_file():
             if self._should_include_file(self.base_path):
                 files_context.append(self._read_file(self.base_path))
                 project_structure.append(str(self.base_path.name))
         else:
             for root, dirs, files in os.walk(self.base_path):
-                # Filter out excluded directories
                 dirs[:] = [d for d in dirs if d not in self.exclude_dirs]
                 
                 rel_root = os.path.relpath(root, self.base_path)
@@ -54,7 +79,11 @@ class ContextBuilder:
                                 files_context.append(context)
                                 project_structure.append(os.path.join(rel_root, file))
         
-        return ContextSnapshot(files=files_context, project_structure=project_structure)
+        return ContextSnapshot(
+            files=files_context, 
+            project_structure=project_structure,
+            rag_snippets=rag_snippets
+        )
 
     def _should_include_file(self, path: Path) -> bool:
         if path.name == ".env" or path.suffix == ".env":
@@ -66,8 +95,12 @@ class ContextBuilder:
         if not path.exists():
             return False
             
-        # Check size
-        if path.stat().st_size > self.max_file_size_kb * 1024:
+        # Check size (we can be more lenient if using semantic summaries)
+        size_limit = self.max_file_size_kb * 1024
+        if self.use_semantical_context:
+            size_limit *= 2 # Allow larger files if we can summarize them
+            
+        if path.stat().st_size > size_limit:
             return False
             
         return True
@@ -76,10 +109,22 @@ class ContextBuilder:
         try:
             with open(path, 'r', encoding='utf-8') as f:
                 content = f.read()
+                
+                summary = None
+                if self.use_semantical_context and self.ast_parser and path.suffix == ".py":
+                    try:
+                        defs = self.ast_parser.get_definitions(str(path))
+                        if defs:
+                            summary_lines = [f"{d['type'].upper()} {d['name']} (L{d['start_line']}-L{d['end_line']})" for d in defs]
+                            summary = "\n".join(summary_lines)
+                    except Exception:
+                        pass
+
                 return FileContext(
                     path=str(path.relative_to(self.base_path) if self.base_path.is_dir() else path.name),
                     content=content,
-                    extension=path.suffix
+                    extension=path.suffix,
+                    summary=summary
                 )
         except Exception:
             return None

akita/tools/diff.py

@@ -1,35 +1,110 @@
 import os
+import shutil
+import pathlib
 from pathlib import Path
-import re
+import whatthepatch
+from typing import List, Tuple, Optional
 
 class DiffApplier:
     @staticmethod
-    def apply_unified_diff(diff_text: str, base_path: str = "."):
+    def apply_unified_diff(diff_text: str, base_path: str = ".") -> bool:
         """
-        Simplistic Unified Diff applier. 
-        In a real scenario, this would use a robust library like 'patch-py' or 'whatthepatch'.
-        For AkitaLLM, we keep it simple for now.
+        Applies a unified diff to files in the base_path.
+        Includes backup and rollback logic for atomicity.
         """
-        # Split by file
-        file_diffs = re.split(r'--- (.*?)\n\+\+\+ (.*?)\n', diff_text)
-        
-        # Pattern extraction is tricky with regex, let's try a safer approach
-        lines = diff_text.splitlines()
-        current_file = None
-        new_content = []
-        
-        # This is a VERY placeholder implementation for safety.
-        # Applying diffs manually is high risk without a dedicated library.
-        # For the MVP, we will log what would happen.
-        
-        print(f"DEBUG: DiffApplier would process {len(lines)} lines of diff.")
-        
-        # Real logic would:
-        # 1. Path identification (--- / +++)
-        # 2. Hunk identification (@@)
-        # 3. Line modification
-        
-        return True
+        patches = list(whatthepatch.parse_patch(diff_text))
+        if not patches:
+            print("ERROR: No valid patches found in the diff text.")
+            return False
+
+        backups: List[Tuple[Path, Path]] = []
+        base = Path(base_path)
+        backup_dir = base / ".akita" / "backups"
+        backup_dir.mkdir(parents=True, exist_ok=True)
+
+        try:
+            for patch in patches:
+                if not patch.header:
+                    continue
+                
+                # whatthepatch identifies the target file in the header
+                # We usually want the 'new' filename (the +++ part)
+                rel_path = patch.header.new_path
+                is_new = (patch.header.old_path == "/dev/null")
+                is_delete = (patch.header.new_path == "/dev/null")
+
+                if is_new:
+                    rel_path = patch.header.new_path
+                elif is_delete:
+                    rel_path = patch.header.old_path
+                else:
+                    rel_path = patch.header.new_path or patch.header.old_path
+
+                if not rel_path or rel_path == "/dev/null":
+                    continue
+                
+                # Clean up path (sometimes they have a/ or b/ prefixes)
+                if rel_path.startswith("a/") or rel_path.startswith("b/"):
+                    rel_path = rel_path[2:]
+                
+                target_file = (base / rel_path).resolve()
+                
+                if not is_new and not target_file.exists():
+                    print(f"ERROR: Target file {target_file} does not exist for patching.")
+                    return False
+
+                # 1. Create backup
+                if target_file.exists():
+                    backup_file = backup_dir / f"{target_file.name}.bak"
+                    shutil.copy2(target_file, backup_file)
+                    backups.append((target_file, backup_file))
+                else:
+                    backups.append((target_file, None)) # Mark for deletion on rollback if it's a new file
+
+                # 2. Apply patch
+                content = ""
+                if target_file.exists():
+                    with open(target_file, "r", encoding="utf-8") as f:
+                        content = f.read()
+                
+                lines = content.splitlines()
+                # whatthepatch apply_diff returns a generator of lines
+                patched_lines = whatthepatch.apply_diff(patch, lines)
+                
+                if patched_lines is None:
+                    print(f"ERROR: Failed to apply patch to {rel_path}.")
+                    raise Exception(f"Patch failure on {rel_path}")
+
+                # 3. Write new content
+                target_file.parent.mkdir(parents=True, exist_ok=True)
+                with open(target_file, "w", encoding="utf-8") as f:
+                    f.write("\n".join(patched_lines) + "\n")
+
+            print(f"SUCCESS: Applied {len(patches)} patches successfully.")
+            
+            # 4. Pre-flight Validation
+            # Run tests to ensure the patch didn't break anything
+            if (base / "tests").exists():
+                print("🧪 Running pre-flight validation (pytest)...")
+                import subprocess
+                # Run pytest in the base_path
+                result = subprocess.run(["pytest"], cwd=str(base), capture_output=True, text=True)
+                if result.returncode != 0:
+                    print(f"❌ Validation FAILED:\n{result.stdout}")
+                    raise Exception("Pre-flight validation failed. Tests are broken.")
+                else:
+                    print("✅ Pre-flight validation passed!")
+            
+            return True
+
+        except Exception as e:
+            print(f"CRITICAL ERROR: {e}. Starting rollback...")
+            for target, backup in backups:
+                if backup and backup.exists():
+                    shutil.move(str(backup), str(target))
+                elif not backup and target.exists():
+                    target.unlink() # Delete newly created file
+            return False
 
     @staticmethod
     def apply_whole_file(file_path: str, content: str):