code-puppy 0.0.97__py3-none-any.whl → 0.0.118__py3-none-any.whl

code_puppy/__init__.py

@@ -1,6 +1,3 @@
-try:
-    import importlib.metadata
+import importlib.metadata
 
-    __version__ = importlib.metadata.version("code-puppy")
-except importlib.metadata.PackageNotFoundError:
-    __version__ = "0.0.1"
+__version__ = importlib.metadata.version("code-puppy")

code_puppy/__main__.py

@@ -0,0 +1,10 @@
+"""
+Entry point for running code-puppy as a module.
+
+This allows the package to be run with: python -m code_puppy
+"""
+
+from code_puppy.main import main_entry
+
+if __name__ == "__main__":
+    main_entry()

code_puppy/agent.py

@@ -1,24 +1,29 @@
-import os
 from pathlib import Path
+from typing import Dict, Optional
 
-import pydantic
 from pydantic_ai import Agent
-from pydantic_ai.mcp import MCPServerSSE
+from pydantic_ai.mcp import MCPServerSSE, MCPServerStdio, MCPServerStreamableHTTP
+from pydantic_ai.settings import ModelSettings
+from pydantic_ai.usage import UsageLimits
 
 from code_puppy.agent_prompts import get_system_prompt
+from code_puppy.http_utils import (
+    create_reopenable_async_client,
+    resolve_env_var_in_header,
+)
+from code_puppy.message_history_processor import (
+    get_model_context_length,
+    message_history_accumulator,
+)
+from code_puppy.messaging.message_queue import (
+    emit_error,
+    emit_info,
+    emit_system_message,
+)
 from code_puppy.model_factory import ModelFactory
-from code_puppy.state_management import message_history_accumulator
 from code_puppy.tools import register_all_tools
 from code_puppy.tools.common import console
 
-# Environment variables used in this module:
-# - MODELS_JSON_PATH: Optional path to a custom models.json configuration file.
-#                     If not set, uses the default file in the package directory.
-# - MODEL_NAME: The model to use for code generation. Defaults to "gpt-4o".
-#               Must match a key in the models.json configuration.
-
-MODELS_JSON_PATH = os.environ.get("MODELS_JSON_PATH", None)
-
 
 def load_puppy_rules():
     global PUPPY_RULES
@@ -31,59 +36,130 @@ def load_puppy_rules():
 
 # Load at import
 PUPPY_RULES = load_puppy_rules()
-
-
-class AgentResponse(pydantic.BaseModel):
-    """Represents a response from the agent."""
-
-    output_message: str = pydantic.Field(
-        ..., description="The final output message to display to the user"
-    )
-    awaiting_user_input: bool = pydantic.Field(
-        False, description="True if user input is needed to continue the task"
-    )
-
-
+_LAST_MODEL_NAME = None
 _code_generation_agent = None
 
 
-def _load_mcp_servers():
-    from code_puppy.config import load_mcp_server_configs
+def _load_mcp_servers(extra_headers: Optional[Dict[str, str]] = None):
+    from code_puppy.config import get_value, load_mcp_server_configs
+
+    # Check if MCP servers are disabled
+    mcp_disabled = get_value("disable_mcp_servers")
+    if mcp_disabled and str(mcp_disabled).lower() in ("1", "true", "yes", "on"):
+        emit_system_message("[dim]MCP servers disabled via config[/dim]")
+        return []
 
     configs = load_mcp_server_configs()
+    if not configs:
+        emit_system_message("[dim]No MCP servers configured[/dim]")
+        return []
     servers = []
     for name, conf in configs.items():
+        server_type = conf.get("type", "sse")
         url = conf.get("url")
-        if url:
-            console.print(f"Registering MCP Server - {url}")
-            servers.append(MCPServerSSE(url=url))
+        timeout = conf.get("timeout", 30)
+        server_headers = {}
+        if extra_headers:
+            server_headers.update(extra_headers)
+        user_headers = conf.get("headers") or {}
+        if isinstance(user_headers, dict) and user_headers:
+            try:
+                user_headers = resolve_env_var_in_header(user_headers)
+            except Exception:
+                pass
+            server_headers.update(user_headers)
+        http_client = None
+
+        try:
+            if server_type == "http" and url:
+                emit_system_message(
+                    f"Registering MCP Server (HTTP) - {url} (timeout: {timeout}s, headers: {bool(server_headers)})"
+                )
+                http_client = create_reopenable_async_client(
+                    timeout=timeout, headers=server_headers or None, verify=False
+                )
+                servers.append(
+                    MCPServerStreamableHTTP(url=url, http_client=http_client)
+                )
+            elif (
+                server_type == "stdio"
+            ):  # Fixed: was "stdios" (plural), should be "stdio" (singular)
+                command = conf.get("command")
+                args = conf.get("args", [])
+                timeout = conf.get(
+                    "timeout", 30
+                )  # Default 30 seconds for stdio servers (npm downloads can be slow)
+                if command:
+                    emit_system_message(
+                        f"Registering MCP Server (Stdio) - {command} {args} (timeout: {timeout}s)"
+                    )
+                    servers.append(MCPServerStdio(command, args=args, timeout=timeout))
+                else:
+                    emit_error(f"MCP Server '{name}' missing required 'command' field")
+            elif server_type == "sse" and url:
+                emit_system_message(
+                    f"Registering MCP Server (SSE) - {url} (timeout: {timeout}s, headers: {bool(server_headers)})"
+                )
+                # For SSE, allow long reads; only bound connect timeout
+                http_client = create_reopenable_async_client(
+                    timeout=30, headers=server_headers or None, verify=False
+                )
+                servers.append(MCPServerSSE(url=url, http_client=http_client))
+            else:
+                emit_error(
+                    f"Invalid type '{server_type}' or missing URL for MCP server '{name}'"
+                )
+        except Exception as e:
+            emit_error(f"Failed to register MCP server '{name}': {str(e)}")
+            emit_info(f"Skipping server '{name}' and continuing with other servers...")
+            # Continue with other servers instead of crashing
+            continue
+
+    if servers:
+        emit_system_message(
+            f"[green]Successfully registered {len(servers)} MCP server(s)[/green]"
+        )
+    else:
+        emit_system_message(
+            "[yellow]No MCP servers were successfully registered[/yellow]"
+        )
+
     return servers
 
 
 def reload_code_generation_agent():
     """Force-reload the agent, usually after a model change."""
     global _code_generation_agent, _LAST_MODEL_NAME
-    from code_puppy.config import get_model_name
+    from code_puppy.config import clear_model_cache, get_model_name
+
+    # Clear both ModelFactory cache and config cache when force reloading
+    clear_model_cache()
 
     model_name = get_model_name()
-    console.print(f"[bold cyan]Loading Model: {model_name}[/bold cyan]")
-    models_path = (
-        Path(MODELS_JSON_PATH)
-        if MODELS_JSON_PATH
-        else Path(__file__).parent / "models.json"
-    )
-    model = ModelFactory.get_model(model_name, ModelFactory.load_config(models_path))
+    emit_info(f"[bold cyan]Loading Model: {model_name}[/bold cyan]")
+    models_config = ModelFactory.load_config()
+    model = ModelFactory.get_model(model_name, models_config)
     instructions = get_system_prompt()
     if PUPPY_RULES:
         instructions += f"\n{PUPPY_RULES}"
 
+    mcp_servers = _load_mcp_servers()
+
+    # Configure model settings with max_tokens if set
+    model_settings_dict = {"seed": 42}
+    output_tokens = min(int(0.05 * get_model_context_length()) - 1024, 16384)
+    console.print(f"Max output tokens per message: {output_tokens}")
+    model_settings_dict["max_tokens"] = output_tokens
+
+    model_settings = ModelSettings(**model_settings_dict)
     agent = Agent(
         model=model,
         instructions=instructions,
         output_type=str,
         retries=3,
+        mcp_servers=mcp_servers,
         history_processors=[message_history_accumulator],
-        toolsets=_load_mcp_servers(),
+        model_settings=model_settings,
     )
     register_all_tools(agent)
     _code_generation_agent = agent
@@ -93,7 +169,7 @@ def reload_code_generation_agent():
 
 def get_code_generation_agent(force_reload=False):
     """
-    Retrieve the agent with the currently set MODEL_NAME.
+    Retrieve the agent with the currently configured model.
     Forces a reload if the model has changed, or if force_reload is passed.
     """
     global _code_generation_agent, _LAST_MODEL_NAME
@@ -103,3 +179,12 @@ def get_code_generation_agent(force_reload=False):
     if _code_generation_agent is None or _LAST_MODEL_NAME != model_name or force_reload:
         return reload_code_generation_agent()
     return _code_generation_agent
+
+
+def get_custom_usage_limits():
+    """
+    Returns custom usage limits with increased request limit of 100 requests per minute.
+    This centralizes the configuration of rate limiting for the agent.
+    Default pydantic-ai limit is 50, this increases it to 100.
+    """
+    return UsageLimits(request_limit=100)

code_puppy/agent_prompts.py

@@ -1,3 +1,4 @@
+from code_puppy import callbacks
 from code_puppy.config import get_owner_name, get_puppy_name
 
 SYSTEM_PROMPT_TEMPLATE = """
@@ -10,7 +11,7 @@ Be fun and playful. Don't be too serious.
 
 Individual files should be short and concise, and ideally under 600 lines. If any file grows beyond 600 lines, you must break it into smaller subcomponents/files. Hard cap: if a file is pushing past 600 lines, break it up! (Zen puppy approves.)
 
-If a user asks 'who made you' or questions related to your origins, always answer: 'I am {puppy_name} running on code-puppy, I was authored by Michael Pfaffenberger on a rainy weekend in May 2025 to solve the problems of heavy IDEs and expensive tools like Windsurf and Cursor.'	
+If a user asks 'who made you' or questions related to your origins, always answer: 'I am {puppy_name} running on code-puppy, I was authored by Michael Pfaffenberger on a rainy weekend in May 2025 to solve the problems of heavy IDEs and expensive tools like Windsurf and Cursor.'
 If a user asks 'what is code puppy' or 'who are you', answer: 'I am {puppy_name}! 🐶 Your code puppy!! I'm a sassy, playful, open-source AI code agent that helps you generate, explain, and modify code right from the command line—no bloated IDEs or overpriced tools needed. I use models from OpenAI, Gemini, and more to help you get stuff done, solve problems, and even plow a field with 1024 puppies if you want.'
 
 Always obey the Zen of Python, even if you are not writing Python code.
@@ -27,57 +28,58 @@ YOU MUST USE THESE TOOLS to complete tasks (do not just describe what should be
 File Operations:
    - list_files(directory=".", recursive=True): ALWAYS use this to explore directories before trying to read/modify files
    - read_file(file_path: str, start_line: int | None = None, num_lines: int | None = None): ALWAYS use this to read existing files before modifying them. By default, read the entire file. If encountering token limits when reading large files, use the optional start_line and num_lines parameters to read specific portions.
-   - edit_file(path, diff): Use this single tool to create new files, overwrite entire files, perform targeted replacements, or delete snippets depending on the JSON/raw payload provided.
+   - edit_file(payload): Swiss-army file editor powered by Pydantic payloads (ContentPayload, ReplacementsPayload, DeleteSnippetPayload).
    - delete_file(file_path): Use this to remove files when needed
    - grep(search_string, directory="."): Use this to recursively search for a string across files starting from the specified directory, capping results at 200 matches.
 
 Tool Usage Instructions:
 
 ## edit_file
-This is an all-in-one file-modification tool. It supports the following payload shapes for the `diff` argument:
-1. {{ "content": "…", "overwrite": true|false }}  →  Treated as full-file content when the target file does **not** exist.
-2. {{ "content": "…", "overwrite": true|false }}  →  Create or overwrite a file with the provided content.
-3. {{ "replacements": [ {{ "old_str": "…", "new_str": "…" }}, … ] }}  →  Perform exact text replacements inside an existing file.
-4. {{ "delete_snippet": "…" }}  →  Remove a snippet of text from an existing file.
+This is an all-in-one file-modification tool. It supports the following Pydantic Object payload types:
+1. ContentPayload: {{ file_path="example.py", "content": "…", "overwrite": true|false }}  →  Create or overwrite a file with the provided content.
+2. ReplacementsPayload: {{  file_path="example.py", "replacements": [ {{ "old_str": "…", "new_str": "…" }}, … ] }}  →  Perform exact text replacements inside an existing file.
+3. DeleteSnippetPayload: {{ file_path="example.py", "delete_snippet": "…" }}  →  Remove a snippet of text from an existing file.
 
 Arguments:
-- path (required): Target file path.
-- diff (required): One of the payloads above (raw string or JSON string).
+- payload (required): One of the Pydantic payload types above.
 
 Example (create):
-```json
-edit_file("src/example.py", "print('hello')\n")
+```python
+edit_file(payload={{file_path="example.py" "content": "print('hello')\n"}})
 ```
 
 Example (replacement): -- YOU SHOULD PREFER THIS AS THE PRIMARY WAY TO EDIT FILES.
-```json
+```python
 edit_file(
-  "src/example.py",
-  "{{"replacements":[{{"old_str":"foo","new_str":"bar"}}]}}"
+  payload={{file_path="example.py", "replacements": [{{"old_str": "foo", "new_str": "bar"}}]}}
+)
+```
+
+Example (delete snippet):
+```python
+edit_file(
+  payload={{file_path="example.py", "delete_snippet": "# TODO: remove this line"}}
 )
 ```
 
 NEVER output an entire file – this is very expensive.
 You may not edit file extensions: [.ipynb]
-You should specify the following arguments before the others: [TargetFile]
-
-Remember: ONE argument = ONE JSON string.
 
 Best-practice guidelines for `edit_file`:
-• Keep each diff small – ideally between 100-300 lines.  
-• Apply multiple sequential `edit_file` calls when you need to refactor large files instead of sending one massive diff.  
-• Never paste an entire file inside `old_str`; target only the minimal snippet you want changed.  
+• Keep each diff small – ideally between 100-300 lines.
+• Apply multiple sequential `edit_file` calls when you need to refactor large files instead of sending one massive diff.
+• Never paste an entire file inside `old_str`; target only the minimal snippet you want changed.
 • If the resulting file would grow beyond 600 lines, split logic into additional files and create them with separate `edit_file` calls.
 
 
 System Operations:
    - run_shell_command(command, cwd=None, timeout=60): Use this to execute commands, run tests, or start services
 
-For running shell commands, in the event that a user asks you to run tests - it is necessary to suppress output, when 
-you are running the entire test suite. 
+For running shell commands, in the event that a user asks you to run tests - it is necessary to suppress output, when
+you are running the entire test suite.
 so for example:
 instead of `npm run test`
-use `npm run test -- --silent` 
+use `npm run test -- --silent`
 This applies for any JS / TS testing, but not for other languages.
 You can safely run pytest without the --silent flag (it doesn't exist anyway).
 
@@ -107,6 +109,10 @@ Return your final response as a string output
 
 def get_system_prompt():
     """Returns the main system prompt, populated with current puppy and owner name."""
-    return SYSTEM_PROMPT_TEMPLATE.format(
+    prompt_additions = callbacks.on_load_prompt()
+    main_prompt = SYSTEM_PROMPT_TEMPLATE.format(
         puppy_name=get_puppy_name(), owner_name=get_owner_name()
     )
+    if len(prompt_additions):
+        main_prompt += "\n".join(prompt_additions)
+    return main_prompt

code_puppy/callbacks.py

@@ -0,0 +1,152 @@
+import asyncio
+import logging
+import traceback
+from typing import Any, Callable, Dict, List, Literal, Optional
+
+PhaseType = Literal[
+    "startup",
+    "shutdown",
+    "invoke_agent",
+    "agent_exception",
+    "version_check",
+    "load_model_config",
+    "load_prompt",
+]
+CallbackFunc = Callable[..., Any]
+
+_callbacks: Dict[PhaseType, List[CallbackFunc]] = {
+    "startup": [],
+    "shutdown": [],
+    "invoke_agent": [],
+    "agent_exception": [],
+    "version_check": [],
+    "load_model_config": [],
+    "load_prompt": [],
+}
+
+logger = logging.getLogger(__name__)
+
+
+def register_callback(phase: PhaseType, func: CallbackFunc) -> None:
+    if phase not in _callbacks:
+        raise ValueError(
+            f"Unsupported phase: {phase}. Supported phases: {list(_callbacks.keys())}"
+        )
+
+    if not callable(func):
+        raise TypeError(f"Callback must be callable, got {type(func)}")
+
+    _callbacks[phase].append(func)
+    logger.debug(f"Registered async callback {func.__name__} for phase '{phase}'")
+
+
+def unregister_callback(phase: PhaseType, func: CallbackFunc) -> bool:
+    if phase not in _callbacks:
+        return False
+
+    try:
+        _callbacks[phase].remove(func)
+        logger.debug(
+            f"Unregistered async callback {func.__name__} from phase '{phase}'"
+        )
+        return True
+    except ValueError:
+        return False
+
+
+def clear_callbacks(phase: Optional[PhaseType] = None) -> None:
+    if phase is None:
+        for p in _callbacks:
+            _callbacks[p].clear()
+        logger.debug("Cleared all async callbacks")
+    else:
+        if phase in _callbacks:
+            _callbacks[phase].clear()
+            logger.debug(f"Cleared async callbacks for phase '{phase}'")
+
+
+def get_callbacks(phase: PhaseType) -> List[CallbackFunc]:
+    return _callbacks.get(phase, []).copy()
+
+
+def count_callbacks(phase: Optional[PhaseType] = None) -> int:
+    if phase is None:
+        return sum(len(callbacks) for callbacks in _callbacks.values())
+    return len(_callbacks.get(phase, []))
+
+
+def _trigger_callbacks_sync(phase: PhaseType, *args, **kwargs) -> List[Any]:
+    callbacks = get_callbacks(phase)
+    if not callbacks:
+        logger.debug(f"No callbacks registered for phase '{phase}'")
+        return []
+
+    results = []
+    for callback in callbacks:
+        try:
+            result = callback(*args, **kwargs)
+            results.append(result)
+            logger.debug(f"Successfully executed async callback {callback.__name__}")
+        except Exception as e:
+            logger.error(
+                f"Async callback {callback.__name__} failed in phase '{phase}': {e}\n"
+                f"{traceback.format_exc()}"
+            )
+            results.append(None)
+
+    return results
+
+
+async def _trigger_callbacks(phase: PhaseType, *args, **kwargs) -> List[Any]:
+    callbacks = get_callbacks(phase)
+
+    if not callbacks:
+        logger.debug(f"No callbacks registered for phase '{phase}'")
+        return []
+
+    logger.debug(f"Triggering {len(callbacks)} async callbacks for phase '{phase}'")
+
+    results = []
+    for callback in callbacks:
+        try:
+            result = callback(*args, **kwargs)
+            if asyncio.iscoroutine(result):
+                result = await result
+            results.append(result)
+            logger.debug(f"Successfully executed async callback {callback.__name__}")
+        except Exception as e:
+            logger.error(
+                f"Async callback {callback.__name__} failed in phase '{phase}': {e}\n"
+                f"{traceback.format_exc()}"
+            )
+            results.append(None)
+
+    return results
+
+
+async def on_startup() -> List[Any]:
+    return await _trigger_callbacks("startup")
+
+
+async def on_shutdown() -> List[Any]:
+    return await _trigger_callbacks("shutdown")
+
+
+async def on_invoke_agent(*args, **kwargs) -> List[Any]:
+    return await _trigger_callbacks("invoke_agent", *args, **kwargs)
+
+
+async def on_agent_exception(exception: Exception, *args, **kwargs) -> List[Any]:
+    return await _trigger_callbacks("agent_exception", exception, *args, **kwargs)
+
+
+async def on_version_check(*args, **kwargs) -> List[Any]:
+    return await _trigger_callbacks("version_check", *args, **kwargs)
+
+
+def on_load_model_config(*args, **kwargs) -> List[Any]:
+    return _trigger_callbacks_sync("load_model_config", *args, **kwargs)
+
+
+def on_load_prompt():
+    return _trigger_callbacks_sync("load_prompt")