diff --git a/backend/managers/search_manager.py b/backend/managers/search_manager.py index e69de29..96be7db 100644 --- a/backend/managers/search_manager.py +++ b/backend/managers/search_manager.py @@ -0,0 +1,122 @@ +"""Handles internet search requests and page fetching for use as AI chat context.""" + +import ipaddress +import socket +from urllib.parse import urlparse + +import requests +from bs4 import BeautifulSoup +from ddgs import DDGS + +# Maximum characters extracted from a fetched page before truncating. +MAX_PAGE_CHARS = 3000 + +_HEADERS = {"User-Agent": "Mozilla/5.0 (compatible; AICodeEditor/1.0)"} + + +class SearchManager: + """Performs DuckDuckGo searches and fetches web pages for AI context injection. + + All outbound requests are validated against an SSRF blocklist so that + localhost and private network addresses can never be reached. + """ + + def perform_search(self, query: str, max_results: int = 5) -> list[dict]: + """Execute a DuckDuckGo text search and return normalised results. + + Args: + query: The search query string. + max_results: Maximum number of results to return. + + Returns: + List of {"title": str, "url": str, "snippet": str} dicts, + or an empty list if the search fails. + """ + try: + with DDGS() as ddgs: + raw = list(ddgs.text(query, max_results=max_results)) + return self.parse_results(raw) + except Exception: + return [] + + def parse_results(self, raw_results: list[dict]) -> list[dict]: + """Normalise raw DDGS result dicts to a consistent {"title", "url", "snippet"} shape. + + Args: + raw_results: List of raw dicts returned by ddgs.text(). + + Returns: + Normalised list of result dicts. + """ + results = [] + for r in raw_results: + results.append({ + "title": r.get("title", ""), + "url": r.get("href", r.get("url", "")), + "snippet": r.get("body", r.get("snippet", "")), + }) + return results + + def fetch_page(self, url: str) -> str: + """Fetch a web page and return its plain text content, truncated to MAX_PAGE_CHARS. + + Args: + url: The URL to fetch. + + Returns: + Plain text extracted from the page, or an error message string. + + Raises: + ValueError: if the URL fails the SSRF safety check. + """ + self._validate_url(url) + try: + response = requests.get(url, timeout=10, headers=_HEADERS) + response.raise_for_status() + + soup = BeautifulSoup(response.text, "html.parser") + + # Remove non-content elements before extracting text. + for tag in soup(["script", "style", "nav", "footer"]): + tag.decompose() + + text = soup.get_text(separator="\n", strip=True) + + if len(text) > MAX_PAGE_CHARS: + text = text[:MAX_PAGE_CHARS] + "\n... [truncated]" + + return text + + except ValueError: + raise + except Exception as e: + return f"Error fetching page: {e}" + + def _validate_url(self, url: str) -> None: + """Block localhost, private IPs, and non-http(s) schemes to prevent SSRF attacks. + + Args: + url: The URL to validate. + + Raises: + ValueError: if the URL is considered unsafe. + """ + parsed = urlparse(url) + + if parsed.scheme not in ("http", "https"): + raise ValueError(f"Blocked: only http/https allowed, got '{parsed.scheme}'") + + hostname = parsed.hostname or "" + + if hostname.lower() in ("localhost", "127.0.0.1", "::1"): + raise ValueError("Blocked: localhost access denied") + + try: + ip = ipaddress.ip_address(socket.gethostbyname(hostname)) + if ip.is_private or ip.is_loopback or ip.is_link_local: + raise ValueError(f"Blocked: private/loopback IP denied ({ip})") + except (socket.gaierror, ValueError) as e: + # Re-raise our own ValueError; ignore DNS resolution failures + # (let requests handle unknown hostnames naturally). + if isinstance(e, ValueError): + raise diff --git a/backend/managers/system_prompter.py b/backend/managers/system_prompter.py index d9610b7..e3c01e6 100644 --- a/backend/managers/system_prompter.py +++ b/backend/managers/system_prompter.py @@ -12,12 +12,17 @@ class SystemPrompter: """ @staticmethod - def generate_prompt(file_context: dict | None = None) -> str: - """Build a system prompt, optionally embedding a file's content. + def generate_prompt( + file_context: dict | None = None, + search_context: list[dict] | None = None, + ) -> str: + """Build a system prompt, optionally embedding a file and/or web search results. Args: - file_context: dict with keys 'name' (filename) and 'content' (raw text), - or None if no file should be included. + file_context: dict with keys 'name' (filename) and 'content' (raw text), + or None if no file should be included. + search_context: list of {"title", "url", "snippet"} dicts from SearchManager, + or None if no search results should be included. Returns: A ready-to-use system prompt string. @@ -28,6 +33,8 @@ class SystemPrompter: "Be concise and precise. Use markdown and fenced code blocks where appropriate." ) + prompt = base + if file_context: name = file_context.get("name", "unknown") content = file_context.get("content", "") @@ -36,13 +43,23 @@ class SystemPrompter: if len(content) > MAX_FILE_CHARS: content = content[:MAX_FILE_CHARS] + "\n... [truncated]" - file_section = ( + prompt += ( f"\n\nThe user currently has the following file open in the editor:\n" f"\n" f"\n{content}\n\n" f"\n" f"Refer to this file when answering questions about the code." ) - return base + file_section - return base + if search_context: + search_section = "\n\nThe user has performed a web search. Use the results below as additional context if relevant:\n\n" + for i, r in enumerate(search_context, 1): + search_section += ( + f"[{i}] {r.get('title', '')}\n" + f"URL: {r.get('url', '')}\n" + f"{r.get('snippet', '')}\n\n" + ) + search_section += "" + prompt += search_section + + return prompt diff --git a/frontend/chat.py b/frontend/chat.py index 43ab8bc..9c310fd 100644 --- a/frontend/chat.py +++ b/frontend/chat.py @@ -3,6 +3,7 @@ import streamlit as st from backend.managers.chat_manager import ChatManager from backend.managers.system_prompter import SystemPrompter +from backend.managers.search_manager import SearchManager import asyncio @@ -212,13 +213,59 @@ def render_agent_mode(): # ── Normal Chat ─────────────────────────────────────────────────────────────── +def _render_search_panel(): + """Render the collapsible web search panel above the chat history. + + Stores results in session_state.search_results so they are automatically + injected as context into the next message the user sends. + """ + search_results = st.session_state.get("search_results", []) + label = f"🔍 Web Search ({len(search_results)} result{'s' if len(search_results) != 1 else ''} active)" if search_results else "🔍 Web Search" + + with st.expander(label, expanded=False): + col_input, col_btn = st.columns([5, 1]) + with col_input: + query = st.text_input( + "Search query", + key="search_query_input", + placeholder="e.g. Python asyncio best practices", + label_visibility="collapsed", + ) + with col_btn: + search_clicked = st.button("Search", use_container_width=True) + + if search_clicked and query.strip(): + with st.spinner("Searching..."): + sm = SearchManager() + results = sm.perform_search(query.strip()) + if results: + st.session_state.search_results = results + st.rerun() + else: + st.warning("No results found.") + + # Display active results with a clear button. + if search_results: + st.caption("Results will be injected as context into your next message.") + for r in search_results: + st.markdown(f"**{r['title']}** \n{r['snippet']} \n[{r['url']}]({r['url']})") + st.divider() + if st.button("Clear search results", use_container_width=True): + st.session_state.search_results = [] + st.rerun() + + def render_normal_chat(): """Render the standard multi-turn chat interface. - On the first message the system prompt is injected into the history. + On the first message the system prompt is injected into the history, + including any active search results as context. Each subsequent message appends to the same conversation so the AI retains - full context throughout the session. + full context throughout the session. If search results are active when the + user sends a message, they are prepended to that message as a context block. """ + _render_search_panel() + # Replay the conversation history as chat bubbles (skip system messages). for message in st.session_state.chat_history: role = message["role"] @@ -227,10 +274,54 @@ def render_normal_chat(): with st.chat_message(role): st.markdown(message["content"]) - # Chat input — Enter to send, no extra button needed - user_input = st.chat_input("Type your message here...") + # Chat input — Enter to send, no extra button needed. + # Supports /search and /search clear as special commands. + user_input = st.chat_input("Type a message or /search ...") if user_input: + stripped = user_input.strip() + + # ── /search command ─────────────────────────────────────────────────── + if stripped.lower().startswith("/search"): + arg = stripped[len("/search"):].strip() + + with st.chat_message("user"): + st.markdown(stripped) + + if arg.lower() == "clear" or arg == "": + # /search clear (or bare /search) — remove active results. + st.session_state.search_results = [] + with st.chat_message("assistant"): + st.markdown("Search context cleared.") + st.session_state.chat_history.append({"role": "user", "content": stripped}) + st.session_state.chat_history.append({"role": "assistant", "content": "Search context cleared."}) + else: + # /search — run search and store results in context. + with st.chat_message("assistant"): + with st.spinner(f'Searching for "{arg}"...'): + sm = SearchManager() + results = sm.perform_search(arg) + + if results: + st.session_state.search_results = results + summary = f"Found {len(results)} result(s) for **{arg}**. They are now in context for this chat session.\n\n" + for i, r in enumerate(results, 1): + summary += f"**{i}. [{r['title']}]({r['url']})** \n{r['snippet']}\n\n" + st.markdown(summary) + response_text = summary + else: + msg = f'No results found for "{arg}".' + st.warning(msg) + response_text = msg + + st.session_state.chat_history.append({"role": "user", "content": stripped}) + st.session_state.chat_history.append({"role": "assistant", "content": response_text}) + + st.rerun() + return + + # ── Normal chat message ─────────────────────────────────────────────── chat_manager = st.session_state.chat_manager + search_results = st.session_state.get("search_results", []) # On the very first user message, prepend the system prompt so the AI # knows it is a code assistant embedded in an editor. @@ -238,7 +329,22 @@ def render_normal_chat(): system_prompt = SystemPrompter.generate_prompt() chat_manager.add_message("system", system_prompt) - # Show user message immediately without waiting for response. + # If search results are active, prepend them as a context block so the + # AI can reference them regardless of where in the conversation we are. + if search_results: + context_block = "\n" + for r in search_results: + context_block += ( + f"Title: {r['title']}\n" + f"URL: {r['url']}\n" + f"Snippet: {r['snippet']}\n\n" + ) + context_block += "\n\n" + message_to_send = context_block + user_input + else: + message_to_send = user_input + + # Show the original user text in the UI (not the context-enriched version). with st.chat_message("user"): st.markdown(user_input) @@ -246,7 +352,7 @@ def render_normal_chat(): with st.chat_message("assistant"): with st.spinner("Thinking..."): try: - ai_response = chat_manager.send_message(user_input) + ai_response = chat_manager.send_message(message_to_send) except Exception as e: ai_response = f"Error: {e}" st.markdown(ai_response) diff --git a/frontend/state.py b/frontend/state.py index 2a4b25d..aa05aa3 100644 --- a/frontend/state.py +++ b/frontend/state.py @@ -88,6 +88,11 @@ def init_state(): if "agent_pending_action" not in st.session_state: st.session_state.agent_pending_action = None + # Web search results to be injected as context into the next AI message. + # List of {"title": str, "url": str, "snippet": str} dicts, or empty list. + if "search_results" not in st.session_state: + st.session_state.search_results = [] + if __name__ == "__main__":