main.py

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "openai",
#     "pydantic",
# ]
# ///

"""Single-file CLI agent that surfaces OpenAI Responses reasoning in a TUI."""

import os
import sys
import json
import argparse
import logging
from dataclasses import dataclass, field
from typing import Callable, List, Dict, Any, Optional, Tuple
from openai import OpenAI  # type: ignore
from pydantic import BaseModel  # type: ignore


# Maximum characters returned per read_file invocation.
MAX_READ_CHARS = 4000
# Maximum number of characters shown for thinking/reasoning text.
MAX_THINK_TEXT = 500


@dataclass
class ChatEvent:
    """Represents an event emitted during the assistant's reasoning loop."""

    kind: str
    data: Dict[str, Any]


@dataclass
class ToolInvocation:
    """Captures a fully-specified tool call to execute."""

    name: str
    call_id: str
    arguments: Dict[str, Any]


@dataclass
class ToolCallState:
    """Tracks partial streaming state for a tool call."""

    item_id: str
    call_id: str = ""
    name: str = ""
    _buffer: List[str] = field(default_factory=list)
    arguments: Dict[str, Any] = field(default_factory=dict)
    parsed: bool = False
    emitted: bool = False
    dispatched: bool = False

    def append(self, delta: str) -> None:
        if delta:
            self._buffer.append(delta)
            self.parsed = False

    def set_full(self, payload: str) -> None:
        self._buffer = [payload] if payload else []
        self.parsed = False

    def ensure_arguments(self, decoder: Callable[[str], Dict[str, Any]]) -> None:
        if self.parsed:
            return
        raw = "".join(self._buffer)
        self.arguments = decoder(raw)
        self.parsed = True


@dataclass
class ToolCallRegistry:
    """Index tool call states by item ID and call ID for quick lookups."""

    by_item: Dict[str, ToolCallState] = field(default_factory=dict)
    by_call: Dict[str, ToolCallState] = field(default_factory=dict)

    def get(self, item_id: str, call_id: str) -> Optional[ToolCallState]:
        item_key = item_id.strip()
        call_key = call_id.strip()
        if item_key and item_key in self.by_item:
            return self.by_item[item_key]
        if call_key and call_key in self.by_call:
            return self.by_call[call_key]
        return None

    def put(self, state: ToolCallState, item_id: str, call_id: str) -> None:
        item_key = item_id.strip()
        call_key = call_id.strip()
        if item_key:
            self.by_item[item_key] = state
            state.item_id = item_key
        if call_key:
            self.by_call[call_key] = state
            state.call_id = call_key

    def states(self) -> List[ToolCallState]:
        seen: Dict[int, ToolCallState] = {}
        for state in self.by_item.values():
            seen[id(state)] = state
        for state in self.by_call.values():
            seen.setdefault(id(state), state)
        return list(seen.values())

# Route logs to a file so the CLI remains clean.
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(message)s",
    handlers=[logging.FileHandler("agent.log")],
)

# Suppress verbose HTTP logs
logging.getLogger("httpcore").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)


def read_file(
    path: str, offset: Optional[int] = None, length: Optional[int] = None
) -> str:
    """Read a portion of a file, supporting pagination via offset/length."""
    try:
        with open(path, "r", encoding="utf-8") as handle:
            content = handle.read()

        total_len = len(content)
        start = max(0, offset or 0)
        if start >= total_len:
            return f"File {path} has {total_len} characters; offset {start} is beyond the end."

        requested = length if (length and length > 0) else MAX_READ_CHARS
        end = min(total_len, start + requested)
        snippet = content[start:end]
        truncated = end < total_len

        header = (
            f"File {path}: returning characters {start}-{end} of {total_len}."
            + (" Additional content remains." if truncated else "")
        )

        footer = (
            f"\nTruncated output. Call read_file with offset {end} to continue."
            if truncated
            else ""
        )

        return f"{header}\n{snippet}{footer}"
    except FileNotFoundError:
        return f"File not found: {path}"
    except Exception as exc:
        return f"Error reading file: {str(exc)}"


def list_files(path: str = ".") -> str:
    """List files and directories with explicit markers."""
    try:
        if not os.path.exists(path):
            return f"Path not found: {path}"

        items = []
        for item in sorted(os.listdir(path)):
            item_path = os.path.join(path, item)
            if os.path.isdir(item_path):
                items.append(f"[DIR]  {item}/")
            else:
                items.append(f"[FILE] {item}")

        if not items:
            return f"Empty directory: {path}"

        return f"Contents of {path}:\n" + "\n".join(items)
    except Exception as exc:
        return f"Error listing files: {str(exc)}"


def edit_file(path: str, old_text: str, new_text: str) -> str:
    """Edit an existing file in-place or create a new file when needed."""
    try:
        if os.path.exists(path) and old_text:
            with open(path, "r", encoding="utf-8") as handle:
                content = handle.read()

            if old_text not in content:
                return f"Text not found in file: {old_text}"

            updated = content.replace(old_text, new_text)
            with open(path, "w", encoding="utf-8") as handle:
                handle.write(updated)

            return f"Successfully edited {path}"

        dir_name = os.path.dirname(path)
        if dir_name:
            os.makedirs(dir_name, exist_ok=True)

        with open(path, "w", encoding="utf-8") as handle:
            handle.write(new_text)

        return f"Successfully created {path}"
    except Exception as exc:
        return f"Error editing file: {str(exc)}"


class Tool(BaseModel):
    """Pydantic representation of a tool definition exposed to the model."""

    name: str
    description: str
    input_schema: Dict[str, Any]


class AIAgent:
    """Coordinates conversations with OpenAI and manages tool execution."""

    def __init__(self, api_key: str):
        """Initialize the client connection and hydrate tool metadata."""
        self.client = OpenAI(api_key=api_key)
        self.tools: List[Tool] = []
        self.tool_definitions: List[Dict[str, Any]] = []
        self.instructions = (
            "You are a helpful coding assistant operating in a terminal environment. "
            "Output only plain text without markdown formatting, as your responses appear "
            "directly in the terminal. Be concise but thorough, providing clear, concise "
            "and practical advice with a friendly tone. Don't use any asterisk characters in "
            "your responses. "
            f"To conserve context, the read_file tool returns at most {MAX_READ_CHARS} "
            "characters per call; use its optional offset and length arguments to page "
            "through larger files when needed."
        )
        self.previous_response_id: Optional[str] = None
        self._setup_tools()

    def _setup_tools(self):
        """Populate the tool registry and build OpenAI-compatible schemas."""
        self.tools = [
            Tool(
                name="read_file",
                description="Read the contents of a file at the specified path",
                input_schema={
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "The path to the file to read",
                        },
                        "offset": {
                            "type": "integer",
                            "description": "Optional character offset to start reading from",
                            "minimum": 0,
                        },
                        "length": {
                            "type": "integer",
                            "description": f"Optional max characters to return (defaults to {MAX_READ_CHARS})",
                            "minimum": 1,
                        },
                    },
                    "required": ["path"],
                },
            ),
            Tool(
                name="list_files",
                description="List all files and directories in the specified path",
                input_schema={
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "The directory path to list (defaults to current directory)",
                        }
                    },
                    "required": [],
                },
            ),
            Tool(
                name="edit_file",
                description="Edit a file by replacing old_text with new_text. Creates the file if it doesn't exist.",
                input_schema={
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "The path to the file to edit",
                        },
                        "old_text": {
                            "type": "string",
                            "description": "The text to search for and replace (leave empty to create new file)",
                        },
                        "new_text": {
                            "type": "string",
                            "description": "The text to replace old_text with",
                        },
                    },
                    "required": ["path", "new_text"],
                },
            ),
        ]
        self.tool_definitions = [
            {
                "type": "function",
                "name": tool.name,
                "description": tool.description,
                "parameters": tool.input_schema,
            }
            for tool in self.tools
        ]

    def _trim_text(self, text: str) -> str:
        """Clamp reasoning text to the configured MAX_THINK_TEXT length."""
        cleaned = text.strip()
        if len(cleaned) <= MAX_THINK_TEXT:
            return cleaned
        return cleaned[: MAX_THINK_TEXT - 3].rstrip() + "..."

    def _summarize_arguments(self, args: Dict[str, Any]) -> str:
        """Render tool arguments into a concise, deterministic summary."""
        if not args:
            return "(no arguments)"

        parts = []
        for key, value in args.items():
            parts.append(f"{key}={self._format_arg_value(key, value)}")
        return ", ".join(parts)

    def _format_arg_value(self, key: str, value: Any) -> str:
        """Normalize individual argument values for CLI presentation."""
        if isinstance(value, str):
            if key in {"new_text", "old_text"}:
                return f"<str len {len(value)}>"
            if len(value) <= 60:
                return value
            return value[:57] + "..."
        if isinstance(value, (int, float)):
            return str(value)
        if value is None:
            return "None"
        return "<complex>"

    def _summarize_tool_result(self, result: str) -> str:
        """Capture the first line of tool output for progress updates."""
        first_line = result.splitlines()[0] if result else ""
        return first_line if first_line else "(no result text)"

    def _decode_arguments(self, raw: str) -> Dict[str, Any]:
        """Safely decode JSON argument payloads from streaming deltas."""
        if not raw or not raw.strip():
            return {}

        try:
            parsed = json.loads(raw)
        except json.JSONDecodeError as exc:
            logging.error("Error decoding tool arguments: %s", exc)
            return {}

        if not isinstance(parsed, dict):
            return {}

        return parsed

    def _stream_response(
        self,
        input_payload: List[Dict[str, Any]],
        previous_response_id: Optional[str],
        emit: Callable[[ChatEvent], None],
    ) -> Tuple[Any, List[ToolInvocation]]:
        """Stream a response, emitting intermediate events immediately."""

        registry = ToolCallRegistry()
        tool_calls: List[ToolInvocation] = []
        reasoning_buffer: List[str] = []
        last_thinking = ""
        message_buffer: List[str] = []
        final_response: Any = None

        def emit_thinking() -> None:
            nonlocal last_thinking
            summary = self._trim_text("".join(reasoning_buffer))
            if not summary or summary == last_thinking:
                return
            emit(ChatEvent(kind="thinking", data={"text": summary}))
            last_thinking = summary

        def emit_assistant_chunk(chunk: str, reset: bool) -> None:
            if not chunk:
                return
            data = {"chunk": chunk}
            if reset:
                data["reset"] = "true"
            emit(ChatEvent(kind="assistant_delta", data=data))

        stream_kwargs: Dict[str, Any] = {
            "model": "gpt-5",
            "instructions": self.instructions,
            "input": input_payload,
            "tools": self.tool_definitions,
            "reasoning": {"summary": "auto"},
        }
        if previous_response_id:
            stream_kwargs["previous_response_id"] = previous_response_id

        with self.client.responses.stream(**stream_kwargs) as stream:
            for event in stream:
                event_type = getattr(event, "type", "")

                if event_type == "response.reasoning_summary_text.delta":
                    reasoning_buffer.append(getattr(event, "delta", ""))
                    emit_thinking()
                elif event_type == "response.reasoning_summary_text.done":
                    reasoning_buffer.clear()
                    reasoning_buffer.append(getattr(event, "text", ""))
                    emit_thinking()
                elif event_type == "response.reasoning_summary_part.added":
                    part = getattr(event, "part", None)
                    text = getattr(part, "text", "") if part else ""
                    if text:
                        reasoning_buffer.append(text)
                        emit_thinking()
                elif event_type == "response.reasoning_summary_part.done":
                    part = getattr(event, "part", None)
                    text = getattr(part, "text", "") if part else ""
                    if text:
                        reasoning_buffer.append(text)
                        emit_thinking()
                elif event_type == "response.output_text.delta":
                    chunk = getattr(event, "delta", "")
                    if chunk:
                        message_buffer.append(chunk)
                        emit_assistant_chunk(chunk, reset=False)
                elif event_type == "response.output_text.done":
                    text = getattr(event, "text", "")
                    if text:
                        current = "".join(message_buffer)
                        if text.startswith(current):
                            extra = text[len(current) :]
                            if extra:
                                message_buffer.append(extra)
                                emit_assistant_chunk(extra, reset=False)
                        else:
                            message_buffer.clear()
                            message_buffer.append(text)
                            emit_assistant_chunk(text, reset=True)
                elif event_type == "response.output_item.added":
                    item = getattr(event, "item", None)
                    if item and getattr(item, "type", "") == "function_call":
                        state = registry.get(getattr(item, "id", "") or "", getattr(item, "call_id", "") or "")
                        if state is None:
                            state = ToolCallState(item_id=getattr(item, "id", "") or "")
                        registry.put(state, getattr(item, "id", "") or "", getattr(item, "call_id", "") or "")
                        if getattr(item, "name", ""):
                            state.name = item.name
                        if getattr(item, "arguments", ""):
                            state.set_full(item.arguments)
                elif event_type == "response.function_call_arguments.delta":
                    state = registry.get(getattr(event, "item_id", ""), "")
                    if state is None:
                        state = ToolCallState(item_id=getattr(event, "item_id", ""))
                        registry.put(state, getattr(event, "item_id", ""), "")
                    state.append(getattr(event, "delta", ""))
                elif event_type == "response.function_call_arguments.done":
                    state = registry.get(getattr(event, "item_id", ""), "")
                    if state is None:
                        state = ToolCallState(item_id=getattr(event, "item_id", ""))
                        registry.put(state, getattr(event, "item_id", ""), "")
                    state.set_full(getattr(event, "arguments", ""))
                    if not state.call_id:
                        state.call_id = getattr(event, "item_id", "")
                    registry.put(state, state.item_id, state.call_id)
                    if not state.name:
                        state.name = "tool"
                    state.ensure_arguments(self._decode_arguments)
                    if not state.emitted:
                        emit(
                            ChatEvent(
                                kind="tool_call",
                                data={
                                    "name": state.name,
                                    "summary": self._summarize_arguments(state.arguments),
                                },
                            )
                        )
                        state.emitted = True
                    state.dispatched = True
                    tool_calls.append(
                        ToolInvocation(
                            name=state.name,
                            call_id=state.call_id,
                            arguments=dict(state.arguments),
                        )
                    )
                elif event_type == "response.completed":
                    final_response = getattr(event, "response", None)
                elif event_type == "response.failed":
                    raise RuntimeError("response failed") from None
                elif event_type == "error":
                    code = getattr(event, "code", "unknown")
                    message = getattr(event, "message", "Unknown error")
                    raise RuntimeError(f"api error ({code}): {message}") from None

            if final_response is None:
                final_response = stream.get_final_response()

        existing_by_call_id: Dict[str, int] = {}
        for idx, call in enumerate(tool_calls):
            if call.call_id:
                existing_by_call_id[call.call_id] = idx

        for state in registry.states():
            if state.dispatched or not state.call_id:
                continue
            state.ensure_arguments(self._decode_arguments)
            if not state.name:
                state.name = "tool"
            if not state.emitted:
                emit(
                    ChatEvent(
                        kind="tool_call",
                        data={
                            "name": state.name,
                            "summary": self._summarize_arguments(state.arguments),
                        },
                    )
                )
                state.emitted = True
            tool_calls.append(
                ToolInvocation(
                    name=state.name,
                    call_id=state.call_id,
                    arguments=dict(state.arguments),
                )
            )
            existing_by_call_id[state.call_id] = len(tool_calls) - 1

        if final_response is not None:
            for item in getattr(final_response, "output", []) or []:
                if getattr(item, "type", "") != "function_call":
                    continue
                call_id = getattr(item, "call_id", "") or ""
                args = self._decode_arguments(getattr(item, "arguments", ""))
                name = getattr(item, "name", "") or "tool"
                if call_id and call_id in existing_by_call_id:
                    stored = tool_calls[existing_by_call_id[call_id]]
                    if not stored.name and name:
                        stored.name = name
                    if not stored.arguments and args:
                        stored.arguments = args
                    continue
                emit(
                    ChatEvent(
                        kind="tool_call",
                        data={
                            "name": name,
                            "summary": self._summarize_arguments(args),
                        },
                    )
                )
                tool_calls.append(
                    ToolInvocation(
                        name=name,
                        call_id=call_id,
                        arguments=args,
                    )
                )
                if call_id:
                    existing_by_call_id[call_id] = len(tool_calls) - 1

        if final_response is None:
            raise RuntimeError("missing response from stream")

        return final_response, tool_calls

    def _extract_reasoning_summary(self, response: Any) -> str:
        """Aggregate reasoning summary text from the response payload."""
        try:
            items = getattr(response, "output", []) or []
        except Exception:
            return ""

        parts: List[str] = []
        for item in items:
            if getattr(item, "type", "") != "reasoning":
                continue

            summaries = getattr(item, "summary", []) or []
            for summary in summaries:
                text = getattr(summary, "text", "")
                if isinstance(text, str):
                    cleaned = text.strip()
                    if cleaned:
                        parts.append(cleaned)

        if not parts:
            return ""

        combined = " ".join(parts)
        return self._trim_text(combined)

    def _execute_tool(self, tool_name: str, tool_input: Dict[str, Any]) -> str:
        """Dispatch a tool call and return its result as plain text."""
        logging.info(f"Executing tool: {tool_name} with input: {tool_input}")
        try:
            if tool_name == "read_file":
                return read_file(
                    tool_input["path"],
                    tool_input.get("offset"),
                    tool_input.get("length"),
                )
            elif tool_name == "list_files":
                return list_files(tool_input.get("path", "."))
            elif tool_name == "edit_file":
                return edit_file(
                    tool_input["path"],
                    tool_input.get("old_text", ""),
                    tool_input["new_text"],
                )
            else:
                return f"Unknown tool: {tool_name}"
        except Exception as e:
            logging.error(f"Error executing {tool_name}: {str(e)}")
            return f"Error executing {tool_name}: {str(e)}"

    def chat(
        self,
        user_input: str,
        event_handler: Optional[Callable[[ChatEvent], None]] = None,
    ) -> Dict[str, Any]:
        """Send user input to the model and stream intermediate events."""
        logging.info(f"User input: {user_input}")
        conversation_input: List[Dict[str, Any]] = [
            {"role": "user", "content": user_input}
        ]
        events: List[ChatEvent] = []

        def emit(event: ChatEvent) -> None:
            events.append(event)
            if event_handler:
                event_handler(event)

        previous_id = self.previous_response_id

        while True:
            try:
                response, tool_calls = self._stream_response(
                    conversation_input, previous_id, emit
                )
            except Exception as exc:
                logging.error("Error streaming response: %s", exc)
                return {"final": f"Error: {str(exc)}", "events": events}

            previous_id = response.id

            if not tool_calls:
                self.previous_response_id = response.id
                final_text = (response.output_text or "").strip()
                return {"final": final_text, "events": events}

            tool_inputs: List[Dict[str, Any]] = []
            for call in tool_calls:
                result = self._execute_tool(call.name, call.arguments)
                logging.info(
                    "Tool result: %s",
                    self._summarize_tool_result(result),
                )

                output_payload = {
                    "type": "function_call_output",
                    "call_id": call.call_id,
                    "output": json.dumps({"result": result}),
                }
                tool_inputs.append(output_payload)

                emit(
                    ChatEvent(
                        kind="tool_result",
                        data={
                            "name": call.name,
                            "summary": self._summarize_tool_result(result),
                        },
                    )
                )

            conversation_input = tool_inputs


def resolve_api_key(arg_value: Optional[str]) -> Optional[str]:
    value = (arg_value or "").strip()
    if value:
        return value
    env_value = os.environ.get("OPENAI_API_KEY", "").strip()
    return env_value or None


class TerminalRenderer:
    """Render streaming events to the terminal while preserving state."""

    def __init__(self) -> None:
        self.thinking_active = False
        self.assistant_active = False
        self.last_thinking_text = ""
        self.thinking_buffer = ""
        self.assistant_buffer = ""

    def handle(self, event: ChatEvent) -> None:
        if event.kind == "thinking":
            self._handle_thinking(event)
        elif event.kind == "assistant_delta":
            chunk = event.data.get("chunk", "")
            if not chunk:
                return
            reset = event.data.get("reset") == "true"
            self._handle_assistant(chunk, reset)
        elif event.kind == "tool_call":
            self.finish()
            name = event.data.get("name", "")
            summary = event.data.get("summary", "")
            print(f"  tool: {name} ({summary})")
        elif event.kind == "tool_result":
            self.finish()
            name = event.data.get("name", "")
            summary = event.data.get("summary", "")
            print(f"  tool result: {name} -> {summary}")

    def finish(self) -> None:
        self._clear_thinking()
        self._clear_assistant()

    def _handle_thinking(self, event: ChatEvent) -> None:
        text = str(event.data.get("text", ""))
        if not text or text == self.last_thinking_text:
            return
        sanitized = text.replace("\n", " ")
        if not self.thinking_active:
            self._prepare_thinking_line()
        if sanitized.startswith(self.thinking_buffer):
            delta = sanitized[len(self.thinking_buffer) :]
            if not delta:
                return
            self._append_raw(delta)
        else:
            self._write_line("  thinking: ", sanitized)
        self.thinking_buffer = sanitized
        self.thinking_active = True
        self.last_thinking_text = text

    def _handle_assistant(self, chunk: str, reset: bool) -> None:
        sanitized = chunk.replace("\n", "\n  ")
        if reset or not self.assistant_active:
            self._clear_assistant()
            self._prepare_assistant_line()
        if sanitized.startswith(self.assistant_buffer):
            delta = sanitized[len(self.assistant_buffer) :]
            if not delta:
                return
            self._append_raw(delta)
        else:
            self._write_line("  assistant: ", sanitized)
        self.assistant_buffer = sanitized

    def _prepare_thinking_line(self) -> None:
        self._clear_assistant()
        self._clear_line()
        self._append_raw("  thinking: ")
        self.thinking_active = True
        self.thinking_buffer = ""

    def _prepare_assistant_line(self) -> None:
        self._clear_thinking()
        self._clear_line()
        self._append_raw("  assistant: ")
        self.assistant_active = True
        self.assistant_buffer = ""

    def _clear_thinking(self) -> None:
        if self.thinking_active:
            self._clear_line()
            self.thinking_active = False
        self.last_thinking_text = ""
        self.thinking_buffer = ""

    def _clear_assistant(self) -> None:
        if self.assistant_active:
            self._clear_line()
            self.assistant_active = False
        self.assistant_buffer = ""

    def _write_line(self, prefix: str, text: str) -> None:
        self._clear_line()
        self._append_raw(prefix)
        self._append_raw(text)

    def _append_raw(self, text: str) -> None:
        sys.stdout.write(text)
        sys.stdout.flush()

    def _clear_line(self) -> None:
        sys.stdout.write("\r\033[2K")
        sys.stdout.flush()


def run_cli(agent: AIAgent) -> None:
    print("AI Code Assistant")
    print("================")
    print("A conversational AI agent that can read, list, and edit files.")
    print("Type 'exit' or 'quit' to end the conversation.")
    print()

    while True:
        renderer: Optional[TerminalRenderer] = None
        try:
            user_input = input("You: ").strip()

            if user_input.lower() in {"exit", "quit"}:
                print("Goodbye!")
                break

            if not user_input:
                continue

            print("\nAssistant:")
            renderer = TerminalRenderer()
            result = agent.chat(user_input, event_handler=renderer.handle)
            renderer.finish()

            final_text = result.get("final", "")
            if final_text:
                print(f"  reply: {final_text}")
            else:
                print("  reply: (no response)")

            print()
        except KeyboardInterrupt:
            print("\n\nGoodbye!")
            break
        except EOFError:
            print("\nGoodbye!")
            break
        except Exception as exc:
            if renderer is not None:
                renderer.finish()
            print(f"\nError: {str(exc)}")
            print()


def main() -> None:
    """Entry point for the interactive CLI."""
    parser = argparse.ArgumentParser(
        description="AI Code Assistant - A conversational AI agent with file editing capabilities"
    )
    parser.add_argument(
        "--api-key", help="OpenAI API key (or set OPENAI_API_KEY env var)"
    )
    args = parser.parse_args()

    api_key = resolve_api_key(args.api_key)
    if not api_key:
        print(
            "Error: Please provide an API key via --api-key or OPENAI_API_KEY environment variable"
        )
        sys.exit(1)

    agent = AIAgent(api_key)
    run_cli(agent)


if __name__ == "__main__":
    main()