Enhancement Suite: Anti-loop protection, Kokoro TTS integration, ComfyUI status monitoring, and auto-open image functionality

build/lib/ollama_cli/tools/base.py

@@ -0,0 +1,55 @@
+from typing import Dict, Any, Callable, List, Optional
+import json
+import inspect
+
+class BaseTool:
+    def __init__(self, name: str, description: str, parameters: Dict[str, Any]):
+        self.name = name
+        self.description = description
+        self.parameters = parameters
+
+    def execute(self, **kwargs) -> Any:
+        raise NotImplementedError
+
+class ToolRegistry:
+    def __init__(self):
+        self.tools: Dict[str, BaseTool] = {}
+
+    def register(self, tool: BaseTool):
+        self.tools[tool.name] = tool
+
+    def get_tool(self, name: str) -> Optional[BaseTool]:
+        return self.tools.get(name)
+
+    def list_tools(self) -> List[Dict[str, Any]]:
+        return [
+            {
+                "name": t.name,
+                "description": t.description,
+                "parameters": t.parameters
+            }
+            for t in self.tools.values()
+        ]
+
+    def generate_system_prompt_snippet(self) -> str:
+        snippet = "Available tools:\n"
+        for t in self.tools.values():
+            snippet += f"- {t.name}: {t.description}\n"
+            snippet += f"  Parameters: {json.dumps(t.parameters)}\n"
+        return snippet
+
+# Singleton registry
+registry = ToolRegistry()
+
+def tool(name: str, description: str, parameters: Dict[str, Any]):
+    """Decorator to register a function as a tool"""
+    def decorator(func: Callable):
+        class DecoratedTool(BaseTool):
+            def __init__(self):
+                super().__init__(name, description, parameters)
+            def execute(self, **kwargs):
+                return func(**kwargs)
+        
+        registry.register(DecoratedTool())
+        return func
+    return decorator

build/lib/ollama_cli/tools/code.py

@@ -0,0 +1,96 @@
+import os
+import subprocess
+from .base import tool
+
+@tool(
+    name="code_analyze_file",
+    description="Analyze code structure and suggest improvements",
+    parameters={
+        "file_path": {"type": "string", "description": "Path to the file to analyze"}
+    }
+)
+def code_analyze_file(file_path: str) -> str:
+    try:
+        with open(os.path.expanduser(file_path), 'r') as f:
+            content = f.read()
+        
+        lines = content.split('\n')
+        functions = [l for l in lines if 'def ' in l]
+        classes = [l for l in lines if 'class ' in l]
+        
+        issues = []
+        if len(lines) > 500:
+            issues.append(f"Large file ({len(lines)} lines) - consider splitting")
+        if len(functions) > 30:
+            issues.append(f"Many functions ({len(functions)}) - consider refactoring")
+        
+        return f"""File: {file_path}
+Lines: {len(lines)} | Functions: {len(functions)} | Classes: {len(classes)}
+
+Issues:
+{chr(10).join(f'- {i}' for i in issues) if issues else '- None detected'}"""
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+@tool(
+    name="generate_readme",
+    description="Generate a README.md for the current project",
+    parameters={
+        "path": {"type": "string", "description": "Project root path", "default": "."}
+    }
+)
+def generate_readme(path: str = ".") -> str:
+    try:
+        project_name = os.path.basename(os.path.abspath(path))
+        readme_content = f"# {project_name}\n\nGenerated by Ollama CLI."
+        readme_path = os.path.join(path, "README.md")
+        with open(readme_path, 'w') as f:
+            f.write(readme_content)
+        return f"Generated README.md at {readme_path}"
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+@tool(
+    name="replace_text",
+    description="Surgically replace text in a file",
+    parameters={
+        "path": {"type": "string", "description": "Path to the file"},
+        "old_text": {"type": "string", "description": "The exact text to find"},
+        "new_text": {"type": "string", "description": "The text to replace it with"}
+    }
+)
+def replace_text(path: str, old_text: str, new_text: str) -> str:
+    try:
+        expanded_path = os.path.expanduser(path)
+        with open(expanded_path, 'r') as f:
+            content = f.read()
+        
+        if old_text not in content:
+            return f"Error: The text to replace was not found in {path}"
+            
+        new_content = content.replace(old_text, new_text, 1) # Replace first occurrence
+        with open(expanded_path, 'w') as f:
+            f.write(new_content)
+        return f"Successfully updated {path}"
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+@tool(
+    name="run_linter",
+    description="Run a linter on a file",
+    parameters={
+        "file_path": {"type": "string", "description": "Path to the file"},
+        "linter": {"type": "string", "description": "Linter to use (pylint|eslint|auto)", "default": "auto"}
+    }
+)
+def run_linter(file_path: str, linter: str = "auto") -> str:
+    try:
+        file_path = os.path.expanduser(file_path)
+        if linter == "auto":
+            ext = os.path.splitext(file_path)[1]
+            linter = "pylint" if ext == ".py" else "eslint"
+        
+        result = subprocess.run(f"{linter} {file_path}", shell=True, capture_output=True, text=True)
+        return result.stdout + result.stderr
+    except Exception as e:
+        return f"Error: {str(e)}"

build/lib/ollama_cli/tools/execution.py

@@ -0,0 +1,36 @@
+import subprocess
+import os
+import sys
+from .base import tool
+
+@tool(
+    name="run_python",
+    description="Execute a Python script or snippet",
+    parameters={
+        "code": {"type": "string", "description": "Python code to execute"},
+        "args": {"type": "string", "description": "Optional arguments", "default": ""}
+    }
+)
+def run_python(code: str, args: str = "") -> str:
+    try:
+        # Create a temp file
+        temp_file = ".ollama-cli-exec.py"
+        with open(temp_file, "w") as f:
+            f.write(code)
+        
+        cmd = f"{sys.executable} {temp_file} {args}"
+        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=30)
+        
+        output = result.stdout + result.stderr
+        
+        # Cleanup
+        if os.path.exists(temp_file):
+            os.remove(temp_file)
+            
+        if result.returncode == 0:
+            return f"Success:\n{output}" if output else "Success (no output)"
+        else:
+            return f"Error (Exit {result.returncode}):\n{output}"
+            
+    except Exception as e:
+        return f"Execution Error: {str(e)}"

build/lib/ollama_cli/tools/filesystem.py

@@ -0,0 +1,83 @@
+import os
+import subprocess
+from .base import tool
+
+@tool(
+    name="read_file",
+    description="Read the contents of a file",
+    parameters={
+        "path": {"type": "string", "description": "Path to the file"}
+    }
+)
+def read_file(path: str) -> str:
+    try:
+        with open(os.path.expanduser(path), 'r') as f:
+            return f.read()
+    except Exception as e:
+        return f"Error reading file: {str(e)}"
+
+@tool(
+    name="write_file",
+    description="Write content to a file",
+    parameters={
+        "path": {"type": "string", "description": "Path to the file"},
+        "content": {"type": "string", "description": "Content to write"}
+    }
+)
+def write_file(path: str, content: str) -> str:
+    try:
+        expanded_path = os.path.expanduser(path)
+        os.makedirs(os.path.dirname(expanded_path), exist_ok=True)
+        with open(expanded_path, 'w') as f:
+            f.write(content)
+        return f"Successfully written to {path}"
+    except Exception as e:
+        return f"Error writing file: {str(e)}"
+
+@tool(
+    name="list_directory",
+    description="List contents of a directory",
+    parameters={
+        "path": {"type": "string", "description": "Path to the directory", "default": "."}
+    }
+)
+def list_directory(path: str = ".") -> str:
+    try:
+        expanded = os.path.expanduser(path)
+        files = os.listdir(expanded)
+        return "\n".join(files) if files else "Directory is empty"
+    except Exception as e:
+        return f"Error listing directory: {str(e)}"
+
+@tool(
+    name="grep_search",
+    description="Search for a regular expression pattern in file contents",
+    parameters={
+        "pattern": {"type": "string", "description": "The regex pattern to search for"},
+        "path": {"type": "string", "description": "Directory to search in", "default": "."},
+        "include": {"type": "string", "description": "Glob pattern for files to include", "default": "*"}
+    }
+)
+def grep_search(pattern: str, path: str = ".", include: str = "*") -> str:
+    try:
+        cmd = f'grep -rE --include="{include}" "{pattern}" {path}'
+        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=20)
+        return result.stdout if result.stdout else "No matches found."
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+@tool(
+    name="get_tree",
+    description="Show the directory structure as a tree",
+    parameters={
+        "path": {"type": "string", "description": "Root path", "default": "."},
+        "depth": {"type": "integer", "description": "Max depth", "default": 2}
+    }
+)
+def get_tree(path: str = ".", depth: int = 2) -> str:
+    try:
+        cmd = f"find {path} -maxdepth {depth} -not -path '*/.*'"
+        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+        return result.stdout
+    except Exception as e:
+        return f"Error: {str(e)}"

build/lib/ollama_cli/tools/git.py

@@ -0,0 +1,52 @@
+import subprocess
+import requests
+from typing import Optional
+from .base import tool
+from ..core.config import load_config
+
+@tool(
+    name="git_operation",
+    description="Perform git operations (status, diff, commit, push, pull, list)",
+    parameters={
+        "operation": {"type": "string", "description": "Operation to perform (status|diff|commit|push|pull|list)"},
+        "message": {"type": "string", "description": "Commit message", "default": "Auto commit"}
+    }
+)
+def git_operation(operation: str, message: str = "Auto commit") -> str:
+    config = load_config()
+    git_config = config.get("git", {})
+    gitea_config = config.get("gitea", {})
+    
+    try:
+        if operation == "status":
+            result = subprocess.run("git status", shell=True, capture_output=True, text=True)
+        elif operation == "diff":
+            result = subprocess.run("git diff", shell=True, capture_output=True, text=True)
+        elif operation == "commit":
+            result = subprocess.run(f'git add -A && git commit -m "{message}"', 
+                                  shell=True, capture_output=True, text=True)
+        elif operation == "push":
+            result = subprocess.run("git push", shell=True, capture_output=True, text=True)
+        elif operation == "pull":
+            result = subprocess.run("git pull", shell=True, capture_output=True, text=True)
+        elif operation == "list":
+            if gitea_config.get("url") and gitea_config.get("token"):
+                gitea_url = gitea_config["url"].rstrip("/")
+                token = gitea_config["token"]
+                response = requests.get(
+                    f"{gitea_url}/api/v1/user/repos",
+                    headers={"Authorization": f"token {token}"},
+                    timeout=10
+                )
+                if response.status_code == 200:
+                    repos = response.json()
+                    return "\n".join([f"- {r['full_name']} ({r['clone_url']})" for r in repos])
+                return f"Gitea API error: {response.status_code}"
+            result = subprocess.run("find . -maxdepth 2 -name '.git' -type d | sed 's|/.git||'", 
+                                  shell=True, capture_output=True, text=True)
+        else:
+            return f"Unknown git operation: {operation}"
+        
+        return result.stdout + result.stderr
+    except Exception as e:
+        return f"Git error: {str(e)}"

build/lib/ollama_cli/tools/knowledge.py

@@ -0,0 +1,41 @@
+import json
+import os
+from .base import tool
+
+KB_FILE = os.path.expanduser("~/.ollama-cli-kb.json")
+
+def load_kb():
+    if os.path.exists(KB_FILE):
+        with open(KB_FILE, 'r') as f:
+            return json.load(f)
+    return {}
+
+def save_kb(kb):
+    with open(KB_FILE, 'w') as f:
+        json.dump(kb, f, indent=2)
+
+@tool(
+    name="kb_add",
+    description="Add content to the knowledge base",
+    parameters={
+        "key": {"type": "string", "description": "Search key"},
+        "content": {"type": "string", "description": "Information content"}
+    }
+)
+def kb_add(key: str, content: str) -> str:
+    kb = load_kb()
+    kb[key] = content
+    save_kb(kb)
+    return f"Added '{key}' to knowledge base"
+
+@tool(
+    name="kb_search",
+    description="Search the knowledge base",
+    parameters={
+        "query": {"type": "string", "description": "Search query"}
+    }
+)
+def kb_search(query: str) -> str:
+    kb = load_kb()
+    results = [f"{k}: {v[:100]}..." for k, v in kb.items() if query.lower() in k.lower() or query.lower() in v.lower()]
+    return "\n".join(results) if results else "No matches found"

build/lib/ollama_cli/tools/media.py

@@ -0,0 +1,183 @@
+import subprocess
+import os
+import sys
+import requests
+import json
+import random
+from .base import tool
+from ..core.config import load_config
+
+@tool(
+    name="generate_image",
+    description="Generate an image using the local Stable Diffusion workflow (ComfyUI)",
+    parameters={
+        "prompt": {"type": "string", "description": "The image description"}
+    }
+)
+def generate_image(prompt: str) -> str:
+    """Invokes the ComfyUI API directly with a predefined workflow"""
+    config = load_config()
+    comfy_url = config.get("comfy_url", "http://127.0.0.1:8188").rstrip("/")
+    image_model = config.get("image_model", "dreamshaper_8_pruned.safetensors")
+    
+    # Determine dimensions (SDXL prefers 1024x1024)
+    width, height = (512, 512)
+    if "xl" in image_model.lower():
+        width, height = (1024, 1024)
+
+    try:
+        # Check if ComfyUI is running
+        try:
+            requests.get(comfy_url, timeout=2)
+        except:
+            return f"Error: ComfyUI server ({comfy_url}) appears to be down. Please start it first or configure the URL with '/config comfy <url>'."
+
+        # Define the workflow directly
+        workflow = {
+            "1": {
+                "inputs": {"ckpt_name": image_model},
+                "class_type": "CheckpointLoaderSimple"
+            },
+            "2": {
+                "inputs": {
+                    "text": f"masterpiece, best quality, highly detailed, {prompt}",
+                    "clip": ["1", 1]
+                },
+                "class_type": "CLIPTextEncode"
+            },
+            "3": {
+                "inputs": {
+                    "text": "blurry, low quality, distorted, watermark, nsfw",
+                    "clip": ["1", 1]
+                },
+                "class_type": "CLIPTextEncode"
+            },
+            "4": {
+                "inputs": {"width": width, "height": height, "batch_size": 1},
+                "class_type": "EmptyLatentImage"
+            },
+            "5": {
+                "inputs": {
+                    "seed": random.randint(1, 1000000000), "steps": 20, "cfg": 7.0,
+                    "sampler_name": "dpmpp_2m", "scheduler": "karras",
+                    "denoise": 1.0, "model": ["1", 0],
+                    "positive": ["2", 0], "negative": ["3", 0],
+                    "latent_image": ["4", 0]
+                },
+                "class_type": "KSampler"
+            },
+            "6": {
+                "inputs": {"samples": ["5", 0], "vae": ["1", 2]},
+                "class_type": "VAEDecode"
+            },
+            "7": {
+                "inputs": {
+                    "filename_prefix": "ollama_cli_",
+                    "images": ["6", 0]
+                },
+                "class_type": "SaveImage"
+            }
+        }
+        
+        payload = {"prompt": workflow}
+        response = requests.post(f"{comfy_url}/prompt", json=payload, timeout=10)
+        
+        if response.status_code == 200:
+            data = response.json()
+            prompt_id = data.get('prompt_id', 'unknown')
+            return f"Image generation queued successfully on {comfy_url}.\nPrompt ID: {prompt_id}\nCheck your ComfyUI output folder for 'ollama_cli_' prefixed files."
+        else:
+            return f"Error form ComfyUI: {response.status_code} - {response.text}"
+            
+    except Exception as e:
+        return f"Error generating image: {str(e)}"
+
+@tool(
+    name="speak_text",
+    description="Convert text to speech using Piper TTS (high quality)",
+    parameters={
+        "text": {"type": "string", "description": "Text to speak"},
+        "output_file": {"type": "string", "description": "Output filename", "default": "output.wav"}
+    }
+)
+def speak_text(text: str, output_file: str = "output.wav") -> str:
+    """Uses the best available TTS engine (Kokoro > Piper > System)"""
+    config = load_config()
+    
+    # 1. Try Kokoro (Ultra-Realistic)
+    kokoro_python = os.path.expanduser(config.get("kokoro_path", "~/.ollama-cli/kokoro/venv/bin/python"))
+    if os.path.exists(kokoro_python):
+        try:
+            # Inline script to run in Kokoro venv
+            kokoro_script = f"""
+import sys
+import soundfile as sf
+import os
+try:
+    from kokoro import KPipeline
+    import torch
+except ImportError:
+    print("Dependencies missing")
+    sys.exit(1)
+
+text = sys.argv[1]
+output_file = sys.argv[2]
+
+# Initialize pipeline
+pipeline = KPipeline(lang_code='a') 
+generator = pipeline(text, voice='af_bella', speed=1)
+
+all_audio = []
+for gs, ps, audio in generator:
+    if audio is not None:
+        all_audio.append(audio)
+
+if all_audio:
+    import numpy as np
+    full_audio = np.concatenate(all_audio)
+    sf.write(output_file, full_audio, 24000)
+    print("Success")
+else:
+    print("No audio generated")
+    sys.exit(1)
+"""
+            # Run in venv
+            subprocess.run([kokoro_python, "-c", kokoro_script, text, output_file], 
+                           check=True, capture_output=True)
+            
+            if sys.platform == "darwin":
+                subprocess.run(["afplay", output_file])
+            return f"Audio saved to {output_file} and played (using Ultra-Realistic Kokoro TTS)"
+        except Exception as e:
+            # Continue to fallback
+            pass
+
+    # 2. Try Piper (High-Quality)
+    piper_config = config.get("piper", {})
+    piper_bin = os.path.expanduser(piper_config.get("path", "~/.ollama-cli/piper/piper"))
+    model_path = os.path.expanduser(piper_config.get("model", "~/.ollama-cli/piper/voice.onnx"))
+    
+    if os.path.exists(piper_bin) and os.path.exists(model_path):
+        try:
+            # Piper reads from stdin and writes to file
+            cmd = f'echo "{text}" | {piper_bin} --model {model_path} --output_file {output_file}'
+            subprocess.run(cmd, shell=True, check=True, capture_output=True)
+            
+            # Auto-play on macOS
+            if sys.platform == "darwin":
+                subprocess.run(["afplay", output_file])
+                
+            return f"Audio saved to {output_file} and played (using high-quality Piper TTS)"
+        except Exception as e:
+            pass
+    
+    # 3. Fallback to system 'say' command on macOS
+    if sys.platform == "darwin":
+        try:
+            subprocess.run(["say", "-o", output_file, text], check=True)
+            subprocess.run(["afplay", output_file])
+            return f"Audio saved to {output_file} and played (using system 'say' fallback)"
+        except Exception as e:
+            return f"Error using 'say': {e}"
+            
+    return "No compatible TTS engine found. Please run the installer."

build/lib/ollama_cli/tools/memory.py

@@ -0,0 +1,59 @@
+import json
+import os
+from .base import tool
+
+PROJECT_MEMORY_FILE = ".ollama-cli/memory.json"
+
+def load_project_memory():
+    if os.path.exists(PROJECT_MEMORY_FILE):
+        try:
+            with open(PROJECT_MEMORY_FILE, 'r') as f:
+                return json.load(f)
+        except:
+            pass
+    return {}
+
+def save_project_memory(memory):
+    os.makedirs(os.path.dirname(PROJECT_MEMORY_FILE), exist_ok=True)
+    with open(PROJECT_MEMORY_FILE, 'w') as f:
+        json.dump(memory, f, indent=2)
+
+@tool(
+    name="remember_fact",
+    description="Store a key piece of information about the project",
+    parameters={
+        "fact": {"type": "string", "description": "The fact to remember (e.g., 'API key is in .env', 'Main entry point is src/index.ts')"}
+    }
+)
+def remember_fact(fact: str) -> str:
+    memory = load_project_memory()
+    facts = memory.get("facts", [])
+    if fact not in facts:
+        facts.append(fact)
+        memory["facts"] = facts
+        save_project_memory(memory)
+        return f"Remembered: {fact}"
+    return "Fact already known."
+
+@tool(
+    name="recall_facts",
+    description="Recall all stored facts about the project",
+    parameters={}
+)
+def recall_facts() -> str:
+    memory = load_project_memory()
+    facts = memory.get("facts", [])
+    if not facts:
+        return "No project facts stored."
+    return "Project Facts:\n" + "\n".join(f"- {f}" for f in facts)
+
+@tool(
+    name="clear_memory",
+    description="Clear project memory",
+    parameters={}
+)
+def clear_memory() -> str:
+    if os.path.exists(PROJECT_MEMORY_FILE):
+        os.remove(PROJECT_MEMORY_FILE)
+        return "Project memory cleared."
+    return "Memory was already empty."

build/lib/ollama_cli/tools/system.py

@@ -0,0 +1,46 @@
+import subprocess
+import os
+from .base import tool
+
+@tool(
+    name="execute_bash",
+    description="Execute a bash command and return its output",
+    parameters={
+        "command": {"type": "string", "description": "The command to execute"}
+    }
+)
+def execute_bash(command: str) -> str:
+    # Basic safety check
+    dangerous = ['rm -rf /', 'dd if=', 'mkfs', ':(){:|:&};:']
+    if any(d in command for d in dangerous):
+        return f"⚠️ DANGEROUS COMMAND BLOCKED: {command}"
+    
+    try:
+        result = subprocess.run(
+            command, 
+            shell=True, 
+            capture_output=True, 
+            text=True, 
+            cwd=os.getcwd(), 
+            timeout=60
+        )
+        output = result.stdout + result.stderr
+        return f"Exit code: {result.returncode}\n{output}" if output else f"Exit code: {result.returncode} (no output)"
+    except subprocess.TimeoutExpired:
+        return "Error: Command timed out after 60 seconds"
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+@tool(
+    name="get_env_info",
+    description="Get information about the current system environment",
+    parameters={}
+)
+def get_env_info() -> str:
+    import sys
+    import platform
+    return f"""OS: {platform.system()} {platform.release()}
+Python: {sys.version}
+Shell: {os.environ.get('SHELL', 'unknown')}
+User: {os.environ.get('USER', 'unknown')}
+CWD: {os.getcwd()}"""

build/lib/ollama_cli/tools/web.py

@@ -0,0 +1,67 @@
+import requests
+import urllib.parse
+from .base import tool
+
+@tool(
+    name="web_search",
+    description="Search the web for information",
+    parameters={
+        "query": {"type": "string", "description": "The search query"}
+    }
+)
+def web_search(query: str) -> str:
+    """Enhanced web search with DuckDuckGo"""
+    try:
+        encoded = urllib.parse.quote(query)
+        # Use DuckDuckGo Instant Answer API
+        api_url = f"https://api.duckduckgo.com/?q={encoded}&format=json"
+        response = requests.get(api_url, timeout=10)
+        data = response.json()
+        
+        results = []
+        if data.get("AbstractText"):
+            results.append(f"Summary: {data['AbstractText']}")
+        
+        for topic in data.get("RelatedTopics", [])[:5]:
+            if isinstance(topic, dict) and "Text" in topic:
+                results.append(f"- {topic['Text']}")
+        
+        return "\n".join(results) if results else "No direct results found. Try a different query."
+    except Exception as e:
+        return f"Search error: {str(e)}"
+
+@tool(
+    name="read_url",
+    description="Read and extract text content from a URL",
+    parameters={
+        "url": {"type": "string", "description": "The URL to read"}
+    }
+)
+def read_url(url: str) -> str:
+    try:
+        response = requests.get(url, timeout=10)
+        response.raise_for_status()
+        
+        try:
+            from bs4 import BeautifulSoup
+            soup = BeautifulSoup(response.text, 'html.parser')
+            
+            # Remove scripts and styles
+            for script in soup(["script", "style"]):
+                script.decompose()
+                
+            text = soup.get_text()
+            
+            # Clean up whitespace
+            lines = (line.strip() for line in text.splitlines())
+            chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
+            text = '\n'.join(chunk for chunk in chunks if chunk)
+            
+            return text[:10000] + ("\n... (truncated)" if len(text) > 10000 else "")
+            
+        except ImportError:
+            # Fallback if bs4 is missing (though we checked)
+            return response.text[:10000]
+            
+    except Exception as e:
+        return f"Error reading URL: {str(e)}"