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, })