""" AI Security Tutor for VAPT Agent. Now powered by Nebius Token Factory (OpenAI-compatible API) with Chroma vector search. Key behaviours: - Uses a Nebius-hosted model for chat. - Uses a Nebius embedding model + Chroma to build a vector index over the VAPT Markdown report. - The index is built ONCE per report (per process) and automatically rebuilt if the report content changes. - Ensures that vectors from one report are never reused for another report by hashing the report content and recreating the index when it changes. - Optionally enriches answers with web search via Tavily. Environment variables: # Nebius (required) - NEBIUS_API_KEY : Nebius Token Factory API key - NEBIUS_BASE_URL : (optional) e.g. "https://api.tokenfactory.nebius.com/v1" - NEBIUS_TUTOR_MODEL : (optional) chat model id for tutor - NEBIUS_EMBEDDING_MODEL : embedding model id for vector search (REQUIRED to enable Chroma search) # Optional web search - TAVILY_API_KEY : enables web search if set """ import os import hashlib from typing import List, Tuple, Dict from dataclasses import dataclass import requests from openai import OpenAI from prompt import get_tutor_system_prompt # Try to import Chroma, but degrade gracefully if not installed try: import chromadb CHROMA_AVAILABLE = True except ImportError: chromadb = None CHROMA_AVAILABLE = False # --------------------------------------------------------------------------- # Simple helpers # --------------------------------------------------------------------------- def _normalize(text: str) -> str: return text.lower() def _extract_report_sections( report_md: str, max_section_chars: int = 2000 ) -> List[str]: """ Split the markdown report into logical sections based on '## ' headings. If sections are very large, they are further split into smaller chunks. """ if not report_md: return [] sections: List[str] = [] current: List[str] = [] lines = report_md.splitlines() for line in lines: if line.startswith("## "): if current: sections.append("\n".join(current).strip()) current = [] current.append(line) if current: sections.append("\n".join(current).strip()) if not sections: sections = [report_md] # Further split oversized sections into smaller chunks final_chunks: List[str] = [] for sec in sections: if len(sec) <= max_section_chars: final_chunks.append(sec) else: # naive split by paragraphs paras = sec.split("\n\n") chunk: List[str] = [] size = 0 for p in paras: p_len = len(p) + 2 if size + p_len > max_section_chars and chunk: final_chunks.append("\n\n".join(chunk)) chunk = [p] size = p_len else: chunk.append(p) size += p_len if chunk: final_chunks.append("\n\n".join(chunk)) return final_chunks def _web_search_tavily(query: str, max_results: int = 3) -> str: """ Optional: perform web search via Tavily Search API. Requires TAVILY_API_KEY in env. If not present or call fails, returns an empty string and the tutor will just rely on the report. """ api_key = os.getenv("TAVILY_API_KEY") if not api_key: return "" try: payload = { "api_key": api_key, "query": query, "max_results": max_results, "include_answer": True, "search_depth": "basic", } resp = requests.post( "https://api.tavily.com/search", json=payload, timeout=15, ) resp.raise_for_status() data = resp.json() parts: List[str] = [] answer = data.get("answer") if answer: parts.append(f"Direct answer: {answer}") results = data.get("results") or [] for r in results[:max_results]: title = r.get("title") or "Untitled" url = r.get("url") or "" content = r.get("content") or "" parts.append(f"- {title}\n {content[:300]}...\n Source: {url}") return "\n".join(parts) if parts else "" except Exception: # Fail silently – we don't want the tutor to break if search fails return "" # --------------------------------------------------------------------------- # Security Tutor with Nebius + Chroma # --------------------------------------------------------------------------- @dataclass class TutorConfig: base_url: str api_key: str model: str embedding_model: str | None class SecurityTutor: """AI-powered security education assistant backed by Nebius + Chroma.""" def __init__(self): """ Initialize the Security Tutor with Nebius OpenAI-compatible client. Required: - NEBIUS_API_KEY Optional: - NEBIUS_BASE_URL (defaults to Nebius Token Factory base URL) - NEBIUS_TUTOR_MODEL - NEBIUS_EMBEDDING_MODEL (required to enable Chroma vector search) """ api_key = os.getenv("NEBIUS_API_KEY") if not api_key: self.client = None self.available = False self.config = None self.chroma_client = None self.vector_enabled = False self._raw_report = "" self._report_hash = None self._collection = None return base_url = os.getenv( "NEBIUS_BASE_URL", "https://api.tokenfactory.nebius.com/v1", # Nebius OpenAI-compatible base URL ) model = os.getenv( "NEBIUS_TUTOR_MODEL", "meta-llama/Meta-Llama-3.1-70B-Instruct", # example; override with your chosen model ) embedding_model = os.getenv( "NEBIUS_EMBEDDING_MODEL", # REQUIRED for vector search None, ) self.config = TutorConfig( base_url=base_url, api_key=api_key, model=model, embedding_model=embedding_model, ) self.client = OpenAI(base_url=self.config.base_url, api_key=self.config.api_key) self.available = True # Chroma setup if CHROMA_AVAILABLE and self.config.embedding_model: # Ephemeral in-memory store is fine for a single process self.chroma_client = chromadb.EphemeralClient() self.vector_enabled = True else: self.chroma_client = None self.vector_enabled = False # Per-report state self._raw_report: str = "" self._report_hash: str | None = None self._collection = None # Chroma collection holding current report vectors # ------------------------------------------------------------------ # # Public entry point # ------------------------------------------------------------------ # def chat( self, message: str, report_context: str, history: List[Tuple[str, str]], ) -> str: """ Handle a chat message from the user. Args: message: User's question report_context: Full VAPT report markdown for THIS user/run history: Previous chat messages [(user_msg, assistant_msg), ...] Behaviour: - If the report content has changed since last call, rebuild the vector index just once and store it. - Then run vector search on the stored index for this question. - Never reuse vectors from a previous report for the new report. """ if not self.available or not self.client: return ( "🔧 AI Tutor is not configured yet.\n\n" "Please set NEBIUS_API_KEY (and optionally NEBIUS_TUTOR_MODEL) " "in your environment to enable the tutor." ) # 1) Ensure the index is up-to-date for THIS report self._ensure_report_index(report_context) # 2) Retrieve relevant snippets from the currently indexed report report_snippets = self._search_report_with_vectors(message) # 3) Optional web search (if Tavily API key is configured) web_snippets = _web_search_tavily(message) web_note = ( "\n\n---\n\nWeb search snippets:\n" + web_snippets if web_snippets else "\n\n(Web search not configured or returned no results.)" ) # 4) Build system prompt with clear grounding instructions system_prompt = self._build_system_prompt( report_snippets=report_snippets, include_web=bool(web_snippets), ) # 5) Build messages (system + history + current question) messages: List[Dict[str, str]] = [{"role": "system", "content": system_prompt}] for user_msg, assistant_msg in history: if user_msg: messages.append({"role": "user", "content": user_msg}) if assistant_msg: messages.append({"role": "assistant", "content": assistant_msg}) user_content = ( f"{message}\n\n" "-----\n\n" "Use the following VAPT report excerpts as your primary source of truth:\n\n" f"{report_snippets}\n" f"{web_note}" ) messages.append({"role": "user", "content": user_content}) try: completion = self.client.chat.completions.create( model=self.config.model, messages=messages, temperature=0.4, max_tokens=800, ) return completion.choices[0].message.content except Exception as e: return f"❌ Error communicating with Security Tutor (Nebius): {str(e)}" # ------------------------------------------------------------------ # # Report index management (build once per report) # ------------------------------------------------------------------ # def _ensure_report_index(self, report_md: str) -> None: """ Ensure that the vector index reflects the given report. - If no report or same report as before -> do nothing. - If a different report -> rebuild vectors so we NEVER reuse old vectors for a new report. """ report_md = report_md or "" # If we never had a report and still don't, nothing to do if not report_md and not self._raw_report: return new_hash = hashlib.sha256(report_md.encode("utf-8")).hexdigest() # If hash matches, it's the same report -> keep existing vectors if self._report_hash == new_hash: return # Report changed: update internal state and rebuild index self._raw_report = report_md self._report_hash = new_hash if ( not self.vector_enabled or not self.chroma_client or not self.config.embedding_model ): # We still store the report, so fallback search can use it self._collection = None return # Build a fresh collection just for this report sections = _extract_report_sections(self._raw_report) if not sections: self._collection = None return # Collection name derived from hash to avoid mixing coll_name = f"vapt_report_{self._report_hash[:8]}" # Create or get collection; then clear any previous contents self._collection = self.chroma_client.get_or_create_collection(name=coll_name) try: self._collection.delete(where={}) except Exception: # Older Chroma versions might not like empty filters; safely ignore pass # Embed and add sections ids = [f"chunk-{i}" for i in range(len(sections))] embeddings = self._embed_texts(sections) self._collection.add(ids=ids, documents=sections, embeddings=embeddings) # ------------------------------------------------------------------ # # Vector search over report using Chroma + Nebius embeddings # ------------------------------------------------------------------ # def _embed_texts(self, texts: List[str]) -> List[List[float]]: """ Create embeddings for a list of texts using Nebius embedding model. """ if not self.config.embedding_model: raise RuntimeError( "NEBIUS_EMBEDDING_MODEL is not set; cannot perform vector search." ) resp = self.client.embeddings.create( model=self.config.embedding_model, input=texts, ) return [item.embedding for item in resp.data] def _search_report_with_vectors(self, question: str, top_k: int = 4) -> str: """ Use the current Chroma collection (for the current report) to retrieve the most relevant chunks. Falls back to simple truncation of the report if vector search is not available. """ if not self._raw_report: return "No VAPT report is currently available." # If vector search is not enabled or we have no collection, fallback if not self.vector_enabled or not self._collection: # Cheap fallback: Executive Summary + Key Findings, or first 2000 chars sections = _extract_report_sections(self._raw_report) fallback = [ s for s in sections if "executive summary" in _normalize(s) or "key findings" in _normalize(s) ] if fallback: return "\n\n---\n\n".join(fallback)[:2000] return self._raw_report[:2000] try: # Embed the question and query the collection q_embedding = self._embed_texts([question])[0] results = self._collection.query( query_embeddings=[q_embedding], n_results=top_k, ) docs = results.get("documents", [[]])[0] if results else [] if not docs: return self._raw_report[:2000] joined = "\n\n---\n\n".join(docs) return joined[:2000] except Exception: # Any failure -> fall back to raw report return self._raw_report[:2000] # ------------------------------------------------------------------ # # System prompt builder # ------------------------------------------------------------------ # def _build_system_prompt(self, report_snippets: str, include_web: bool) -> str: """ Build the system prompt for the AI tutor. Args: report_snippets: Text retrieved from the VAPT report include_web: Whether web search snippets are available Returns: System prompt string """ return get_tutor_system_prompt(report_snippets, include_web) # Global tutor instance (shared within the process) _tutor_instance: SecurityTutor | None = None def get_tutor() -> SecurityTutor: """ Get or create the global SecurityTutor instance. Note: Vectors are tied to the report markdown passed into `chat()`. Whenever a new report is used, the tutor automatically rebuilds its index so that vectors from a previous report are never reused. """ global _tutor_instance if _tutor_instance is None: _tutor_instance = SecurityTutor() return _tutor_instance