nano_cowork.py

import streamlit as st
import os
import json
import inspect
import functools
import urllib.request
import urllib.error

# ==========================================
# 1. CORE: Nano-Skills Engine
# ==========================================
st.sidebar.title("Nano Cowork")
REGISTRY = {}


def skill(func):
    """Decorator: Converts a Python function into an Anthropic-style tool schema."""

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)

    sig = inspect.signature(func)
    doc = func.__doc__ or "No description."

    properties = {}
    required = []

    for name, param in sig.parameters.items():
        type_map = {
            int: "integer",
            float: "number",
            bool: "boolean",
            list: "array",
            dict: "object",
            str: "string",
        }
        param_type = type_map.get(param.annotation, "string")
        properties[name] = {"type": param_type, "description": f"Parameter: {name}"}
        if param.default == inspect.Parameter.empty:
            required.append(name)

    REGISTRY[func.__name__] = {
        "func": func,
        "schema": {
            "name": func.__name__,
            "description": doc.strip(),
            "input_schema": {
                "type": "object",
                "properties": properties,
                "required": required,
            },
        },
    }
    return wrapper


def get_tools_schema():
    return [item["schema"] for item in REGISTRY.values()]


def execute_skill(name, arguments):
    if name not in REGISTRY:
        return f"Error: Skill '{name}' not found."
    try:
        return str(REGISTRY[name]["func"](**arguments))
    except Exception as e:
        return f"Error executing '{name}': {str(e)}"


# ==========================================
# 2. OpenRouter API Client
# ==========================================

DEFAULT_MODEL = "anthropic/claude-opus-4.5"
API_URL = "https://openrouter.ai/api/v1/messages"


def call_openrouter(api_key, messages, model=None, system_prompt=None):
    """Call OpenRouter API using Anthropic message format."""
    body = {
        "model": model or DEFAULT_MODEL,
        "max_tokens": 8192,
        "messages": messages,
        "tools": get_tools_schema(),
    }
    if system_prompt:
        body["system"] = system_prompt

    request = urllib.request.Request(
        API_URL,
        data=json.dumps(body).encode("utf-8"),
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}",
            "anthropic-version": "2023-06-01",
        },
    )

    try:
        response = urllib.request.urlopen(request, timeout=120)
        return json.loads(response.read())
    except urllib.error.HTTPError as e:
        error_body = e.read().decode("utf-8", errors="ignore")
        raise Exception(f"API Error {e.code}: {error_body}")


# ==========================================
# 3. CONFIG & STATE
# ==========================================

st.set_page_config(layout="wide", page_title="NanoCowork")

if "messages" not in st.session_state:
    st.session_state.messages = []  # Anthropic format messages
if "preview_content" not in st.session_state:
    st.session_state.preview_content = None
if "preview_type" not in st.session_state:
    st.session_state.preview_type = "markdown"
if "preview_open" not in st.session_state:
    st.session_state.preview_open = False
if "preview_selected_file" not in st.session_state:
    st.session_state.preview_selected_file = None

# --- Sidebar ---
with st.sidebar:
    st.header("⚙️ Config")
    api_key_default = os.getenv("OPENROUTER_API_KEY", "")
    api_key_input = st.text_input(
        "OpenRouter API Key", type="password", value=api_key_default
    )
    model_input = st.text_input("Model", value=DEFAULT_MODEL)
    system_prompt_input = st.text_area(
        "System Prompt",
        value=f"Concise coding assistant. cwd: {os.getcwd()}",
        height=68,
    )

    st.divider()
    custom_workspace = st.text_input("Workspace Dir", value="./")
    WORKSPACE_DIR = custom_workspace
    if not os.path.exists(WORKSPACE_DIR):
        try:
            os.makedirs(WORKSPACE_DIR)
        except Exception:
            pass

    uploaded_file = st.file_uploader("Upload File", label_visibility="collapsed")
    if uploaded_file:
        with open(os.path.join(WORKSPACE_DIR, uploaded_file.name), "wb") as f:
            f.write(uploaded_file.getbuffer())
        st.toast(f"Uploaded {uploaded_file.name}")

    st.divider()
    if st.button("🔄 Refresh"):
        st.rerun()
    if st.button("🗑️ Clear Chat"):
        st.session_state.messages = []
        st.session_state.preview_content = None
        st.rerun()

    st.caption("Workspace Files:")
    try:
        for root, _, filenames in os.walk(WORKSPACE_DIR):
            for filename in filenames:
                f_path = os.path.join(root, filename)
                rel_path = os.path.relpath(f_path, WORKSPACE_DIR)
                if rel_path.startswith("."):
                    continue
                with st.expander(f"📄 {rel_path}"):
                    with open(f_path, "r", encoding="utf-8", errors="ignore") as f:
                        st.code(f.read()[:500])
    except Exception:
        pass


# ==========================================
# 4. SKILLS
# ==========================================


@skill
def list_files(subdirectory: str = "") -> str:
    """Lists files in workspace."""
    target_dir = os.path.join(WORKSPACE_DIR, subdirectory)
    if not os.path.exists(target_dir):
        return "Directory not found."
    files = []
    for root, _, filenames in os.walk(target_dir):
        for filename in filenames:
            files.append(os.path.relpath(os.path.join(root, filename), WORKSPACE_DIR))
    return json.dumps(files)


@skill
def read_file(filename: str) -> str:
    """Reads a file."""
    try:
        with open(os.path.join(WORKSPACE_DIR, filename), "r", encoding="utf-8") as f:
            return f.read()
    except Exception as e:
        return str(e)


@skill
def write_file(filename: str, content: str) -> str:
    """Writes to a file."""
    path = os.path.join(WORKSPACE_DIR, filename)
    os.makedirs(os.path.dirname(path) if os.path.dirname(path) else ".", exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        f.write(content)
    return f"Wrote {filename}"


@skill
def run_command(command: str) -> str:
    """Runs shell command."""
    import subprocess

    try:
        res = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            cwd=WORKSPACE_DIR,
            timeout=30,
        )
        return f"COMMAND:{command}\nSTDOUT:\n{res.stdout}\nSTDERR:\n{res.stderr}"
    except Exception as e:
        return str(e)


@skill
def preview_artifact(filename: str, format: str = "markdown") -> str:
    """
    Renders a file in the preview panel. Input filename should be relative to workspace.
    format: 'markdown' or 'html'.
    """
    with open(os.path.join(WORKSPACE_DIR, filename), "r", encoding="utf-8") as f:
        content = f.read()
    st.session_state.preview_content = content
    st.session_state.preview_type = format
    return "Displayed in preview panel."


# ==========================================
# 5. Message Format Helpers
# ==========================================


def extract_text_from_content(content):
    """Extract display text from Anthropic content blocks."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                parts.append(block["text"])
        return "\n".join(parts)
    return ""


def has_tool_use(content):
    """Check if content blocks contain tool_use."""
    if not isinstance(content, list):
        return False
    return any(b.get("type") == "tool_use" for b in content if isinstance(b, dict))


def get_tool_uses(content):
    """Extract tool_use blocks from content."""
    if not isinstance(content, list):
        return []
    return [b for b in content if isinstance(b, dict) and b.get("type") == "tool_use"]


# ==========================================
# 6. MAIN LAYOUT
# ==========================================

has_preview = st.session_state.preview_content is not None

if has_preview:
    chat_col, preview_col = st.columns([0.5, 0.5], gap="large")
else:
    chat_col = st.container()
    preview_col = None

# --- Preview Controls (Sidebar) ---
with st.sidebar:
    if st.button("Open Preview", key="open_prev"):
        st.session_state.preview_open = True

    if st.session_state.preview_open:
        files = []
        for root, _, filenames in os.walk(WORKSPACE_DIR):
            for f in filenames:
                files.append(os.path.relpath(os.path.join(root, f), WORKSPACE_DIR))
        files = [f for f in files if not f.startswith(".")]
        if files:
            st.session_state.preview_selected_file = st.selectbox(
                "Select file to preview",
                options=files,
                key="preview_file_selectbox",
            )
            if st.button("Load into Preview", key="load_prev"):
                selected_file = st.session_state.preview_selected_file
                filetype = os.path.splitext(selected_file)[1][1:].lower()
                st.session_state.preview_content = read_file(selected_file)
                st.session_state.preview_type = (
                    "html" if filetype == "html" else "markdown"
                )
                st.rerun()
        else:
            st.info("No files in workspace yet.")

# --- Render Preview ---
if has_preview and preview_col:
    with preview_col:
        subcol1, subcol2, subcol3 = st.columns([0.2, 0.5, 0.3])
        with subcol1:
            st.subheader("👁️ Preview")
        with subcol3:
            if st.button("Close Preview", key="close_prev"):
                st.session_state.preview_content = None
                st.session_state.preview_open = False
                st.rerun()
        with subcol2:
            artifact_height = st.slider("Preview Height", min_value=300, max_value=1200, value=950)

        with st.container(border=True, height=artifact_height):
            if st.session_state.preview_type == "html":
                st.html(st.session_state.preview_content)
            else:
                st.markdown(st.session_state.preview_content)

# --- Render Chat History ---
with chat_col:
    for msg in st.session_state.messages:
        role = msg.get("role", "")
        content = msg.get("content", "")

        if role == "user":
            # User messages: could be plain string or list with tool_result blocks
            if isinstance(content, str):
                with st.chat_message("user"):
                    st.markdown(content)
            # tool_result blocks are rendered below as assistant context
            elif isinstance(content, list):
                tool_results = [
                    b for b in content if isinstance(b, dict) and b.get("type") == "tool_result"
                ]
                if tool_results:
                    with st.chat_message("assistant"):
                        for tr in tool_results:
                            with st.expander(
                                f"Tool Result (id: {tr.get('tool_use_id', '?')[:12]}...)",
                                expanded=False,
                            ):
                                st.code(tr.get("content", "")[:500])

        elif role == "assistant":
            with st.chat_message("assistant"):
                text = extract_text_from_content(content)
                if text:
                    st.markdown(text)
                tool_uses = get_tool_uses(content)
                if tool_uses:
                    with st.status("🛠️ Used Tools", state="complete"):
                        for tu in tool_uses:
                            st.code(f"{tu['name']}({json.dumps(tu.get('input', {}))[:80]})")

# --- Chat Input ---
if prompt := st.chat_input("Ex: 'Create a visualization of sales data'"):
    if not api_key_input:
        st.error("OpenRouter API Key required.")
        st.stop()

    # Append user message (Anthropic format: string content)
    st.session_state.messages.append({"role": "user", "content": prompt})

    with chat_col:
        with st.chat_message("user"):
            st.markdown(prompt)

        # --- Agent Loop ---
        with st.chat_message("assistant"):
            step = 0
            needs_rerun = False
            while step < 10:
                step += 1
                try:
                    response = call_openrouter(
                        api_key=api_key_input,
                        messages=st.session_state.messages,
                        model=model_input,
                        system_prompt=system_prompt_input,
                    )

                    content_blocks = response.get("content", [])

                    # Store assistant message
                    st.session_state.messages.append(
                        {"role": "assistant", "content": content_blocks}
                    )

                    # Display text blocks
                    text = extract_text_from_content(content_blocks)
                    if text:
                        st.markdown(text)

                    # Handle tool calls
                    tool_uses = get_tool_uses(content_blocks)
                    if tool_uses:
                        tool_results = []
                        needs_rerun = False
                        with st.status(
                            f"⚙️ Step {step}: Working...", expanded=True
                        ) as status:
                            for tu in tool_uses:
                                name = tu["name"]
                                args = tu.get("input", {})
                                tool_id = tu["id"]
                                status.write(f"**Action:** `{name}`")

                                result = execute_skill(name, args)

                                status.write("**Result:**")
                                status.code(result[:500] + ("..." if len(result) > 500 else ""))

                                tool_results.append(
                                    {
                                        "type": "tool_result",
                                        "tool_use_id": tool_id,
                                        "content": result,
                                    }
                                )

                                if name == "preview_artifact":
                                    needs_rerun = True

                        # Append tool results as user message (Anthropic format)
                        st.session_state.messages.append(
                            {"role": "user", "content": tool_results}
                        )

                        # If this was the last step, do one final call for summary
                        if step >= 10:
                            break

                        continue
                    else:
                        # No tool calls — this is the final text response
                        break

                except Exception as e:
                    st.error(f"Error: {e}")
                    break

            # If the loop ended after tool calls without a final text response,
            # make one more API call to get a conclusion
            last_msg = st.session_state.messages[-1] if st.session_state.messages else {}
            if last_msg.get("role") == "user" and isinstance(last_msg.get("content"), list):
                try:
                    response = call_openrouter(
                        api_key=api_key_input,
                        messages=st.session_state.messages,
                        model=model_input,
                        system_prompt=system_prompt_input,
                    )
                    content_blocks = response.get("content", [])
                    st.session_state.messages.append(
                        {"role": "assistant", "content": content_blocks}
                    )
                    text = extract_text_from_content(content_blocks)
                    if text:
                        st.markdown(text)
                except Exception as e:
                    st.error(f"Error getting summary: {e}")

            # Rerun to refresh preview panel if needed
            if needs_rerun:
                st.rerun()