Skip to content

Commit 42f21ce

Browse files
authored
Merge branch 'main' into test/multisearch-unit-tests
2 parents 846d8e4 + a15ab00 commit 42f21ce

File tree

30 files changed

+2615
-1597
lines changed

30 files changed

+2615
-1597
lines changed

.github/workflows/run-ci-cd.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,7 @@ jobs:
176176

177177
- name: Run backend tests
178178
run: |
179-
docker run -e DJANGO_CONFIGURATION=Test owasp/nest:test-backend-latest pytest
179+
docker run -e DJANGO_SETTINGS_MODULE=settings.test --env-file backend/.env.example owasp/nest:test-backend-latest pytest
180180
181181
run-frontend-unit-tests:
182182
name: Run frontend unit tests

backend/Makefile

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,10 @@ test-backend:
116116
--cache-from nest-test-backend \
117117
-f backend/docker/Dockerfile.test backend \
118118
-t nest-test-backend
119-
@docker run -e DJANGO_CONFIGURATION=Test --rm nest-test-backend pytest
119+
@docker run \
120+
-e DJANGO_SETTINGS_MODULE=settings.test \
121+
--env-file backend/.env.example \
122+
--rm nest-test-backend pytest
120123

121124
update-backend-dependencies:
122125
@cd backend && poetry update

backend/apps/ai/Makefile

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,3 +17,7 @@ ai-create-project-chunks:
1717
ai-create-slack-message-chunks:
1818
@echo "Creating Slack message chunks"
1919
@CMD="python manage.py ai_create_slack_message_chunks" $(MAKE) exec-backend-command
20+
21+
ai-run-rag-tool:
22+
@echo "Running RAG tool"
23+
@CMD="python manage.py ai_run_rag_tool" $(MAKE) exec-backend-command

backend/apps/ai/agent/__init__.py

Whitespace-only changes.

backend/apps/ai/agent/tools/__init__.py

Whitespace-only changes.

backend/apps/ai/agent/tools/rag/__init__.py

Whitespace-only changes.
Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
"""Generator for the RAG system."""
2+
3+
import logging
4+
import os
5+
from typing import Any
6+
7+
import openai
8+
9+
logger = logging.getLogger(__name__)
10+
11+
12+
class Generator:
13+
"""Generates answers to user queries based on retrieved context."""
14+
15+
MAX_TOKENS = 2000
16+
SYSTEM_PROMPT = """
17+
You are a helpful and professional AI assistant for the OWASP Foundation.
18+
Your task is to answer user queries based ONLY on the provided context.
19+
Follow these rules strictly:
20+
1. Base your entire answer on the information given in the "CONTEXT" section. Do not use any
21+
external knowledge unless and until it is about OWASP.
22+
2. Do not mention or refer to the word "context", "based on context", "provided information",
23+
"Information given to me" or similar phrases in your responses.
24+
3. you will answer questions only related to OWASP and within the scope of OWASP.
25+
4. Be concise and directly answer the user's query.
26+
5. Provide the necessary link if the context contains a URL.
27+
6. If there is any query based on location, you need to look for latitude and longitude in the
28+
context and provide the nearest OWASP chapter based on that.
29+
7. You can ask for more information if the query is very personalized or user-centric.
30+
8. after trying all of the above, If the context does not contain the information or you think that
31+
it is out of scope for OWASP, you MUST state: "please ask question related to OWASP."
32+
"""
33+
TEMPERATURE = 0.4
34+
35+
def __init__(self, chat_model: str = "gpt-4o"):
36+
"""Initialize the Generator.
37+
38+
Args:
39+
chat_model (str): The name of the OpenAI chat model to use for generation.
40+
41+
Raises:
42+
ValueError: If the OpenAI API key is not set.
43+
44+
"""
45+
if not (openai_api_key := os.getenv("DJANGO_OPEN_AI_SECRET_KEY")):
46+
error_msg = "DJANGO_OPEN_AI_SECRET_KEY environment variable not set"
47+
raise ValueError(error_msg)
48+
49+
self.chat_model = chat_model
50+
self.openai_client = openai.OpenAI(api_key=openai_api_key)
51+
logger.info("Generator initialized with chat model: %s", self.chat_model)
52+
53+
def prepare_context(self, context_chunks: list[dict[str, Any]]) -> str:
54+
"""Format the list of retrieved context chunks into a single string for the LLM.
55+
56+
Args:
57+
context_chunks: A list of chunk dictionaries from the retriever.
58+
59+
Returns:
60+
A formatted string containing the context.
61+
62+
"""
63+
if not context_chunks:
64+
return "No context provided"
65+
66+
formatted_context = []
67+
for i, chunk in enumerate(context_chunks):
68+
source_name = chunk.get("source_name", f"Unknown Source {i + 1}")
69+
text = chunk.get("text", "")
70+
71+
context_block = f"Source Name: {source_name}\nContent: {text}"
72+
formatted_context.append(context_block)
73+
74+
return "\n\n---\n\n".join(formatted_context)
75+
76+
def generate_answer(self, query: str, context_chunks: list[dict[str, Any]]) -> str:
77+
"""Generate an answer to the user's query using provided context chunks.
78+
79+
Args:
80+
query: The user's query text.
81+
context_chunks: A list of context chunks retrieved by the retriever.
82+
83+
Returns:
84+
The generated answer as a string.
85+
86+
"""
87+
formatted_context = self.prepare_context(context_chunks)
88+
89+
user_prompt = f"""
90+
- You are an assistant for question-answering tasks related to OWASP.
91+
- Use the following pieces of retrieved context to answer the question.
92+
- If the question is related to OWASP then you can try to answer based on your knowledge, if you
93+
don't know the answer, just say that you don't know.
94+
- Try to give answer and keep the answer concise, but you really think that the response will be
95+
longer and better you will provide more information.
96+
- Ask for the current location if the query is related to location.
97+
- Ask for the information you need if the query is very personalized or user-centric.
98+
- Do not mention or refer to the word "context", "based on context", "provided information",
99+
"Information given to me" or similar phrases in your responses.
100+
Question: {query}
101+
Context: {formatted_context}
102+
Answer:
103+
"""
104+
105+
try:
106+
response = self.openai_client.chat.completions.create(
107+
model=self.chat_model,
108+
messages=[
109+
{"role": "system", "content": self.SYSTEM_PROMPT},
110+
{"role": "user", "content": user_prompt},
111+
],
112+
temperature=self.TEMPERATURE,
113+
max_tokens=self.MAX_TOKENS,
114+
)
115+
answer = response.choices[0].message.content.strip()
116+
except openai.OpenAIError:
117+
logger.exception("OpenAI API error")
118+
answer = "I'm sorry, I'm currently unable to process your request."
119+
120+
return answer
Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
"""A tool for orchestrating the components of RAG process."""
2+
3+
import logging
4+
5+
from apps.ai.common.constants import DEFAULT_CHUNKS_RETRIEVAL_LIMIT, DEFAULT_SIMILARITY_THRESHOLD
6+
7+
from .generator import Generator
8+
from .retriever import Retriever
9+
10+
logger = logging.getLogger(__name__)
11+
12+
13+
class RagTool:
14+
"""Main RAG tool that orchestrates the retrieval and generation process."""
15+
16+
def __init__(
17+
self,
18+
embedding_model: str = "text-embedding-3-small",
19+
chat_model: str = "gpt-4o",
20+
):
21+
"""Initialize the RAG tool.
22+
23+
Args:
24+
embedding_model (str, optional): The model to use for embeddings.
25+
chat_model (str, optional): The model to use for chat generation.
26+
27+
Raises:
28+
ValueError: If the OpenAI API key is not set.
29+
30+
"""
31+
try:
32+
self.retriever = Retriever(embedding_model=embedding_model)
33+
self.generator = Generator(chat_model=chat_model)
34+
except Exception:
35+
logger.exception("Failed to initialize RAG tool")
36+
raise
37+
38+
def query(
39+
self,
40+
question: str,
41+
content_types: list[str] | None = None,
42+
limit: int = DEFAULT_CHUNKS_RETRIEVAL_LIMIT,
43+
similarity_threshold: float = DEFAULT_SIMILARITY_THRESHOLD,
44+
) -> str:
45+
"""Process a user query using the complete RAG pipeline.
46+
47+
Args:
48+
question (str): The user's question.
49+
content_types (Optional[list[str]]): Content types to filter by.
50+
limit (int): Maximum number of context chunks to retrieve.
51+
similarity_threshold (float): Minimum similarity score for retrieval.
52+
53+
Returns:
54+
The generated answer as a string.
55+
56+
"""
57+
logger.info("Retrieving context for query")
58+
59+
return self.generator.generate_answer(
60+
context_chunks=self.retriever.retrieve(
61+
content_types=content_types,
62+
limit=limit,
63+
query=question,
64+
similarity_threshold=similarity_threshold,
65+
),
66+
query=question,
67+
)

0 commit comments

Comments
 (0)