Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 9 additions & 18 deletions backend/apps/ai/agent/tools/rag/generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,9 @@
from typing import Any

import openai
from django.core.exceptions import ObjectDoesNotExist

from apps.core.models.prompt import Prompt

logger = logging.getLogger(__name__)

Expand All @@ -13,23 +16,6 @@ class Generator:
"""Generates answers to user queries based on retrieved context."""

MAX_TOKENS = 2000
SYSTEM_PROMPT = """
You are a helpful and professional AI assistant for the OWASP Foundation.
Your task is to answer user queries based ONLY on the provided context.
Follow these rules strictly:
1. Base your entire answer on the information given in the "CONTEXT" section. Do not use any
external knowledge unless and until it is about OWASP.
2. Do not mention or refer to the word "context", "based on context", "provided information",
"Information given to me" or similar phrases in your responses.
3. you will answer questions only related to OWASP and within the scope of OWASP.
4. Be concise and directly answer the user's query.
5. Provide the necessary link if the context contains a URL.
6. If there is any query based on location, you need to look for latitude and longitude in the
context and provide the nearest OWASP chapter based on that.
7. You can ask for more information if the query is very personalized or user-centric.
8. after trying all of the above, If the context does not contain the information or you think that
it is out of scope for OWASP, you MUST state: "please ask question related to OWASP."
"""
TEMPERATURE = 0.4

def __init__(self, chat_model: str = "gpt-4o"):
Expand Down Expand Up @@ -103,10 +89,15 @@ def generate_answer(self, query: str, context_chunks: list[dict[str, Any]]) -> s
"""

try:
system_prompt = Prompt.get_rag_system_prompt()
if not system_prompt or not system_prompt.strip():
error_msg = "Prompt with key 'rag-system-prompt' not found."
raise ObjectDoesNotExist(error_msg)

response = self.openai_client.chat.completions.create(
model=self.chat_model,
messages=[
{"role": "system", "content": self.SYSTEM_PROMPT},
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
temperature=self.TEMPERATURE,
Expand Down
4 changes: 2 additions & 2 deletions backend/apps/ai/common/constants.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
"""AI app constants."""

DEFAULT_LAST_REQUEST_OFFSET_SECONDS = 2
DEFAULT_CHUNKS_RETRIEVAL_LIMIT = 5
DEFAULT_SIMILARITY_THRESHOLD = 0.4
DEFAULT_CHUNKS_RETRIEVAL_LIMIT = 8
DEFAULT_SIMILARITY_THRESHOLD = 0.1
DELIMITER = "\n\n"
GITHUB_REQUEST_INTERVAL_SECONDS = 0.5
MIN_REQUEST_INTERVAL_SECONDS = 1.2
Expand Down
20 changes: 20 additions & 0 deletions backend/apps/core/models/prompt.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,3 +139,23 @@ def get_owasp_project_summary() -> str:

"""
return Prompt.get_text("owasp-project-summary")

@staticmethod
def get_rag_system_prompt() -> str:
"""Return RAG system prompt.

Returns
str: The RAG system prompt text.

"""
return Prompt.get_text("rag-system-prompt")

@staticmethod
def get_slack_question_detector_prompt() -> str:
"""Return Slack question detector prompt.

Returns
str: The Slack question detector prompt text.

"""
return Prompt.get_text("slack-question-detector-system-prompt")
4 changes: 3 additions & 1 deletion backend/apps/slack/blocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@

from typing import Any

from apps.slack.utils import format_links_for_slack

DIVIDER = "{{ DIVIDER }}"
SECTION_BREAK = "{{ SECTION_BREAK }}"

Expand All @@ -30,7 +32,7 @@ def markdown(text: str) -> dict:
"""
return {
"type": "section",
"text": {"type": "mrkdwn", "text": text},
"text": {"type": "mrkdwn", "text": format_links_for_slack(text)},
}


Expand Down
21 changes: 8 additions & 13 deletions backend/apps/slack/common/question_detector.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,9 @@
import re

import openai
from django.core.exceptions import ObjectDoesNotExist

from apps.core.models.prompt import Prompt
from apps.slack.constants import OWASP_KEYWORDS

logger = logging.getLogger(__name__)
Expand All @@ -20,18 +22,6 @@ class QuestionDetector:
TEMPERATURE = 0.1
CHAT_MODEL = "gpt-4o"

SYSTEM_PROMPT = """
You are an expert in cybersecurity and OWASP (Open Web Application Security Project).
Your task is to determine if a given question is related to OWASP, cybersecurity,
web application security, or similar topics.

Key OWASP-related terms: {keywords}

Respond with only "YES" if the question is related to OWASP/cybersecurity,
or "NO" if it's not.
Do not provide any explanation or additional text.
"""

def __init__(self):
"""Initialize the question detector.

Expand Down Expand Up @@ -98,7 +88,12 @@ def is_owasp_question_with_openai(self, text: str) -> bool | None:
- None: If the API call fails or the response is unexpected.

"""
system_prompt = self.SYSTEM_PROMPT.format(keywords=", ".join(self.owasp_keywords))
prompt_template = Prompt.get_slack_question_detector_prompt()
if not prompt_template or not prompt_template.strip():
error_msg = "Prompt with key 'slack-question-detector-system-prompt' not found."
raise ObjectDoesNotExist(error_msg)

system_prompt = prompt_template.format(keywords=", ".join(self.owasp_keywords))
user_prompt = f'Question: "{text}"'

try:
Expand Down
17 changes: 17 additions & 0 deletions backend/apps/slack/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,23 @@ def escape(content) -> str:
return escape_html(content, quote=False)


def format_links_for_slack(text: str) -> str:
"""Convert Markdown links to Slack markdown link format.

Args:
text (str): The input text that may include Markdown links.

Returns:
str: Text with Markdown links converted to Slack markdown links.

"""
if not text:
return text

markdown_link_pattern = re.compile(r"\[([^\]]+)\]\((https?://[^\s)]+)\)")
return markdown_link_pattern.sub(r"<\2|\1>", text)


@lru_cache
def get_gsoc_projects(year: int) -> list:
"""Get GSoC projects.
Expand Down
140 changes: 133 additions & 7 deletions backend/tests/apps/ai/agent/tools/rag/generator_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

import openai
import pytest
from django.core.exceptions import ObjectDoesNotExist

from apps.ai.agent.tools.rag.generator import Generator

Expand Down Expand Up @@ -102,6 +103,10 @@ def test_generate_answer_success(self):
with (
patch.dict(os.environ, {"DJANGO_OPEN_AI_SECRET_KEY": "test-key"}),
patch("openai.OpenAI") as mock_openai,
patch(
"apps.core.models.prompt.Prompt.get_rag_system_prompt",
return_value="System prompt",
),
):
mock_client = MagicMock()
mock_response = MagicMock()
Expand All @@ -128,6 +133,10 @@ def test_generate_answer_with_custom_model(self):
with (
patch.dict(os.environ, {"DJANGO_OPEN_AI_SECRET_KEY": "test-key"}),
patch("openai.OpenAI") as mock_openai,
patch(
"apps.core.models.prompt.Prompt.get_rag_system_prompt",
return_value="System prompt",
),
):
mock_client = MagicMock()
mock_response = MagicMock()
Expand All @@ -150,6 +159,10 @@ def test_generate_answer_openai_error(self):
with (
patch.dict(os.environ, {"DJANGO_OPEN_AI_SECRET_KEY": "test-key"}),
patch("openai.OpenAI") as mock_openai,
patch(
"apps.core.models.prompt.Prompt.get_rag_system_prompt",
return_value="System prompt",
),
):
mock_client = MagicMock()
mock_client.chat.completions.create.side_effect = openai.OpenAIError("API Error")
Expand All @@ -167,6 +180,10 @@ def test_generate_answer_with_empty_chunks(self):
with (
patch.dict(os.environ, {"DJANGO_OPEN_AI_SECRET_KEY": "test-key"}),
patch("openai.OpenAI") as mock_openai,
patch(
"apps.core.models.prompt.Prompt.get_rag_system_prompt",
return_value="System prompt",
),
):
mock_client = MagicMock()
mock_response = MagicMock()
Expand All @@ -184,14 +201,123 @@ def test_generate_answer_with_empty_chunks(self):
assert "No context provided" in call_args[1]["messages"][1]["content"]

def test_system_prompt_content(self):
"""Test that system prompt contains expected content."""
assert "OWASP Foundation" in Generator.SYSTEM_PROMPT
assert "context" in Generator.SYSTEM_PROMPT.lower()
assert "professional" in Generator.SYSTEM_PROMPT.lower()
assert "latitude and longitude" in Generator.SYSTEM_PROMPT.lower()
"""Test that system prompt passed to OpenAI comes from Prompt getter."""
with (
patch.dict(os.environ, {"DJANGO_OPEN_AI_SECRET_KEY": "test-key"}),
patch("openai.OpenAI") as mock_openai,
patch(
"apps.core.models.prompt.Prompt.get_rag_system_prompt",
return_value="OWASP Foundation system prompt",
) as mock_prompt_getter,
):
mock_client = MagicMock()
mock_response = MagicMock()
mock_response.choices = [MagicMock()]
mock_response.choices[0].message.content = "Answer"
mock_client.chat.completions.create.return_value = mock_response
mock_openai.return_value = mock_client

generator = Generator()
generator.generate_answer("Q", [])

call_args = mock_client.chat.completions.create.call_args
assert call_args[1]["messages"][0]["role"] == "system"
assert call_args[1]["messages"][0]["content"] == "OWASP Foundation system prompt"
mock_prompt_getter.assert_called_once()

def test_generate_answer_missing_system_prompt(self):
"""Test answer generation when system prompt is missing."""
with (
patch.dict(os.environ, {"DJANGO_OPEN_AI_SECRET_KEY": "test-key"}),
patch("openai.OpenAI") as mock_openai,
patch(
"apps.core.models.prompt.Prompt.get_rag_system_prompt",
return_value=None,
),
):
mock_client = MagicMock()
mock_openai.return_value = mock_client

generator = Generator()

chunks = [{"source_name": "Test", "text": "Test content"}]

with pytest.raises(
ObjectDoesNotExist, match="Prompt with key 'rag-system-prompt' not found"
):
generator.generate_answer("Test query", chunks)

def test_generate_answer_empty_system_prompt(self):
"""Test answer generation when system prompt is empty."""
with (
patch.dict(os.environ, {"DJANGO_OPEN_AI_SECRET_KEY": "test-key"}),
patch("openai.OpenAI") as mock_openai,
patch(
"apps.core.models.prompt.Prompt.get_rag_system_prompt",
return_value=" ",
),
):
mock_client = MagicMock()
mock_openai.return_value = mock_client

generator = Generator()

chunks = [{"source_name": "Test", "text": "Test content"}]

with pytest.raises(
ObjectDoesNotExist, match="Prompt with key 'rag-system-prompt' not found"
):
generator.generate_answer("Test query", chunks)

def test_generate_answer_empty_openai_response(self):
"""Test answer generation when OpenAI returns empty content."""
with (
patch.dict(os.environ, {"DJANGO_OPEN_AI_SECRET_KEY": "test-key"}),
patch("openai.OpenAI") as mock_openai,
patch(
"apps.core.models.prompt.Prompt.get_rag_system_prompt",
return_value="System prompt",
),
):
mock_client = MagicMock()
mock_response = MagicMock()
mock_response.choices = [MagicMock()]
mock_response.choices[0].message.content = ""
mock_client.chat.completions.create.return_value = mock_response
mock_openai.return_value = mock_client

generator = Generator()

chunks = [{"source_name": "Test", "text": "Test content"}]
result = generator.generate_answer("Test query", chunks)

assert result == ""

def test_generate_answer_none_openai_response(self):
"""Test answer generation when OpenAI returns None content."""
with (
patch.dict(os.environ, {"DJANGO_OPEN_AI_SECRET_KEY": "test-key"}),
patch("openai.OpenAI") as mock_openai,
patch(
"apps.core.models.prompt.Prompt.get_rag_system_prompt",
return_value="System prompt",
),
):
mock_client = MagicMock()
mock_response = MagicMock()
mock_response.choices = [MagicMock()]
mock_response.choices[0].message.content = None
mock_client.chat.completions.create.return_value = mock_response
mock_openai.return_value = mock_client

generator = Generator()

chunks = [{"source_name": "Test", "text": "Test content"}]

with pytest.raises(AttributeError):
generator.generate_answer("Test query", chunks)

def test_constants(self):
"""Test class constants have expected values."""
assert Generator.MAX_TOKENS == 2000
assert isinstance(Generator.SYSTEM_PROMPT, str)
assert len(Generator.SYSTEM_PROMPT) > 0
assert Generator.TEMPERATURE == 0.4
8 changes: 4 additions & 4 deletions backend/tests/apps/ai/agent/tools/rag/rag_tool_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -124,9 +124,9 @@ def test_query_with_defaults(self):
assert result == "Default answer"
mock_retriever.retrieve.assert_called_once_with(
content_types=None,
limit=5,
limit=8,
query="Test question",
similarity_threshold=0.4,
similarity_threshold=0.1,
)

def test_query_empty_content_types(self):
Expand All @@ -152,9 +152,9 @@ def test_query_empty_content_types(self):
assert result == "Answer"
mock_retriever.retrieve.assert_called_once_with(
content_types=[],
limit=5,
limit=8,
query="Test question",
similarity_threshold=0.4,
similarity_threshold=0.1,
)

@patch("apps.ai.agent.tools.rag.rag_tool.logger")
Expand Down
Loading