app.py

import re
import tempfile
from pathlib import Path
from typing import Iterator

import streamlit as st
from agno.agent import RunEvent, Agent
import os
from dotenv import load_dotenv

load_dotenv()

# -- Working directory setup --
OUTPUT_DIR = Path("outputs")
OUTPUT_DIR.mkdir(exist_ok=True)

# -- Avatars (DiceBear API) --
USER_AVATAR = "https://api.dicebear.com/9.x/avataaars/png?seed=user-chat&size=128"
BOT_AVATAR = "https://api.dicebear.com/9.x/bottts/png?seed=agno-bot&size=128"

# -- Allowed file types --
ALLOWED_TYPES = ["pdf", "docx", "pptx", "png", "jpg", "jpeg", "gif", "webp", "xlsx"]
ALLOWED_MIME = [
    "application/pdf",
    "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    "application/vnd.openxmlformats-officedocument.presentationml.presentation",
    "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    "image/png",
    "image/jpeg",
    "image/gif",
    "image/webp",
]
IMAGE_UPLOAD_EXTS = {".png", ".jpg", ".jpeg", ".gif", ".webp"}

# -- Available models --
MODELS = {
    "GPT-5.1": ("openai", "gpt-5.1"),
    "GPT-5.2": ("openai", "gpt-5.2"),
    "GPT-5.3": ("openai", "gpt-5.3"),
    "Claude Sonnet 4.6": ("anthropic", "claude-sonnet-4-6"),
    "Gemini 3.1 Pro": ("google", "gemini-3.1-pro-preview"),
}


@st.cache_resource
def create_agent(model_provider: str, model_id: str):
    """Create agent with specified model."""
    from agno.agent import Agent
    from agno.skills import Skills, LocalSkills
    from agno.models.anthropic import Claude
    from agno.models.google import Gemini
    from agno.models.openai import OpenAIResponses
    from agno.tools.file import FileTools
    from agno.tools.shell import ShellTools
    
    # Select model based on provider
    if model_provider == "openai":
        model = OpenAIResponses(id=model_id)
    elif model_provider == "anthropic":
        model = Claude(id=model_id)
    elif model_provider == "google":
        model = Gemini(id=model_id)
    else:
        model = OpenAIResponses(id="gpt-5.2")  # default
    
    return Agent(
        model=model,
        skills=Skills(loaders=[LocalSkills("agent_skills")]),
        tools=[FileTools(), ShellTools()],
        instructions=[
            "You have access to specialized skills. When a user request matches a skill, ALWAYS call get_skill_instructions() first to load the full instructions, then follow them exactly.",
            "Use the docx skill for Word documents (.docx files) - creating, reading, editing, or manipulating.",
            "Use the pptx skill for PowerPoint presentations (.pptx files) - creating, reading, editing, or analyzing.",
            "Use the pdf skill for PDF files - reading, creating, merging, splitting, filling forms, or manipulating.",
            "Use the xlsx skill for Excel files - reading, creating, calculating formulae, data visualization.",
            "To create files use save_file(). To run commands use run_shell_command(). To run skill scripts use get_skill_script(). These are your ONLY tools for file I/O and execution — do NOT call tools that do not exist.",
            "Save all output files to the 'outputs/' directory. Save temporary/intermediate scripts to 'outputs/.tmp/'. The outputs folder is created automatically.",
            "For data visualizations: use plotly for interactive charts (save as .html), matplotlib for static charts (save as .png). Both libraries are installed.",
            "IMPORTANT: For charts in Excel, use matplotlib to generate PNG images (dpi=300) and embed them with openpyxl.drawing.image.Image. Do NOT use openpyxl's native Chart API - it produces low-quality charts.",
            "When using get_skill_script(), ALL file path args MUST be absolute paths (e.g. /home/user/project/file.docx). The scripts run in the skill directory, not your working directory. Use run_shell_command('pwd') to get the current directory first.",
        ],
        markdown=True,
    )

def stream_response(agent, message: str) -> Iterator[str]:
    """Yield content chunks from a streaming agent run."""
    stream = agent.run(message, stream=True)
    for chunk in stream:
        if chunk.event == RunEvent.run_content and chunk.content:
            yield chunk.content


FILE_PATH_RE = re.compile(r"(?:/[\w.\-]+)+\.(?:pdf|docx|pptx|xlsx|png|jpg|jpeg|svg|html|txt)")
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".svg"}
VIZ_EXTS = {".html"}  # Plotly visualizations


def render_output_files(text: str):
    """Scan response text for file paths and render download buttons / previews."""
    paths = [Path(m) for m in FILE_PATH_RE.findall(text) if Path(m).is_file()]
    for p in paths:
        if p.suffix.lower() in IMAGE_EXTS:
            st.image(str(p), caption=p.name)
        elif p.suffix.lower() in VIZ_EXTS:
            # Render interactive Plotly visualizations
            with open(p, "r", encoding="utf-8") as f:
                html_content = f.read()
            st.components.v1.html(html_content, height=600, scrolling=True)
            st.download_button(
                label=f"Download {p.name}",
                data=p.read_bytes(),
                file_name=p.name,
                mime="text/html",
                key=f"dl-{p.name}-{hash(str(p))}",
            )
        else:
            st.download_button(
                label=f"📥 Download {p.name}",
                data=p.read_bytes(),
                file_name=p.name,
                mime=get_mime_type(p),
                key=f"dl-{p.name}-{hash(str(p))}",
            )


def save_uploaded_file(uploaded_file) -> Path:
    """Persist an uploaded file to a temp directory and return its path."""
    tmp_dir = Path(tempfile.mkdtemp())
    dest = tmp_dir / uploaded_file.name
    dest.write_bytes(uploaded_file.getvalue())
    return dest


def get_mime_type(file_path: Path) -> str:
    """Get MIME type for a file."""
    mime_types = {
        ".pdf": "application/pdf",
        ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
        ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        ".txt": "text/plain",
        ".html": "text/html",
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".svg": "image/svg+xml",
    }
    return mime_types.get(file_path.suffix.lower(), "application/octet-stream")


def list_output_files() -> list[Path]:
    """List all files in the outputs directory, excluding temp files."""
    if not OUTPUT_DIR.exists():
        return []
    files = []
    for f in OUTPUT_DIR.rglob('*'):
        if f.is_file() and '.tmp' not in f.parts:
            files.append(f)
    return sorted(files, key=lambda x: x.stat().st_mtime, reverse=True)


def cleanup_output_files() -> int:
    """Clean up output files in the OUTPUT_DIR. Returns count of deleted files."""
    if not OUTPUT_DIR.exists():
        return 0
    deleted = 0
    for item in OUTPUT_DIR.rglob("*"):
        if item.is_file():
            try:
                item.unlink()
                deleted += 1
            except Exception:
                pass
    return deleted

def cleanup_temp_files() -> int:
    """Clean up temporary files in .tmp directory. Returns count of deleted files."""
    tmp_dir = OUTPUT_DIR / ".tmp"
    if not tmp_dir.exists():
        return 0
    
    deleted_count = 0
    for item in tmp_dir.rglob('*'):
        if item.is_file():
            try:
                item.unlink()
                deleted_count += 1
            except Exception:
                pass
    
    # Remove empty directories
    for item in sorted(tmp_dir.rglob('*'), key=lambda x: len(x.parts), reverse=True):
        if item.is_dir() and not any(item.iterdir()):
            try:
                item.rmdir()
            except Exception:
                pass
    
    return deleted_count


# ── Page config ──────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Agno Skills Agent",
    page_icon="📄",
    layout="centered",
)

# ── Sidebar: file upload ─────────────────────────────────────────────────────
with st.sidebar:
    st.header("⚙️ Settings")
    
    # Model selection
    selected_model = st.selectbox(
        "Select Model",
        options=list(MODELS.keys()),
        index=1,  # Default to GPT-5.2
        help="Choose the AI model to use"
    )
    
    model_provider, model_id = MODELS[selected_model]
    st.caption(f"Provider: {model_provider.upper()}")
    
    st.divider()
    
    st.header("Upload files")
    uploaded = st.file_uploader(
        "Documents (Word/Excel/PPT/PDF) AND/OR Image (PNG, JPG, GIF, WEBP)",
        type=ALLOWED_TYPES,
        accept_multiple_files=True,
    )

    if uploaded:
        # Generate a list of file names for comparison
        uploaded_names = [f.name for f in uploaded]
        
        # Check if the uploaded files have changed
        if "uploaded_names" not in st.session_state or st.session_state.uploaded_names != uploaded_names:
            # Save all uploaded files
            paths = [save_uploaded_file(f) for f in uploaded]
            st.session_state.uploaded_paths = [str(p) for p in paths]
            st.session_state.uploaded_names = uploaded_names
        
        # Display success messages for each file
        for i, file in enumerate(uploaded):
            st.success(f"**{file.name}** ready")
            if i < len(st.session_state.uploaded_paths):
                st.caption(f"`{st.session_state.uploaded_paths[i]}`")

    st.divider()
    
    # ── Output files browser ──
    st.header("📥 Generated Files")
    output_files = list_output_files()
    
    if output_files:
        st.caption(f"{len(output_files)} file(s) in outputs/")
        for file_path in output_files[:10]:  # Show last 10 files
            relative_path = file_path.relative_to(OUTPUT_DIR)
            col1, col2 = st.columns([3, 1])
            with col1:
                st.caption(f"📄 {relative_path}")
            with col2:
                st.download_button(
                    label="⬇️",
                    data=file_path.read_bytes(),
                    file_name=file_path.name,
                    mime=get_mime_type(file_path),
                    key=f"sidebar-dl-{file_path.name}-{file_path.stat().st_mtime}",
                )
        if len(output_files) > 10:
            st.caption(f"... and {len(output_files) - 10} more files")
    else:
        st.caption("No files generated yet")
    
    st.divider()
    
    # Cleanup buttons
    col1, col2, col3 = st.columns(3)
    with col1:
        if st.button("🧹 Clean .tmp"):
            deleted = cleanup_temp_files()
            if deleted > 0:
                st.success(f"Deleted {deleted} temp file(s)")
            else:
                st.info("No temp files to clean")
            st.rerun()
    with col2:
        if st.button("🧹 Clean outputs"):
            deleted = cleanup_output_files()
            if deleted > 0:
                st.success(f"Deleted {deleted} file(s)")
            else:
                st.info("No files to clean")
            st.rerun()
    with col3:
        if st.button("Clear chat"):
            st.session_state.messages = []
            st.rerun()

# ── Session state ─────────────────────────────────────────────────────────────
if "messages" not in st.session_state:
    st.session_state.messages = []

# ── Title ─────────────────────────────────────────────────────────────────────
st.title("Agno Skills Agent")
st.caption(f"Using **{selected_model}** • Process PDF, DOCX, PPTX & XLSX files")

# ── Chat history ──────────────────────────────────────────────────────────────
for msg in st.session_state.messages:
    avatar = USER_AVATAR if msg["role"] == "user" else BOT_AVATAR
    with st.chat_message(msg["role"], avatar=avatar):
        st.markdown(msg["content"])
        if msg["role"] == "assistant":
            render_output_files(msg["content"])

# ── Chat input ────────────────────────────────────────────────────────────────
if prompt := st.chat_input("Ask something about your files…"):
    # Prepend file path context when files are uploaded
    full_prompt = prompt
    if "uploaded_paths" in st.session_state:
        file_context = []
        for uploaded_path in st.session_state.uploaded_paths:
            if Path(uploaded_path).suffix.lower() in IMAGE_UPLOAD_EXTS:
                file_context.append(f"[Image to embed: {uploaded_path}]")
            else:
                file_context.append(f"[File: {uploaded_path}]")
        full_prompt = "\n".join(file_context) + f"\n\n{prompt}"

    # Show user message
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user", avatar=USER_AVATAR):
        st.markdown(prompt)

    # Stream assistant response
    model_provider, model_id = MODELS[selected_model]
    agent = create_agent(model_provider, model_id)
    with st.chat_message("assistant", avatar=BOT_AVATAR):
        response = st.write_stream(stream_response(agent, full_prompt))
        render_output_files(response)

    st.session_state.messages.append({"role": "assistant", "content": response})