start web search in chat
This commit is contained in:
parent
f91707b5b5
commit
bdbdbc9b3a
@ -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
|
||||||
@ -12,12 +12,17 @@ class SystemPrompter:
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def generate_prompt(file_context: dict | None = None) -> str:
|
def generate_prompt(
|
||||||
"""Build a system prompt, optionally embedding a file's content.
|
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:
|
Args:
|
||||||
file_context: dict with keys 'name' (filename) and 'content' (raw text),
|
file_context: dict with keys 'name' (filename) and 'content' (raw text),
|
||||||
or None if no file should be included.
|
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:
|
Returns:
|
||||||
A ready-to-use system prompt string.
|
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."
|
"Be concise and precise. Use markdown and fenced code blocks where appropriate."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
prompt = base
|
||||||
|
|
||||||
if file_context:
|
if file_context:
|
||||||
name = file_context.get("name", "unknown")
|
name = file_context.get("name", "unknown")
|
||||||
content = file_context.get("content", "")
|
content = file_context.get("content", "")
|
||||||
@ -36,13 +43,23 @@ class SystemPrompter:
|
|||||||
if len(content) > MAX_FILE_CHARS:
|
if len(content) > MAX_FILE_CHARS:
|
||||||
content = content[:MAX_FILE_CHARS] + "\n... [truncated]"
|
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\nThe user currently has the following file open in the editor:\n"
|
||||||
f"<file name=\"{name}\">\n"
|
f"<file name=\"{name}\">\n"
|
||||||
f"<code>\n{content}\n</code>\n"
|
f"<code>\n{content}\n</code>\n"
|
||||||
f"</file>\n"
|
f"</file>\n"
|
||||||
f"Refer to this file when answering questions about the code."
|
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<search_results>\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 += "</search_results>"
|
||||||
|
prompt += search_section
|
||||||
|
|
||||||
|
return prompt
|
||||||
|
|||||||
118
frontend/chat.py
118
frontend/chat.py
@ -3,6 +3,7 @@
|
|||||||
import streamlit as st
|
import streamlit as st
|
||||||
from backend.managers.chat_manager import ChatManager
|
from backend.managers.chat_manager import ChatManager
|
||||||
from backend.managers.system_prompter import SystemPrompter
|
from backend.managers.system_prompter import SystemPrompter
|
||||||
|
from backend.managers.search_manager import SearchManager
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
|
||||||
@ -212,13 +213,59 @@ def render_agent_mode():
|
|||||||
|
|
||||||
# ── Normal Chat ───────────────────────────────────────────────────────────────
|
# ── 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():
|
def render_normal_chat():
|
||||||
"""Render the standard multi-turn chat interface.
|
"""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
|
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).
|
# Replay the conversation history as chat bubbles (skip system messages).
|
||||||
for message in st.session_state.chat_history:
|
for message in st.session_state.chat_history:
|
||||||
role = message["role"]
|
role = message["role"]
|
||||||
@ -227,10 +274,54 @@ def render_normal_chat():
|
|||||||
with st.chat_message(role):
|
with st.chat_message(role):
|
||||||
st.markdown(message["content"])
|
st.markdown(message["content"])
|
||||||
|
|
||||||
# Chat input — Enter to send, no extra button needed
|
# Chat input — Enter to send, no extra button needed.
|
||||||
user_input = st.chat_input("Type your message here...")
|
# Supports /search <query> and /search clear as special commands.
|
||||||
|
user_input = st.chat_input("Type a message or /search <query>...")
|
||||||
if user_input:
|
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 <query> — 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
|
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
|
# On the very first user message, prepend the system prompt so the AI
|
||||||
# knows it is a code assistant embedded in an editor.
|
# knows it is a code assistant embedded in an editor.
|
||||||
@ -238,7 +329,22 @@ def render_normal_chat():
|
|||||||
system_prompt = SystemPrompter.generate_prompt()
|
system_prompt = SystemPrompter.generate_prompt()
|
||||||
chat_manager.add_message("system", system_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 = "<search_context>\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 += "</search_context>\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"):
|
with st.chat_message("user"):
|
||||||
st.markdown(user_input)
|
st.markdown(user_input)
|
||||||
|
|
||||||
@ -246,7 +352,7 @@ def render_normal_chat():
|
|||||||
with st.chat_message("assistant"):
|
with st.chat_message("assistant"):
|
||||||
with st.spinner("Thinking..."):
|
with st.spinner("Thinking..."):
|
||||||
try:
|
try:
|
||||||
ai_response = chat_manager.send_message(user_input)
|
ai_response = chat_manager.send_message(message_to_send)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
ai_response = f"Error: {e}"
|
ai_response = f"Error: {e}"
|
||||||
st.markdown(ai_response)
|
st.markdown(ai_response)
|
||||||
|
|||||||
@ -88,6 +88,11 @@ def init_state():
|
|||||||
if "agent_pending_action" not in st.session_state:
|
if "agent_pending_action" not in st.session_state:
|
||||||
st.session_state.agent_pending_action = None
|
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__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user