lambda-agent 0.1.0__py3-none-any.whl

lambda_agent/main.py

@@ -0,0 +1,318 @@
+from .agent import Agent, TokenUsage
+from . import config
+from .spinner import console
+import os
+import getpass
+from pathlib import Path
+
+from rich.panel import Panel
+from rich.text import Text
+from rich.rule import Rule
+from rich.markdown import Markdown
+from rich.prompt import Prompt
+from rich import box
+from rich.align import Align
+from rich.table import Table
+
+
+BANNER = r"""
+██╗      █████╗ ███╗   ███╗██████╗ ██████╗  █████╗
+██║     ██╔══██╗████╗ ████║██╔══██╗██╔══██╗██╔══██╗
+██║     ███████║██╔████╔██║██████╔╝██║  ██║███████║
+██║     ██╔══██║██║╚██╔╝██║██╔══██╗██║  ██║██╔══██║
+███████╗██║  ██║██║ ╚═╝ ██║██████╔╝██████╔╝██║  ██║
+╚══════╝╚═╝  ╚═╝╚═╝     ╚═╝╚═════╝ ╚═════╝ ╚═╝  ╚═╝
+"""
+
+
+SLASH_COMMANDS = {
+    "/models": "List available models and switch between them",
+    "/config": "Update API key and save to config",
+    "/help": "Show available slash commands",
+}
+
+
+def print_banner():
+    banner_text = Text(BANNER, style="bold cyan", justify="center")
+    subtitle = Text(
+        "  Minimal AI Coding Agent  ·  Type '/help' for commands  ",
+        style="dim white",
+        justify="center",
+    )
+
+    panel = Panel(
+        Align.center(Text.assemble(banner_text, "\n", subtitle)),
+        border_style="cyan",
+        box=box.DOUBLE_EDGE,
+        padding=(0, 2),
+    )
+    console.print(panel)
+
+
+def print_user_message(text: str):
+    label = Text(" YOU ", style="bold black on bright_yellow")
+    content = Text(f"  {text}", style="bright_white")
+    console.print()
+    console.print(Text.assemble(label, content))
+
+
+def print_lambda_message(text: str):
+    console.print()
+    label = Text(" LAMBDA ", style="bold black on cyan")
+    console.print(label)
+    console.print(
+        Panel(
+            Markdown(text),
+            border_style="cyan",
+            box=box.ROUNDED,
+            padding=(0, 2),
+        )
+    )
+
+
+def print_token_stats(turn: TokenUsage, session: TokenUsage):
+    """Render a compact token usage line under the Lambda response."""
+    console.print(
+        Text.assemble(
+            ("  ▶ tokens  ", "dim"),
+            ("this turn: ", "dim"),
+            (f"↑{turn.prompt:,}", "dim cyan"),
+            (" in  ", "dim"),
+            (f"↓{turn.completion:,}", "dim cyan"),
+            (" out     ", "dim"),
+            ("session total: ", "dim"),
+            (f"{session.total:,}", "bold cyan"),
+            (" tokens", "dim"),
+        )
+    )
+
+
+def handle_models_command(agent: Agent):
+    """Display available models and let the user pick one."""
+    table = Table(
+        title="Available Models",
+        title_style="bold cyan",
+        border_style="cyan",
+        box=box.ROUNDED,
+        padding=(0, 2),
+    )
+    table.add_column("#", style="dim", width=4)
+    table.add_column("Model", style="white")
+    table.add_column("Status", justify="center")
+
+    models = config.AVAILABLE_MODELS
+    for i, model in enumerate(models, 1):
+        is_active = model == agent.model_name
+        status = "[bold green]● active[/bold green]" if is_active else "[dim]—[/dim]"
+        name_style = "bold cyan" if is_active else "white"
+        table.add_row(str(i), f"[{name_style}]{model}[/{name_style}]", status)
+
+    console.print()
+    console.print(table)
+    console.print()
+
+    choice = Prompt.ask(
+        "[bold bright_yellow]  Select model #[/bold bright_yellow] (or Enter to cancel)",
+        default="",
+        console=console,
+    )
+
+    if not choice.strip():
+        console.print("  [dim]No change.[/dim]")
+        return
+
+    try:
+        idx = int(choice) - 1
+        if 0 <= idx < len(models):
+            selected = models[idx]
+            if selected == agent.model_name:
+                console.print(f"  [dim]Already using[/dim] [cyan]{selected}[/cyan]")
+            else:
+                msg = agent.switch_model(selected)
+                console.print(f"  {msg}")
+        else:
+            console.print("  [red]Invalid selection.[/red]")
+    except ValueError:
+        console.print("  [red]Please enter a number.[/red]")
+
+
+def handle_help_command():
+    """Show available slash commands."""
+    table = Table(
+        title="Slash Commands",
+        title_style="bold cyan",
+        border_style="cyan",
+        box=box.ROUNDED,
+        padding=(0, 1),
+    )
+    table.add_column("Command", style="bold bright_yellow", min_width=12)
+    table.add_column("Description", style="white")
+
+    for cmd, desc in SLASH_COMMANDS.items():
+        table.add_row(cmd, desc)
+
+    # Also list the built-in exit commands
+    table.add_row("exit / quit", "End the session")
+
+    console.print()
+    console.print(table)
+
+
+def handle_config_command(agent: Agent):
+    """Let the user update their API key mid-session."""
+    console.print()
+    console.print(
+        Panel(
+            Text.assemble(
+                ("Current API key: ", "dim"),
+                (f"{config.API_KEY[:8]}...{config.API_KEY[-4:]}", "cyan"),
+            ),
+            border_style="cyan",
+            box=box.ROUNDED,
+            title="[bold cyan]⚙ Configuration[/bold cyan]",
+            title_align="left",
+        )
+    )
+    console.print()
+
+    new_key = getpass.getpass("  Enter new API key (or press Enter to keep current): ")
+
+    if not new_key.strip():
+        console.print("  [dim]No change.[/dim]")
+        return
+
+    # Update in-memory config
+    config.API_KEY = new_key.strip()
+    os.environ["API_KEY"] = config.API_KEY
+
+    # Re-create the API client with the new key
+    from google import genai
+
+    agent.client = genai.Client(api_key=config.API_KEY)
+
+    # Re-create the chat session so the new client is used
+    from google.genai import types
+    from .tools import TOOL_FUNCTIONS
+
+    agent.chat_session = agent.client.chats.create(
+        model=agent.model_name,
+        config=types.GenerateContentConfig(
+            system_instruction=agent.system_instruction,
+            tools=TOOL_FUNCTIONS,
+            automatic_function_calling=types.AutomaticFunctionCallingConfig(
+                disable=True
+            ),
+        ),
+    )
+    agent.is_first_message = True
+
+    # Persist to config file
+    config_file = Path.home() / ".config" / "lambda-agent" / "config.env"
+    try:
+        config_file.parent.mkdir(parents=True, exist_ok=True)
+        with open(config_file, "w") as f:
+            f.write(f"API_KEY={config.API_KEY}\n")
+            f.write(f"MODEL_NAME={config.MODEL_NAME}\n")
+        os.chmod(config_file, 0o600)
+        console.print("  [green]✓[/green] API key updated and saved to config.")
+    except Exception as e:
+        console.print("  [green]✓[/green] API key updated in memory.")
+        console.print(f"  [yellow]⚠[/yellow] Could not save to disk: {e}")
+
+
+def main():
+    print_banner()
+
+    try:
+        if not config.API_KEY:
+            from .cli_setup import run_setup
+
+            config.API_KEY, config.MODEL_NAME = run_setup()
+            os.environ["API_KEY"] = config.API_KEY
+            os.environ["MODEL_NAME"] = config.MODEL_NAME
+
+        agent = Agent()
+
+        console.print(
+            Rule("[bold cyan]Session Started[/bold cyan]", style="cyan"),
+        )
+
+        while True:
+            try:
+                # Styled prompt — uses plain input to keep cursor on same line
+                user_input = Prompt.ask(
+                    "\n[bold bright_yellow]  You[/bold bright_yellow]",
+                    console=console,
+                )
+
+                if user_input.lower() in ["exit", "quit"]:
+                    console.print()
+                    # Show session token summary before quitting
+                    if agent.token_usage.total > 0:
+                        console.print(
+                            Panel(
+                                Text.assemble(
+                                    ("Session token usage\n", "bold white"),
+                                    ("  Prompt (in):      ", "dim"),
+                                    (f"{agent.token_usage.prompt:>10,}\n", "cyan"),
+                                    ("  Completion (out): ", "dim"),
+                                    (f"{agent.token_usage.completion:>10,}\n", "cyan"),
+                                    ("  Total:            ", "dim"),
+                                    (f"{agent.token_usage.total:>10,}", "bold cyan"),
+                                ),
+                                border_style="cyan",
+                                box=box.ROUNDED,
+                                title="[bold cyan]⚡ Token Summary[/bold cyan]",
+                                title_align="left",
+                            )
+                        )
+                    console.print(
+                        Panel(
+                            "[bold cyan]Goodbye! Lambda signing off.[/bold cyan]",
+                            border_style="cyan",
+                            box=box.ROUNDED,
+                        )
+                    )
+                    break
+
+                if not user_input.strip():
+                    continue
+
+                # Handle slash commands
+                if user_input.strip().lower() == "/models":
+                    handle_models_command(agent)
+                    continue
+                elif user_input.strip().lower() == "/config":
+                    handle_config_command(agent)
+                    continue
+                elif user_input.strip().lower() == "/help":
+                    handle_help_command()
+                    continue
+                elif user_input.strip().startswith("/"):
+                    console.print(
+                        f"  [red]Unknown command:[/red] {user_input.strip()}  "
+                        "[dim]Type /help for available commands.[/dim]"
+                    )
+                    continue
+
+                response, turn_usage = agent.chat(user_input)
+                print_lambda_message(response)
+                print_token_stats(turn_usage, agent.token_usage)
+
+            except KeyboardInterrupt:
+                console.print()
+                console.print("[bold cyan]\nGoodbye![/bold cyan]")
+                break
+
+    except Exception as e:
+        console.print(
+            Panel(
+                f"[bold red]Failed to initialize Lambda:[/bold red]\n{str(e)}",
+                border_style="red",
+                box=box.ROUNDED,
+            )
+        )
+
+
+if __name__ == "__main__":
+    main()

lambda_agent/scratchpad.py

@@ -0,0 +1,139 @@
+"""
+Scratchpad Module
+=================
+Provides tools for the agent to maintain a persistent, human-readable scratchpad
+file (.lambda_scratchpad.md) in the user's working directory.
+
+The scratchpad lets the agent plan complex tasks, track progress, and keep
+implementation notes — all visible to the user as a markdown file in the repo.
+"""
+
+import os
+from datetime import datetime
+
+AGENT_DIR = ".agent"
+SCRATCHPAD_FILE = os.path.join(AGENT_DIR, "scratchpad.md")
+
+_HEADER_TEMPLATE = """\
+<!-- This file is managed by the Lambda coding agent. -->
+<!-- Feel free to read it, but edits may be overwritten by the agent. -->
+
+# Lambda Scratchpad
+
+"""
+
+
+def _ensure_scratchpad() -> str:
+    """Return the absolute path to the scratchpad, creating it if it doesn't exist."""
+    agent_dir = os.path.abspath(AGENT_DIR)
+    os.makedirs(agent_dir, exist_ok=True)
+    path = os.path.abspath(SCRATCHPAD_FILE)
+    if not os.path.exists(path):
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(_HEADER_TEMPLATE)
+    return path
+
+
+def read_scratchpad() -> str:
+    """Reads the full contents of the Lambda scratchpad file (.lambda_scratchpad.md).
+
+    Use this to recall your previous plans, task lists, and implementation notes.
+    """
+    path = _ensure_scratchpad()
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            return f.read()
+    except Exception as e:
+        return f"Error reading scratchpad: {e}"
+
+
+def write_scratchpad(content: str) -> str:
+    """Overwrites the entire Lambda scratchpad file with the provided content.
+
+    Use this when you need to replace the scratchpad with a fresh plan or when
+    starting a new major task. For incremental updates, prefer update_scratchpad.
+
+    Args:
+        content: The full markdown content to write to the scratchpad.
+    """
+    path = _ensure_scratchpad()
+    try:
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(_HEADER_TEMPLATE + content)
+        return f"Scratchpad written successfully → {path}"
+    except Exception as e:
+        return f"Error writing scratchpad: {e}"
+
+
+def update_scratchpad(note: str, section: str = "Notes") -> str:
+    """Appends a timestamped note to a specific section in the scratchpad.
+
+    This is ideal for incrementally logging progress, decisions, and discoveries
+    without replacing existing content.
+
+    Args:
+        note: The text to append (supports markdown).
+        section: The section heading to append under (e.g. 'Plan', 'Progress', 'Notes').
+    """
+    path = _ensure_scratchpad()
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            existing = f.read()
+
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M")
+        entry = f"\n- **[{timestamp}]** {note}"
+
+        section_heading = f"## {section}"
+        if section_heading in existing:
+            # Append under the existing section
+            parts = existing.split(section_heading, 1)
+            # Find the next section heading (##) or end of file
+            rest = parts[1]
+            next_section = rest.find("\n## ")
+            if next_section == -1:
+                # No next section — just append at the end
+                updated = existing + entry
+            else:
+                # Insert before the next section
+                insert_pos = len(parts[0]) + len(section_heading) + next_section
+                updated = existing[:insert_pos] + entry + "\n" + existing[insert_pos:]
+        else:
+            # Create the section at the end
+            updated = existing.rstrip() + f"\n\n{section_heading}\n{entry}\n"
+
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(updated)
+
+        return f"Scratchpad updated (section: {section}) → {path}"
+    except Exception as e:
+        return f"Error updating scratchpad: {e}"
+
+
+def clear_scratchpad() -> str:
+    """Clears the scratchpad, resetting it to a blank state.
+
+    Use this when a major task is fully complete and the scratchpad is no longer needed.
+    """
+    path = _ensure_scratchpad()
+    try:
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(_HEADER_TEMPLATE)
+        return f"Scratchpad cleared → {path}"
+    except Exception as e:
+        return f"Error clearing scratchpad: {e}"
+
+
+# Tool registrations for the agent
+SCRATCHPAD_EXECUTORS = {
+    "read_scratchpad": read_scratchpad,
+    "write_scratchpad": write_scratchpad,
+    "update_scratchpad": update_scratchpad,
+    "clear_scratchpad": clear_scratchpad,
+}
+
+SCRATCHPAD_FUNCTIONS = [
+    read_scratchpad,
+    write_scratchpad,
+    update_scratchpad,
+    clear_scratchpad,
+]

lambda_agent/spinner.py

@@ -0,0 +1,50 @@
+import random
+from rich.console import Console
+from rich.spinner import Spinner as RichSpinner
+from rich.live import Live
+from rich.text import Text
+
+# Shared console instance used across the whole app
+console = Console()
+
+QUOTES = [
+    "Consulting the mainframe…",
+    "Synthesizing logic…",
+    "Bending the matrix…",
+    "Drinking virtual coffee…",
+    "Compiling thoughts…",
+    "Evaluating your code…",
+    "Simulating outcomes…",
+    "Reversing the polarity…",
+    "Aligning the vectors…",
+    "Traversing the graph…",
+    "Reading your source code…",
+    "Assembling the bytes…",
+]
+
+
+class Spinner:
+    """Context manager that shows a rich animated spinner while Lambda is thinking."""
+
+    def __init__(self, message: str = "Thinking"):
+        self._label = random.choice(QUOTES) if random.random() > 0.3 else message
+        self._live = None
+
+    def __enter__(self):
+        renderable = RichSpinner(
+            "dots",
+            text=Text(f" {self._label}", style="dim cyan italic"),
+            style="bold cyan",
+        )
+        self._live = Live(
+            renderable,
+            console=console,
+            refresh_per_second=15,
+            transient=True,
+        )
+        self._live.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self._live:
+            self._live.__exit__(exc_type, exc_val, exc_tb)