pw-agent 0.3.1__tar.gz

pw_agent-0.3.1/PKG-INFO

@@ -0,0 +1,67 @@
+Metadata-Version: 2.4
+Name: pw-agent
+Version: 0.3.1
+Summary: CLI coding assistant powered by your Ollama GPUs via PastaWater
+Home-page: https://pastawater.io
+Author: PastaWater
+Author-email: support@pastawater.io
+Project-URL: Homepage, https://pastawater.io
+Project-URL: GPU Setup, https://pastawater.io/gpu-setup
+Classifier: Development Status :: 4 - Beta
+Classifier: Environment :: Console
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Software Development :: Libraries
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: requests>=2.28.0
+Requires-Dist: rich>=13.0.0
+Requires-Dist: prompt_toolkit>=3.0.0
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: project-url
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# PW Agent
+
+CLI coding assistant powered by your Ollama GPUs via [PastaWater](https://pastawater.io).
+
+## Install
+
+```bash
+pip install pw-agent
+```
+
+## Usage
+
+```bash
+pw-agent
+```
+
+First run guides you through setup — paste your API token, pick a GPU, start chatting.
+
+## Features
+
+- Interactive REPL with streaming responses
+- Tab autocomplete for commands and file paths
+- Session persistence (resume where you left off)
+- `/add file.py` or `@file.py` — inject files into context
+- `/models` — view your GPU fleet
+- `/switch N` — hot-switch between GPUs
+- `/commit` — AI-generated git commit messages
+- `-y` flag for auto-approve mode
+- `-p "prompt"` for one-shot non-interactive mode
+
+## Connect
+
+- **Cloud mode**: Use your PastaWater API token
+- **Direct mode**: Point at a local Ollama instance (`--brain http://localhost:11434`)
+
+Get your token at [pastawater.io/settings](https://pastawater.io/settings?tab=cli)

pw_agent-0.3.1/README.md

@@ -0,0 +1,36 @@
+# PW Agent
+
+CLI coding assistant powered by your Ollama GPUs via [PastaWater](https://pastawater.io).
+
+## Install
+
+```bash
+pip install pw-agent
+```
+
+## Usage
+
+```bash
+pw-agent
+```
+
+First run guides you through setup — paste your API token, pick a GPU, start chatting.
+
+## Features
+
+- Interactive REPL with streaming responses
+- Tab autocomplete for commands and file paths
+- Session persistence (resume where you left off)
+- `/add file.py` or `@file.py` — inject files into context
+- `/models` — view your GPU fleet
+- `/switch N` — hot-switch between GPUs
+- `/commit` — AI-generated git commit messages
+- `-y` flag for auto-approve mode
+- `-p "prompt"` for one-shot non-interactive mode
+
+## Connect
+
+- **Cloud mode**: Use your PastaWater API token
+- **Direct mode**: Point at a local Ollama instance (`--brain http://localhost:11434`)
+
+Get your token at [pastawater.io/settings](https://pastawater.io/settings?tab=cli)

pw_agent-0.3.1/agent.py

@@ -0,0 +1,363 @@
+"""ReAct agent loop — the brain of pw-agent."""
+
+import json
+import os
+import re
+import subprocess
+from typing import Optional
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.panel import Panel
+from rich.syntax import Syntax
+from rich.text import Text
+from llm_client import LLMClient
+from tools import TOOL_DEFINITIONS, execute_tool
+from config import save_session
+
+
+MAX_ITERATIONS = 25
+# ~4 chars per token is a reasonable estimate for most models
+CHARS_PER_TOKEN = 4
+
+# ─── Supported PastaWater Ollama models and their context limits ──────────
+# These are the exact models available on the GPU Setup page.
+# Context limits are the model's native max, but we apply a safe budget
+# based on available VRAM to prevent OOM.
+SUPPORTED_MODELS = {
+    # id                  native_ctx   safe_ctx   description
+    "gpt-oss:20b":       {"native": 32768, "safe": 8192,  "name": "GPT-OSS 20B",   "vram_gb": 16},
+    "qwen2.5:14b":       {"native": 32768, "safe": 8192,  "name": "Qwen 2.5 14B",  "vram_gb": 10},
+    "qwen3.5:4b":        {"native": 32768, "safe": 16384, "name": "Qwen 3.5 4B",   "vram_gb": 3.5},
+    "qwen3.5:2b":        {"native": 32768, "safe": 16384, "name": "Qwen 3.5 2B",   "vram_gb": 2},
+}
+# Fallback for unknown models
+DEFAULT_SAFE_CONTEXT = 4096
+
+# Reserve tokens for the model's response
+RESPONSE_RESERVE = 2048
+
+console = Console()
+
+
+SYSTEM_PROMPT = """You are PW Agent — a helpful AI assistant running on the user's own GPU hardware via PastaWater. You can have natural conversations, answer questions, AND help with coding tasks.
+
+Be yourself. Answer questions naturally. If the user asks about you, your model, your training, etc — answer honestly based on what you know. You're not limited to only coding topics.
+
+When the user needs help with files or code, you have tools available. For casual conversation, just respond normally.
+
+## Project Context
+Working directory: {cwd}
+{git_context}
+
+## Tools (use only when needed for file/code tasks)
+
+{tool_list}
+
+When you need a tool, output EXACTLY this on its own line:
+ACTION: {{"tool": "tool_name", "args": {{"param1": "value1"}}}}
+
+Tool rules:
+- ONE action at a time, then wait for the result.
+- Read a file before editing it.
+- For edit_file, old_str must match EXACTLY (including whitespace).
+- After a RESULT, continue or use another tool.
+- When done, respond normally without ACTION.
+"""
+
+
+def _build_tool_list() -> str:
+    """Format tool definitions for the system prompt."""
+    lines = []
+    for tool in TOOL_DEFINITIONS:
+        params = ", ".join(
+            f"{k}: {v['type']}" for k, v in tool["parameters"].items()
+        )
+        lines.append(f"- {tool['name']}({params}) — {tool['description']}")
+    return "\n".join(lines)
+
+
+def _get_git_context() -> str:
+    """Gather git status for the system prompt."""
+    parts = []
+    try:
+        branch = subprocess.run(
+            ["git", "branch", "--show-current"],
+            capture_output=True, text=True, timeout=5
+        )
+        if branch.returncode == 0 and branch.stdout.strip():
+            parts.append(f"Git branch: {branch.stdout.strip()}")
+
+        status = subprocess.run(
+            ["git", "status", "--short"],
+            capture_output=True, text=True, timeout=5
+        )
+        if status.returncode == 0 and status.stdout.strip():
+            changed = len(status.stdout.strip().split("\n"))
+            parts.append(f"Git status: {changed} changed files")
+
+        log = subprocess.run(
+            ["git", "log", "--oneline", "-5"],
+            capture_output=True, text=True, timeout=5
+        )
+        if log.returncode == 0 and log.stdout.strip():
+            parts.append(f"Recent commits:\n{log.stdout.strip()}")
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        pass
+
+    return "\n".join(parts) if parts else "Not a git repository"
+
+
+def _build_messages(conversation: list[dict], cwd: str) -> list[dict]:
+    """Build Ollama chat messages from conversation history."""
+    git_ctx = _get_git_context()
+    system = SYSTEM_PROMPT.format(
+        cwd=cwd,
+        git_context=git_ctx,
+        tool_list=_build_tool_list(),
+    )
+
+    messages = [{"role": "system", "content": system}]
+
+    for msg in conversation:
+        role = msg["role"]
+        content = msg["content"]
+        if role == "user":
+            messages.append({"role": "user", "content": content})
+        elif role == "assistant":
+            messages.append({"role": "assistant", "content": content})
+        elif role == "tool_result":
+            messages.append({"role": "user", "content": content})
+
+    return messages
+
+
+def _parse_tool_call(response: str) -> Optional[tuple[dict, str]]:
+    """Extract a tool call from the model's response.
+
+    Returns (tool_call_dict, text_before_action) or None.
+    """
+    match = re.search(r"ACTION:\s*(\{.*?\})\s*$", response, re.MULTILINE | re.DOTALL)
+    if not match:
+        match = re.search(r'ACTION:\s*(\{[^}]*"tool"[^}]*"args"[^}]*\{[^}]*\}[^}]*\})', response, re.DOTALL)
+    if not match:
+        return None
+
+    try:
+        call = json.loads(match.group(1))
+        if "tool" in call and "args" in call:
+            before = response[:match.start()].strip()
+            return call, before
+    except json.JSONDecodeError:
+        json_str = match.group(1).replace("'", '"')
+        try:
+            call = json.loads(json_str)
+            if "tool" in call and "args" in call:
+                before = response[:match.start()].strip()
+                return call, before
+        except json.JSONDecodeError:
+            pass
+
+    return None
+
+
+def _estimate_tokens(text: str) -> int:
+    """Rough token estimate (~4 chars per token)."""
+    return len(text) // CHARS_PER_TOKEN + 1
+
+
+def _get_context_budget(model_name: str) -> int:
+    """Get the safe input token budget for a model.
+
+    Uses the 'safe' context limit (conservative for VRAM) minus response reserve.
+    Larger models use more VRAM per token, so their safe limit is lower than native.
+    """
+    model_lower = model_name.lower().strip()
+    # Exact match first
+    if model_lower in SUPPORTED_MODELS:
+        return SUPPORTED_MODELS[model_lower]["safe"] - RESPONSE_RESERVE
+    # Prefix match (e.g. "qwen3.5" matches "qwen3.5:4b")
+    for model_id, info in SUPPORTED_MODELS.items():
+        base = model_id.split(":")[0]
+        if model_lower.startswith(base):
+            return info["safe"] - RESPONSE_RESERVE
+    return DEFAULT_SAFE_CONTEXT - RESPONSE_RESERVE
+
+
+def _truncate_history(conversation: list[dict], model: str = "default") -> list[dict]:
+    """Sliding window: keep recent turns within the model's context budget.
+
+    Always preserves:
+    1. The first user message (original task context)
+    2. The most recent turns (working memory)
+    Drops middle messages when the budget is exceeded.
+    """
+    if len(conversation) <= 3:
+        return conversation
+
+    budget = _get_context_budget(model)
+    total_tokens = sum(_estimate_tokens(m["content"]) for m in conversation)
+
+    if total_tokens <= budget:
+        return conversation
+
+    # Always keep the first message
+    first = conversation[0]
+    first_tokens = _estimate_tokens(first["content"])
+    remaining_budget = budget - first_tokens
+
+    # Walk backwards from the end, adding messages until budget is hit
+    remaining = conversation[1:]
+    kept = []
+    used = 0
+    for msg in reversed(remaining):
+        msg_tokens = _estimate_tokens(msg["content"])
+        if used + msg_tokens > remaining_budget:
+            break
+        kept.insert(0, msg)
+        used += msg_tokens
+
+    dropped = len(remaining) - len(kept)
+    result = [first]
+    if dropped > 0:
+        result.append({"role": "tool_result", "content": f"[... {dropped} earlier messages truncated to fit context window ...]"})
+    result.extend(kept)
+    return result
+
+
+class Agent:
+    """ReAct agent that uses an LLM to perform coding tasks."""
+
+    def __init__(self, client: LLMClient, stream: bool = True):
+        self.client = client
+        self.stream = stream
+        self.conversation: list[dict] = []
+        self.cwd = os.getcwd()
+        self.files_in_context: list[str] = []
+
+    def run(self, user_input: str) -> None:
+        """Process a user message through the ReAct loop."""
+        self.conversation.append({"role": "user", "content": user_input})
+
+        for iteration in range(MAX_ITERATIONS):
+            trimmed = _truncate_history(self.conversation, self.client.model)
+            messages = _build_messages(trimmed, self.cwd)
+
+            # Get response (streaming or not)
+            if self.stream:
+                response = self._stream_response(messages)
+            else:
+                with console.status("[cyan]Thinking...", spinner="dots"):
+                    response = self.client.chat(messages)
+
+            if not response or response.startswith("[Error:"):
+                console.print(f"  [red]{response or '[Empty response from model]'}[/red]")
+                return
+
+            # Check for tool call
+            parsed = _parse_tool_call(response)
+
+            if parsed:
+                tool_call, thinking = parsed
+                tool_name = tool_call["tool"]
+                tool_args = tool_call["args"]
+
+                # Print thinking
+                if thinking:
+                    console.print()
+                    console.print(Markdown(thinking))
+
+                # Print tool call
+                args_display = []
+                for k, v in tool_args.items():
+                    if isinstance(v, str) and len(v) > 80:
+                        args_display.append(f'{k}="...{len(v)} chars..."')
+                    else:
+                        args_display.append(f'{k}="{v}"' if isinstance(v, str) else f'{k}={v}')
+                args_str = ", ".join(args_display)
+                console.print(f"  [cyan]> {tool_name}({args_str})[/cyan]")
+
+                # Execute tool
+                result = execute_tool(tool_name, tool_args)
+
+                # Print result
+                display = result[:800] + "..." if len(result) > 800 else result
+                if tool_name == "read_file" and len(result) > 200:
+                    # Show file content with line numbers
+                    console.print(f"  [dim]{display}[/dim]")
+                else:
+                    console.print(f"  [dim]{display}[/dim]")
+
+                # Add to conversation
+                if thinking:
+                    self.conversation.append({"role": "assistant", "content": thinking})
+                self.conversation.append({"role": "assistant", "content": f"ACTION: used {tool_name}"})
+                self.conversation.append({"role": "tool_result", "content": f"RESULT of {tool_name}:\n{result}"})
+
+            else:
+                # Final answer
+                self.conversation.append({"role": "assistant", "content": response})
+                if not self.stream:
+                    console.print()
+                    console.print(Markdown(response))
+                self._auto_save()
+                return
+
+        console.print(f"  [yellow][Reached max iterations ({MAX_ITERATIONS}), stopping][/yellow]")
+        self._auto_save()
+
+    def _stream_response(self, messages: list[dict]) -> str:
+        """Stream response tokens to terminal, return full text."""
+        chunks = []
+        console.print()
+        for chunk in self.client.chat_stream(messages):
+            chunks.append(chunk)
+            # Print raw chunks for streaming effect
+            console.print(chunk, end="", highlight=False)
+
+        full = "".join(chunks)
+        console.print()  # newline after stream
+
+        # If it contains a tool call, re-render won't be needed (agent loop handles it)
+        return full
+
+    def _auto_save(self):
+        """Save session to disk after each turn."""
+        try:
+            save_session(self.conversation, self.cwd)
+        except Exception:
+            pass
+
+    def load_session(self, conversation: list[dict]):
+        """Resume a previous session."""
+        self.conversation = conversation
+        turns = sum(1 for m in conversation if m["role"] == "user")
+        console.print(f"  [dim]Resumed session ({turns} turns)[/dim]")
+
+    def add_file(self, path: str):
+        """Add a file's contents to the conversation context."""
+        full = os.path.abspath(path)
+        if not os.path.exists(full):
+            console.print(f"  [red]File not found: {path}[/red]")
+            return
+        try:
+            with open(full, "r", encoding="utf-8", errors="replace") as f:
+                content = f.read()
+            # Truncate large files
+            if len(content) > 15000:
+                content = content[:15000] + f"\n... [truncated, {len(content)} chars total]"
+            self.conversation.append({
+                "role": "user",
+                "content": f"[File added to context: {path}]\n```\n{content}\n```"
+            })
+            self.files_in_context.append(path)
+            console.print(f"  [green]+ {path}[/green] [dim]({len(content)} chars)[/dim]")
+        except Exception as e:
+            console.print(f"  [red]Error reading {path}: {e}[/red]")
+
+    def reset(self):
+        """Clear conversation history and session."""
+        self.conversation = []
+        self.files_in_context = []
+        from config import clear_session
+        clear_session(self.cwd)
+        console.print("  [dim]Conversation cleared.[/dim]")

pw_agent-0.3.1/config.py

@@ -0,0 +1,88 @@
+"""Config management — persists token, slot, model preferences."""
+
+import json
+import os
+
+DEFAULT_CONFIG_DIR = os.path.expanduser("~/.config/pw-agent")
+CONFIG_FILE = "config.json"
+
+
+def _config_path(config_dir: str = "") -> str:
+    d = config_dir or DEFAULT_CONFIG_DIR
+    return os.path.join(d, CONFIG_FILE)
+
+
+def load_config(config_dir: str = "") -> dict:
+    """Load saved config, or return empty dict."""
+    path = _config_path(config_dir)
+    if not os.path.exists(path):
+        return {}
+    try:
+        with open(path, "r") as f:
+            return json.load(f)
+    except (json.JSONDecodeError, OSError):
+        return {}
+
+
+def save_config(data: dict, config_dir: str = "") -> str:
+    """Save config and return the file path."""
+    d = config_dir or DEFAULT_CONFIG_DIR
+    os.makedirs(d, exist_ok=True)
+    path = os.path.join(d, CONFIG_FILE)
+    with open(path, "w") as f:
+        json.dump(data, f, indent=2)
+    return path
+
+
+def has_config(config_dir: str = "") -> bool:
+    """Check if a config file exists."""
+    return os.path.exists(_config_path(config_dir))
+
+
+# ─── Session persistence ─────────────────────────────────────────────────────
+
+def _sessions_dir(config_dir: str = "") -> str:
+    d = config_dir or DEFAULT_CONFIG_DIR
+    return os.path.join(d, "sessions")
+
+
+def save_session(conversation: list[dict], cwd: str, config_dir: str = "") -> str:
+    """Save conversation to a session file. Returns the path."""
+    d = _sessions_dir(config_dir)
+    os.makedirs(d, exist_ok=True)
+    # Use cwd hash as session key so each project has its own session
+    import hashlib
+    key = hashlib.md5(cwd.encode()).hexdigest()[:10]
+    path = os.path.join(d, f"{key}.json")
+    data = {"cwd": cwd, "conversation": conversation}
+    with open(path, "w") as f:
+        json.dump(data, f)
+    return path
+
+
+def load_session(cwd: str, config_dir: str = "") -> list[dict]:
+    """Load the session for a given cwd, or return empty list."""
+    d = _sessions_dir(config_dir)
+    import hashlib
+    key = hashlib.md5(cwd.encode()).hexdigest()[:10]
+    path = os.path.join(d, f"{key}.json")
+    if not os.path.exists(path):
+        return []
+    try:
+        with open(path, "r") as f:
+            data = json.load(f)
+        if data.get("cwd") == cwd:
+            return data.get("conversation", [])
+    except (json.JSONDecodeError, OSError):
+        pass
+    return []
+
+
+def clear_session(cwd: str, config_dir: str = ""):
+    """Delete the session for a given cwd."""
+    d = _sessions_dir(config_dir)
+    import hashlib
+    key = hashlib.md5(cwd.encode()).hexdigest()[:10]
+    path = os.path.join(d, f"{key}.json")
+    if os.path.exists(path):
+        os.remove(path)