RAG tool for agent (PoC for Nestbot) #1780
Conversation
Summary by CodeRabbit
WalkthroughThis change introduces a Retrieval-Augmented Generation (RAG) tool to the AI backend. It adds new classes for retrieval and generation, integrates them into a Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~40 minutes
Assessment against linked issues
Assessment against linked issues: Out-of-scope changesNo out-of-scope changes identified. All code additions relate directly to implementing and supporting the RAG tool as a PoC agent for local operation per the linked issue objectives. 📜 Recent review detailsConfiguration used: .coderabbit.yaml 📒 Files selected for processing (2)
✅ Files skipped from review due to trivial changes (2)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (2)
backend/apps/ai/agent/tools/RAG/generator.py (1)
16-32: Fix grammatical issues in system prompt.The system prompt contains several grammatical errors and awkward phrasing that should be corrected for clarity.
Apply these corrections:
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. +external knowledge unless 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. +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." +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 questions related to OWASP." """backend/apps/ai/agent/tools/RAG/retriever.py (1)
243-266: Consider more context-aware content type detectionThe current implementation might match content type names that appear in unrelated contexts (e.g., "message" in "error message"). Consider implementing more sophisticated detection that looks for content-type-specific context clues.
For example, you could look for patterns like "find {content_type}", "search for {content_type}", "get {content_type}", etc., or use a more sophisticated NLP approach to identify when these words are used as entity types rather than common words.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
backend/apps/ai/Makefile(1 hunks)backend/apps/ai/agent/tools/RAG/generator.py(1 hunks)backend/apps/ai/agent/tools/RAG/rag_tool.py(1 hunks)backend/apps/ai/agent/tools/RAG/retriever.py(1 hunks)backend/apps/ai/common/constants.py(1 hunks)backend/apps/ai/management/commands/ai_run_rag_tool.py(1 hunks)
🔇 Additional comments (6)
backend/apps/ai/common/constants.py (1)
4-5: LGTM! Well-defined constants for RAG defaults.The constants follow the existing naming convention and provide sensible defaults for the RAG tool's retrieval functionality.
backend/apps/ai/Makefile (1)
21-23: LGTM! Makefile target follows established patterns.The new target is well-integrated with the existing Makefile structure and provides a convenient way to run the RAG tool.
backend/apps/ai/agent/tools/RAG/rag_tool.py (1)
46-46: Confirmed Python 3.10+ CompatibilityThe project’s backend/pyproject.toml specifies
python = "^3.13"which implies a minimum of Python 3.13. Since Python 3.13 > 3.10, the
list[str] | Noneunion type syntax is fully supported. No changes required.backend/apps/ai/agent/tools/RAG/retriever.py (3)
64-71: Well-designed source name extractionGood approach to handle various content types by checking multiple common identifier attributes with a sensible fallback.
171-242: Well-implemented vector similarity searchExcellent implementation with proper use of pgvector, Django ORM optimizations (select_related/prefetch_related), and comprehensive error handling for missing content objects.
57-59: Update OpenAIError exception catch
Theopenai.error.OpenAIErrorclass doesn’t exist in the installed OpenAI Python SDK (v1.97.1). You should catchopenai.OpenAIErrorinstead to avoid runtime errors.Please update in
backend/apps/ai/agent/tools/RAG/retriever.py(around lines 57–59):- except openai.error.OpenAIError: + except openai.OpenAIError: logger.exception("OpenAI API error") raiseThis aligns with how errors are handled elsewhere (e.g. catching
openai.APIConnectionErrordirectly).Likely an incorrect or invalid review comment.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
backend/apps/ai/agent/tools/RAG/generator.py (2)
89-103: Resolve conflicting instructions between system and user prompts.The system prompt restricts answers to "ONLY the provided context" (line 18), but the user prompt allows answering "based on your knowledge" for OWASP questions (line 92). This contradiction could lead to inconsistent behavior.
Either align both prompts to be context-only or allow knowledge-based answers. Here's a suggested fix for consistency:
user_prompt = f""" - You are an assistant for question-answering tasks related to OWASP. - Use the following pieces of retrieved context to answer the question. -- If the question is related to OWASP then you can try to answer based on your knowledge, if you -don't know the answer, just say that you don't know. -- Try to give answer and keep the answer concise, but you really think that the response will be -longer and better you will provide more information. +- Base your answer ONLY on the provided context below. +- Keep the answer concise, but provide more detail if necessary for clarity. - Ask for the current location if the query is related to location. - Ask for the information you need if the query is very personalized or user-centric. - Do not mention or refer to the word "context", "based on context", "provided information", "Information given to me" or similar phrases in your responses. Question: {query} Context: {formatted_context} Answer: """
105-120: Enhance error handling to catch all exceptions.While the current implementation catches
OpenAIError, it should also handle other potential exceptions (e.g., network issues, JSON parsing errors).try: response = self.openai_client.chat.completions.create( model=self.chat_model, messages=[ {"role": "system", "content": self.SYSTEM_PROMPT}, {"role": "user", "content": user_prompt}, ], temperature=self.TEMPERATURE, max_tokens=self.MAX_TOKENS, ) final_answer = response.choices[0].message.content.strip() except openai.OpenAIError: logger.exception("OpenAI API error") return { "answer": "I'm sorry, I'm currently unable to process your request.", } + except Exception: + logger.exception("Unexpected error during generation") + return { + "answer": "An unexpected error occurred. Please try again later.", + }
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
backend/apps/ai/agent/tools/RAG/generator.py(1 hunks)backend/apps/ai/agent/tools/RAG/rag_tool.py(1 hunks)backend/apps/ai/agent/tools/RAG/retriever.py(1 hunks)backend/apps/ai/management/commands/ai_run_rag_tool.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- backend/apps/ai/agent/tools/RAG/rag_tool.py
- backend/apps/ai/management/commands/ai_run_rag_tool.py
- backend/apps/ai/agent/tools/RAG/retriever.py
🔇 Additional comments (3)
backend/apps/ai/agent/tools/RAG/generator.py (3)
1-14: LGTM!Imports are appropriate and the class structure is well-defined with a clear docstring.
35-52: LGTM!Well-implemented initialization with proper error handling for missing API key and appropriate logging.
53-75: LGTM!Clean implementation of context formatting with proper handling of edge cases and clear separation between context chunks.
|
@Dishant1804 it seems your PRs not longer trigger CI/CD required steps. No idea why. |
It is strange for me as well, I don't have any idea why CI/CD is not triggered when I push, it requires workflow approval everytime |
| """Generates answers to user queries based on retrieved context.""" | ||
|
|
||
| MAX_TOKENS = 2000 | ||
| SYSTEM_PROMPT = """ |
There was a problem hiding this comment.
We have a core.Prompt model for this. It's in core but will be migrated ot ai later.
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (5)
backend/apps/ai/agent/tools/rag/rag_tool.py (1)
28-30: Clarify the Raises section in docstringThe
ValueErroris not directly raised in this method but propagated fromRetrieverorGenerator. Consider updating the docstring to be more accurate.Raises: - ValueError: If the OpenAI API key is not set. + ValueError: If the OpenAI API key is not set (raised by Retriever or Generator). + Exception: Any other exception during initialization.backend/apps/ai/agent/tools/rag/generator.py (2)
17-32: Fix grammar and clarity issues in SYSTEM_PROMPTSeveral grammatical improvements needed for clarity:
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. +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." +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 questions related to OWASP."
89-103: Fix grammar issues in user promptSimilar grammar improvements needed as in the system prompt:
- You are an assistant for question-answering tasks related to OWASP. - Use the following pieces of retrieved context to answer the question. -- If the question is related to OWASP then you can try to answer based on your knowledge, if you -don't know the answer, just say that you don't know. -- Try to give answer and keep the answer concise, but you really think that the response will be -longer and better you will provide more information. +- If the question is related to OWASP, then you can try to answer based on your knowledge. If you +don't know the answer, just say that you don't know. +- Try to give an answer and keep it concise, but if you think that the response will be +better with more information, you can provide it.backend/apps/ai/agent/tools/rag/retriever.py (2)
75-183: Consider refactoring into smaller methods for better maintainabilityWhile the implementation is correct, this method is quite long (108 lines). Consider extracting each content type's context building into separate methods for better readability and maintainability.
Example refactor approach:
def get_additional_context(self, content_object, content_type: str) -> dict[str, Any]: """Get additional context information based on content type.""" clean_content_type = content_type.split(".")[-1] if "." in content_type else content_type context_methods = { "chapter": self._get_chapter_context, "project": self._get_project_context, "event": self._get_event_context, "committee": self._get_committee_context, "message": self._get_message_context, } if method := context_methods.get(clean_content_type): context = method(content_object) else: context = {} return {k: v for k, v in context.items() if v is not None} def _get_chapter_context(self, content_object) -> dict[str, Any]: """Extract context for chapter objects.""" return { "location": getattr(content_object, "suggested_location", None), "region": getattr(content_object, "region", None), # ... rest of chapter attributes } # Similar methods for other content types...
268-272: Consider improving readabilityWhile the list comprehension works, it might be more readable as a regular loop given the multiple conditions:
- detected_types = [ - content_type - for content_type in self.SUPPORTED_CONTENT_TYPES - if content_type in query_words or f"{content_type}s" in query_words - ] + detected_types = [] + for content_type in self.SUPPORTED_CONTENT_TYPES: + if content_type in query_words or f"{content_type}s" in query_words: + detected_types.append(content_type)
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
backend/apps/ai/agent/tools/rag/generator.py(1 hunks)backend/apps/ai/agent/tools/rag/rag_tool.py(1 hunks)backend/apps/ai/agent/tools/rag/retriever.py(1 hunks)backend/apps/ai/common/constants.py(1 hunks)backend/apps/ai/management/commands/ai_run_rag_tool.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- backend/apps/ai/common/constants.py
- backend/apps/ai/management/commands/ai_run_rag_tool.py
🔇 Additional comments (4)
backend/apps/ai/agent/tools/rag/generator.py (2)
35-52: LGTM!Good implementation with proper environment variable validation and error handling.
53-75: LGTM!Well-structured method with proper handling of empty context and clear formatting.
backend/apps/ai/agent/tools/rag/retriever.py (2)
67-74: LGTM!Nice implementation that gracefully handles different object types by trying multiple common name attributes.
184-254: LGTM!Excellent implementation with efficient database queries, proper content type filtering, and good null safety checks.
| """Generates answers to user queries based on retrieved context.""" | ||
|
|
||
| MAX_TOKENS = 2000 | ||
| SYSTEM_PROMPT = """ |
There was a problem hiding this comment.
This still needs to be moved into DB
| try: | ||
| self.retriever = Retriever(embedding_model=embedding_model) | ||
| self.generator = Generator(chat_model=chat_model) | ||
| except Exception: |
There was a problem hiding this comment.
This is not a good practice.
|
Proposed change
Resolves #1736
Checklist
make check-testlocally; all checks and tests passed.