128 lines
5.3 KiB
Python
128 lines
5.3 KiB
Python
"""Builds the system prompt that is sent to the AI at the start of each chat session."""
|
|
|
|
import ast
|
|
|
|
from backend.managers.debug_logger import get_logger
|
|
logger = get_logger(__name__)
|
|
|
|
# Prevents very large files from flooding the context window with tokens.
|
|
MAX_FILE_CHARS = 4000
|
|
|
|
# Per-task base prompts — selected via the task_type parameter.
|
|
_TASK_PROMPTS: dict[str, str] = {
|
|
"debug": (
|
|
"You are a debugging expert integrated into a lightweight code editor. "
|
|
"Focus on identifying and fixing errors. "
|
|
"Be concise and precise. Use markdown and fenced code blocks where appropriate."
|
|
),
|
|
"explain": (
|
|
"You are a code explainer integrated into a lightweight code editor. "
|
|
"Use simple language and examples. "
|
|
"Be concise and precise. Use markdown and fenced code blocks where appropriate."
|
|
),
|
|
"optimize": (
|
|
"You are a code optimization expert integrated into a lightweight code editor. "
|
|
"Focus on performance and readability. "
|
|
"Be concise and precise. Use markdown and fenced code blocks where appropriate."
|
|
),
|
|
"default": (
|
|
"You are an expert code assistant integrated into a lightweight code editor. "
|
|
"Help the user with code suggestions, debugging, explanations, and improvements. "
|
|
"Be concise and precise. Use markdown and fenced code blocks where appropriate."
|
|
),
|
|
}
|
|
|
|
|
|
def _extract_relevant_context(content: str, user_message: str) -> str:
|
|
"""Return the most relevant part of a Python file for the given user message.
|
|
|
|
Parses the file with ast and checks whether any top-level function or class
|
|
name appears in the user message. If a match is found only that definition
|
|
is returned, keeping the context focused. Falls back to simple truncation
|
|
when parsing fails or no name matches.
|
|
"""
|
|
try:
|
|
tree = ast.parse(content)
|
|
except SyntaxError:
|
|
# Not valid Python (or not Python at all) — fall back to truncation.
|
|
if len(content) > MAX_FILE_CHARS:
|
|
return content[:MAX_FILE_CHARS] + "\n... [truncated]"
|
|
return content
|
|
|
|
lower_msg = user_message.lower()
|
|
for node in tree.body:
|
|
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
|
|
if node.name.lower() in lower_msg:
|
|
segment = ast.get_source_segment(content, node)
|
|
if segment:
|
|
return segment
|
|
|
|
# No specific symbol matched — fall back to truncation.
|
|
if len(content) > MAX_FILE_CHARS:
|
|
return content[:MAX_FILE_CHARS] + "\n... [truncated]"
|
|
return content
|
|
|
|
|
|
class SystemPrompter:
|
|
"""Generates system prompts for the chat assistant.
|
|
|
|
When a file is open in the editor it can be embedded in the prompt so the
|
|
AI has direct context of the code the user is currently working on.
|
|
"""
|
|
|
|
@staticmethod
|
|
def generate_prompt(
|
|
user_message: str = "",
|
|
file_context: dict | None = None,
|
|
search_context: list[dict] | None = None,
|
|
task_type: str = "default",
|
|
) -> str:
|
|
"""Build a system prompt, optionally embedding a file and/or web search results.
|
|
|
|
Args:
|
|
user_message: The current user input — used for task-type detection
|
|
and selective context extraction. Reserved for future
|
|
task-specific prompt tuning beyond what task_type covers.
|
|
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.
|
|
task_type: One of "debug", "explain", "optimize", "default".
|
|
Selects the matching base prompt from _TASK_PROMPTS.
|
|
|
|
Returns:
|
|
A ready-to-use system prompt string.
|
|
"""
|
|
logger.info("Generating system prompt (task_type=%s).", task_type)
|
|
prompt = _TASK_PROMPTS.get(task_type, _TASK_PROMPTS["default"])
|
|
|
|
if file_context:
|
|
logger.info("Appending file context.")
|
|
name = file_context.get("name", "unknown")
|
|
content = file_context.get("content", "")
|
|
|
|
# Extract only the relevant function/class when the user mentions one;
|
|
# otherwise fall back to simple truncation at MAX_FILE_CHARS.
|
|
content = _extract_relevant_context(content, user_message)
|
|
|
|
prompt += (
|
|
f"\n\nThe user currently has the following file open in the editor:\n"
|
|
f"<file name=\"{name}\">\n"
|
|
f"<code>\n{content}\n</code>\n"
|
|
f"</file>\n"
|
|
f"Refer to this file when answering questions about the code."
|
|
)
|
|
|
|
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
|