src/app.py

import os
import warnings
warnings.filterwarnings("ignore")

import streamlit as st

# ── Page config ────────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="IT Knowledge Base Assistant",
    page_icon="🤖",
    layout="wide",
)

st.title("🤖 IT Knowledge Base Assistant")
st.caption("Agentic RAG powered by smolagents · GPT-4o-mini · ChromaDB")

# ── Credentials ────────────────────────────────────────────────────────────────
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")

if not OPENAI_API_KEY:
    st.error(
        "**OPENAI_API_KEY** is not set. "
        "Add it as a Space secret (Settings → Variables and secrets)."
    )
    st.stop()

os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
os.environ["OPENAI_BASE_URL"] = OPENAI_BASE_URL

# ── Agent bootstrap (cached for the lifetime of the Space) ────────────────────
@st.cache_resource(show_spinner="Building knowledge base — this takes ~60 s on first start…")
def build_agent():
    from langchain.text_splitter import RecursiveCharacterTextSplitter
    from langchain_community.document_loaders import PyPDFDirectoryLoader
    from langchain_community.vectorstores import Chroma
    from langchain_openai import OpenAIEmbeddings
    from smolagents import CodeAgent, OpenAIServerModel, Tool

    content_dir = os.path.join(os.path.dirname(__file__), "content")
    loader = PyPDFDirectoryLoader(path=content_dir)
    splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
        encoding_name="cl100k_base",
        chunk_size=400,
        chunk_overlap=50,
    )
    chunks = loader.load_and_split(splitter)

    embeddings = OpenAIEmbeddings(model="text-embedding-ada-002")
    vectorstore = Chroma.from_documents(
        chunks, embeddings, collection_name="kb_articles"
    )

    class RetrieverTool(Tool):
        name = "retriever"
        description = (
            "Leverages semantic search to retrieve the most contextually relevant "
            "sections from the documentation based on a user query."
        )
        inputs = {
            "query": {
                "type": "string",
                "description": (
                    "The search query to match against documentation. "
                    "Phrase it as a natural language statement that reflects "
                    "the kind of information you're seeking."
                ),
            }
        }
        output_type = "string"

        def __init__(self, vector_db, **kwargs):
            super().__init__(**kwargs)
            self.vector_db = vector_db

        def forward(self, query: str) -> str:
            if not isinstance(query, str):
                raise ValueError("Query must be a string.")
            results = self.vector_db.similarity_search(query, k=10)
            formatted = "\n".join(
                f"\n\n===== Document {i} =====\n{doc.page_content}"
                for i, doc in enumerate(results)
            )
            return f"\nRetrieved documents:\n{formatted}"

    retriever_tool = RetrieverTool(vector_db=vectorstore)

    llm = OpenAIServerModel(
        model_id="gpt-4o-mini",
        api_base=OPENAI_BASE_URL,
        api_key=OPENAI_API_KEY,
    )
    agent = CodeAgent(
        tools=[retriever_tool],
        model=llm,
        max_steps=6,
        verbosity_level=0,
    )
    return agent


# ── Extract a serialisable snapshot from agent.logs ───────────────────────────
def snapshot_logs(logs: list) -> list:
    """Convert smolagents log objects into plain dicts for session-state storage."""
    steps = []
    for entry in logs:
        kind = type(entry).__name__
        if "Planning" in kind:
            steps.append({
                "type": "planning",
                "output": getattr(entry, "model_output", "") or "",
            })
        elif "Action" in kind:
            obs = getattr(entry, "observations", "") or ""
            steps.append({
                "type": "action",
                "step_number": getattr(entry, "step_number", len(steps)),
                "output": getattr(entry, "model_output", "") or "",
                "observations": obs,
                "error": str(getattr(entry, "error", "") or ""),
                "duration": getattr(entry, "duration", None),
            })
    return steps


# ── Render a single observation ───────────────────────────────────────────────
def render_observation(obs: str):
    """
    If the observation is a RetrieverTool result (contains '===== Document N ====='),
    display each chunk as a collapsed card.  Otherwise fall back to plain text.
    """
    if not obs:
        return

    if "===== Document" in obs:
        st.markdown("*📚 Retrieved documents:*")
        # Split on the document-header lines produced by RetrieverTool
        raw_parts = obs.split("===== Document ")
        docs = []
        for part in raw_parts[1:]:          # first element is the preamble
            header, _, body = part.partition("=====")
            docs.append((header.strip(), body.strip()))

        for doc_num, content in docs:
            with st.expander(f"Document {doc_num}", expanded=False):
                st.markdown(content)
    else:
        st.markdown("*Observation:*")
        st.text(obs)


# ── Render reasoning steps below an assistant message ─────────────────────────
def render_reasoning(steps: list):
    if not steps:
        return
    with st.expander("🧠 Agent reasoning", expanded=True):
        for step in steps:
            if step["type"] == "planning":
                st.markdown("**📋 Planning**")
                st.markdown(step["output"])
                st.divider()

            elif step["type"] == "action":
                dur = f" — {step['duration']:.1f}s" if step.get("duration") else ""
                st.markdown(f"**⚙️ Step {step['step_number']}{dur}**")

                if step["output"]:
                    st.markdown("*Generated code:*")
                    st.code(step["output"], language="python")

                if step["observations"]:
                    render_observation(step["observations"])

                if step["error"]:
                    st.error(step["error"])

                st.divider()


# ── Sidebar – example questions ───────────────────────────────────────────────
EXAMPLES = [
    "A new hire starts Monday. Walk me through every IT step before their first login.",
    "My laptop BSODs with DRIVER_IRQL_NOT_LESS_OR_EQUAL when connecting to VPN. What should I check?",
    "What is the exact password complexity policy — length, character types, reuse limit, expiry?",
    "Files on my desktop are renaming with weird extensions after a CrowdStrike alert. What do I do?",
    "An employee is terminated today at 3 pm. Give me a complete IT checklist in order.",
]

with st.sidebar:
    st.header("Example questions")
    for q in EXAMPLES:
        if st.button(q, use_container_width=True):
            st.session_state.pending_prompt = q

# ── Chat history ──────────────────────────────────────────────────────────────
if "messages" not in st.session_state:
    st.session_state.messages = []

for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])
        if msg["role"] == "assistant" and msg.get("steps"):
            render_reasoning(msg["steps"])

# ── Handle example-button clicks ──────────────────────────────────────────────
prompt = st.chat_input("Ask anything about IT policies and procedures…")
if "pending_prompt" in st.session_state:
    prompt = st.session_state.pop("pending_prompt")

# ── Run the agent ─────────────────────────────────────────────────────────────
if prompt:
    agent = build_agent()

    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        with st.spinner("Searching knowledge base…"):
            try:
                response = agent.run(prompt)
                # smolagents ≥1.10 stores steps in agent.memory.steps;
                # older builds used agent.logs — try both.
                raw = (
                    getattr(getattr(agent, "memory", None), "steps", None)
                    or getattr(agent, "logs", None)
                    or []
                )
                steps = snapshot_logs(raw)
            except Exception as e:
                response = f"Sorry, something went wrong: {e}"
                steps = []

        st.markdown(response)
        render_reasoning(steps)

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