Rag and Tex-to-SQL Agent done by Safni Usman #102
Conversation
WalkthroughThis update introduces a comprehensive query agent project integrating retrieval-augmented generation and SQL querying. New files include extensive documentation (README and requirements), a Streamlit-based UI application for handling chat and asynchronous queries, and several setup modules for configuring Llama Cloud services, SQL databases, and associated query engine tools. Additionally, an asynchronous routing workflow is implemented using custom event classes to manage chat interactions and tool calls. Changes
Sequence Diagram(s)sequenceDiagram
participant U as User
participant A as RAGSQLAgentApp
participant Q as Query Engine
participant S as Service (SQL/LlamaCloud)
U->>A: Enter query
A->>A: Update chat session state
A->>Q: Execute run_query(query)
Q->>S: Process the query
S-->>Q: Return result
Q-->>A: Return response
A-->>U: Display query response
sequenceDiagram
participant U as User
participant R as RouterOutputAgentWorkflow
participant L as Language Model (LLM)
participant T as Tool Caller
U->>R: Send input query
R->>R: Prepare and validate chat message
R->>L: Request generated response
L-->>R: Provide response with tool call info
R->>T: Dispatch tool call events
T-->>R: Return tool call results
R-->>U: Aggregate results and send final response
Poem
Tip ⚡🧪 Multi-step agentic review comment chat (experimental)
✨ Finishing Touches
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:
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
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (20)
rag-sql-agent/requirements.txt (1)
1-7: Pin package versions for reproducibility and securityWhile the requirements.txt file includes all necessary dependencies for the application, it doesn't pin specific package versions. This could lead to inconsistent environments and potential breaking changes when dependencies are updated.
Consider updating the requirements.txt file with specific versions:
-streamlit -llama-index -sqlalchemy -pyvis -python-dotenv -arize-phoenix -llama-index-callbacks-arize-phoenix +streamlit==1.30.0 +llama-index==0.10.0 +sqlalchemy==2.0.27 +pyvis==0.3.2 +python-dotenv==1.0.0 +arize-phoenix==1.1.0 +llama-index-callbacks-arize-phoenix==0.1.0The versions above are examples - please use the specific versions tested with your application.
rag-sql-agent/README.md (4)
8-18: Fix heading hierarchy in READMEThe heading "Use Case" jumps from h1 (#) to h3 (###) without using h2 (##) first, which violates standard markdown hierarchy.
-### **Use Case** +## **Use Case**🧰 Tools
🪛 markdownlint-cli2 (0.17.2)
8-8: Heading levels should only increment by one level at a time
Expected: h2; Actual: h3(MD001, heading-increment)
24-24: Use consistent formatting for GPT-3.5The model name formatting should be consistent with industry standard.
-**GPT 3.5 turbo** – as the LLM. +**GPT-3.5 Turbo** – as the LLM.🧰 Tools
🪛 LanguageTool
[uncategorized] ~24-~24: Did you mean “GPT-3.5”?
Context: ... - Streamlit – to build the UI. - GPT 3.5 turbo – as the LLM. ## Demo - [...(EN_SIMPLE_REPLACE_GPT____)
10-18: Enhance documentation for PDF preparationThe README mentions using PDFs of Wikipedia pages but doesn't provide instructions on how to obtain, convert, or prepare these PDFs for use with the application.
Consider adding a section explaining how to:
- Download or convert the Wikipedia pages to PDF format
- Where to store these files for the application to access them
- Any preprocessing steps needed before the PDFs can be used with LlamaCloud
59-73: Add SQL database setup instructionsThe installation section covers API keys and dependencies but doesn't mention how to set up the SQL database referenced in the code.
Consider adding instructions for:
- Database creation or configuration
- Schema setup
- Sample data population
- Connection configuration
This would make it easier for users to fully replicate the project environment.
rag-sql-agent/llama_cloud_setup.py (3)
7-11: Add error handling for missing environment variablesThe code doesn't check if the required environment variables are present and valid before attempting to use them.
def __init__(self): load_dotenv() self.llama_cloud_api_key = os.getenv('LLAMA_CLOUD_API_KEY') self.phoenix_api_key = os.getenv('PHOENIX_API_KEY') + + if not self.llama_cloud_api_key: + raise ValueError("LLAMA_CLOUD_API_KEY environment variable is not set") + if not self.phoenix_api_key: + raise ValueError("PHOENIX_API_KEY environment variable is not set")
6-7: Add class docstringThe class is missing a docstring that explains its purpose and usage.
class LlamaCloudSetup: + """ + Sets up and manages the LlamaCloud index for RAG-based retrieval. + + This class handles initialization of the LlamaCloud index, + configures observability via Arize Phoenix, and provides + access to the query engine. + """ def __init__(self):
28-31: Add error handling for query engine retrievalThe
get_query_engine()method doesn't handle potential errors that might occur when accessing the index.def get_query_engine(self): """Returns the query engine instance.""" - return self.index.as_query_engine() + try: + return self.index.as_query_engine() + except Exception as e: + raise RuntimeError(f"Failed to create query engine: {str(e)}")rag-sql-agent/tool_setup.py (4)
1-3: Fix import statements to use explicit relative importsThe import statements use implicit relative imports, which is not recommended in Python and may cause issues in some scenarios.
from llama_index.core.tools import QueryEngineTool -from sql_setup import SQLSetup -from llama_cloud_setup import LlamaCloudSetup +from .sql_setup import SQLSetup +from .llama_cloud_setup import LlamaCloudSetup
5-17: Add error handling and comprehensive class docstringThe class lacks a docstring explaining its purpose and doesn't handle potential initialization errors.
class QueryEngineTools: + """ + Initializes and manages query engine tools for SQL and LlamaCloud. + + This class encapsulates the setup of both SQL and LlamaCloud query engines, + wrapping them as tools that can be used by an agent to answer questions + through either structured data queries or semantic retrieval. + """ def __init__(self): - self.sql_setup = SQLSetup() - self.llama_cloud_setup = LlamaCloudSetup() + try: + self.sql_setup = SQLSetup() + self.llama_cloud_setup = LlamaCloudSetup() - self.sql_query_engine = self.sql_setup.get_sql_query_engine() - self.llama_cloud_query_engine = self.llama_cloud_setup.get_query_engine() + self.sql_query_engine = self.sql_setup.get_sql_query_engine() + self.llama_cloud_query_engine = self.llama_cloud_setup.get_query_engine() + except Exception as e: + raise RuntimeError(f"Failed to initialize query engines: {str(e)}")
18-34: Make tool descriptions configurableThe tool descriptions are hardcoded in the initialization method. Making them configurable would improve flexibility.
+ def __init__(self, sql_description=None, llama_cloud_description=None): + try: + self.sql_setup = SQLSetup() + self.llama_cloud_setup = LlamaCloudSetup() + + self.sql_query_engine = self.sql_setup.get_sql_query_engine() + self.llama_cloud_query_engine = self.llama_cloud_setup.get_query_engine() + + self.sql_description = sql_description or ( + "Useful for translating a natural language query into a SQL query over" + " a table containing: city_stats, containing the population/state of" + " each city located in the USA." + ) + + self.llama_cloud_description = llama_cloud_description or ( + "Useful for answering semantic questions about certain cities in the US." + ) + except Exception as e: + raise RuntimeError(f"Failed to initialize query engines: {str(e)}") def _initialize_tools(self): """Initialize query engine tools.""" self.sql_tool = QueryEngineTool.from_defaults( query_engine=self.sql_query_engine, - description=( - "Useful for translating a natural language query into a SQL query over" - " a table containing: city_stats, containing the population/state of" - " each city located in the USA." - ), + description=self.sql_description, name="sql_tool" ) self.llama_cloud_tool = QueryEngineTool.from_defaults( query_engine=self.llama_cloud_query_engine, - description="Useful for answering semantic questions about certain cities in the US.", + description=self.llama_cloud_description, name="llama_cloud_tool" )
36-42: Add validation in getter methodsThe getter methods don't check if the tools were properly initialized before returning them.
def get_sql_tool(self): """Returns the SQL query tool.""" + if self.sql_tool is None: + raise ValueError("SQL tool was not properly initialized") return self.sql_tool def get_llama_cloud_tool(self): """Returns the LlamaCloud query tool.""" + if self.llama_cloud_tool is None: + raise ValueError("LlamaCloud tool was not properly initialized") return self.llama_cloud_toolrag-sql-agent/sql_setup.py (3)
19-24: Reevaluate the fixed string length for city names.Currently,
city_nameis limited to 16 characters. Some city names and areas might exceed 16 characters. Consider increasing this limit (e.g., toString(64)or higher) to avoid potential truncation.-Column("city_name", String(16), primary_key=True), +Column("city_name", String(64), primary_key=True),
26-31: Replace prints with logging.Using
logging) for consistent message handling.
41-54: Ensure consistency of initial data insertion across runs.While the table creation is skipped if it already exists, the insertion process also depends on this logic. If you plan to add or update rows over time, consider using an incremental or migration-based approach rather than only inserting data on table creation.
rag-sql-agent/app.py (2)
26-31: Validate existence of assets.Images referenced at
./assets/my_image.pngand others might cause errors if the assets folder or file names change. Consider usingpathlib.Pathchecks or a try-catch around image loading to gracefully handle missing assets.
77-79: Prevent unbounded chat history growth.Storing all chat messages in
st.session_statecan lead to high memory usage over long sessions. Consider implementing a maximum history length or offering a way to reset the conversation to keep resource usage manageable.rag-sql-agent/query_agent.py (3)
2-16: Remove unused imports for a cleaner codebase.Static analysis shows the following unused imports:
IPython.display.displayIPython.display.Markdownasynciollama_index.core.base.response.schema.Responsellama_index.core.tools.FunctionToolConsider removing these imports to reduce code clutter and potential confusion.
-from IPython.display import display, Markdown -import asyncio -from llama_index.core.base.response.schema import Response -from llama_index.core.tools import FunctionTool🧰 Tools
🪛 Ruff (0.8.2)
2-2:
IPython.display.displayimported but unusedRemove unused import
(F401)
2-2:
IPython.display.Markdownimported but unusedRemove unused import
(F401)
3-3:
asyncioimported but unusedRemove unused import:
asyncio(F401)
15-15:
llama_index.core.base.response.schema.Responseimported but unusedRemove unused import:
llama_index.core.base.response.schema.Response(F401)
16-16:
llama_index.core.tools.FunctionToolimported but unusedRemove unused import:
llama_index.core.tools.FunctionTool(F401)
111-114: Add error handling for failed tool calls.If a tool raises an exception internally, the
await tool.acall(...)call might fail. You could wrap this in a try-except block to record the failure gracefully in the chat history rather than halting the workflow.- output = await tool.acall(**tool_call.tool_kwargs) + try: + output = await tool.acall(**tool_call.tool_kwargs) + except Exception as e: + output = f"Tool call failed with error: {str(e)}"
140-153: Enforce consistency across multiple tool results.The current gather step simply collects events after the number of tool calls is reached. If any tool fails or returns partial results, the subsequent logic might need special handling (e.g., skipping partial results or combining them). Consider adding custom logic to handle partial completions or cancellations.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (6)
rag-sql-agent/assets/cover.pngis excluded by!**/*.pngrag-sql-agent/assets/image1.pngis excluded by!**/*.pngrag-sql-agent/assets/image2.pngis excluded by!**/*.pngrag-sql-agent/assets/image3.pngis excluded by!**/*.pngrag-sql-agent/assets/my_image.pngis excluded by!**/*.pngrag-sql-agent/demo.mp4is excluded by!**/*.mp4
📒 Files selected for processing (7)
rag-sql-agent/README.md(1 hunks)rag-sql-agent/app.py(1 hunks)rag-sql-agent/llama_cloud_setup.py(1 hunks)rag-sql-agent/query_agent.py(1 hunks)rag-sql-agent/requirements.txt(1 hunks)rag-sql-agent/sql_setup.py(1 hunks)rag-sql-agent/tool_setup.py(1 hunks)
🧰 Additional context used
🪛 LanguageTool
rag-sql-agent/README.md
[uncategorized] ~24-~24: Did you mean “GPT-3.5”?
Context: ... - Streamlit – to build the UI. - GPT 3.5 turbo – as the LLM. ## Demo - [...
(EN_SIMPLE_REPLACE_GPT____)
[style] ~85-~85: Using many exclamation marks might seem excessive (in this case: 4 exclamation marks for a text that’s 1575 characters long)
Context: ... Contribution Contributions are welcome! Please fork the repository and submit a...
(EN_EXCESSIVE_EXCLAMATION)
🪛 markdownlint-cli2 (0.17.2)
rag-sql-agent/README.md
8-8: Heading levels should only increment by one level at a time
Expected: h2; Actual: h3
(MD001, heading-increment)
🪛 Ruff (0.8.2)
rag-sql-agent/query_agent.py
2-2: IPython.display.display imported but unused
Remove unused import
(F401)
2-2: IPython.display.Markdown imported but unused
Remove unused import
(F401)
3-3: asyncio imported but unused
Remove unused import: asyncio
(F401)
15-15: llama_index.core.base.response.schema.Response imported but unused
Remove unused import: llama_index.core.base.response.schema.Response
(F401)
16-16: llama_index.core.tools.FunctionTool imported but unused
Remove unused import: llama_index.core.tools.FunctionTool
(F401)
🔇 Additional comments (4)
rag-sql-agent/sql_setup.py (2)
9-11: Handle missing.envor environment variable gracefully.If
OPENAI_API_KEYis not set in the environment, this code may break or lead to unintended behavior. Consider adding fallback logic or raising a clear exception if the key is missing.Would you like a helper script to confirm that
.envexists and the environment variable is properly loaded?
59-62: Good usage example.The usage example is clear and demonstrates how to construct and use the
SQLSetupclass. Make sure to keep these examples up-to-date as the code evolves.rag-sql-agent/app.py (1)
19-21: Leverage Streamlit's async features carefully.The
run_querymethod is declared async, but is called using a customloop.run_until_complete(...). For future improvements, consider using Streamlit's async-compatible patterns orasyncio.run()with caution, as multiple event loops can be tricky.Are you comfortable with the current approach? If you prefer, we can explore a more streamlined async pattern.
rag-sql-agent/query_agent.py (1)
76-85: Handle concurrency with parallel tool calls carefully.
allow_parallel_tool_calls=Truecan invoke multiple tools simultaneously. Verify that all tools are thread-safe or concurrent-safe to avoid race conditions or shared resource conflicts.
| self.index = LlamaCloudIndex( | ||
| name="city_pdfs_index", | ||
| project_name="Default", | ||
| organization_id="8c1b26ae-5c98-4873-bce9-28d0c67a3dce", | ||
| api_key=self.llama_cloud_api_key, | ||
| ) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Move hardcoded values to environment variables
The organization_id and other configuration parameters are hardcoded in the class. This makes the code less flexible and potentially poses a security risk.
self.index = LlamaCloudIndex(
- name="city_pdfs_index",
- project_name="Default",
- organization_id="8c1b26ae-5c98-4873-bce9-28d0c67a3dce",
+ name=os.getenv('LLAMA_CLOUD_INDEX_NAME', 'city_pdfs_index'),
+ project_name=os.getenv('LLAMA_CLOUD_PROJECT_NAME', 'Default'),
+ organization_id=os.getenv('LLAMA_CLOUD_ORGANIZATION_ID'),
api_key=self.llama_cloud_api_key,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| self.index = LlamaCloudIndex( | |
| name="city_pdfs_index", | |
| project_name="Default", | |
| organization_id="8c1b26ae-5c98-4873-bce9-28d0c67a3dce", | |
| api_key=self.llama_cloud_api_key, | |
| ) | |
| self.index = LlamaCloudIndex( | |
| name=os.getenv('LLAMA_CLOUD_INDEX_NAME', 'city_pdfs_index'), | |
| project_name=os.getenv('LLAMA_CLOUD_PROJECT_NAME', 'Default'), | |
| organization_id=os.getenv('LLAMA_CLOUD_ORGANIZATION_ID'), | |
| api_key=self.llama_cloud_api_key, | |
| ) |
| def __init__(self): | ||
| """Initialize the application and workflow.""" | ||
| self.tools = QueryEngineTools() | ||
| self.workflow = RouterOutputAgentWorkflow( | ||
| tools=[self.tools.get_sql_tool(), self.tools.get_llama_cloud_tool()], | ||
| verbose=True, | ||
| timeout=120 | ||
| ) | ||
| self._setup_ui() |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Consider error handling for missing or failing tools.
The constructor sets up the workflow with get_sql_tool() and get_llama_cloud_tool(). If either fails or isn't configured properly, the workflow may break. You might wrap tool initialization (and environment checks) in try-catch logic to report clearer error messages to users.
| loop = asyncio.new_event_loop() | ||
| asyncio.set_event_loop(loop) | ||
| result = loop.run_until_complete(self.run_query(query)) | ||
|
|
||
| with st.chat_message("assistant"): | ||
| st.write(result) | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add exception handling for failed queries.
If workflow.run(message=query) raises an exception (e.g., a tool error or a network failure), the app might crash. Wrap this call in a try-except block to show a user-friendly error message instead of an internal traceback.
- result = loop.run_until_complete(self.run_query(query))
+ try:
+ result = loop.run_until_complete(self.run_query(query))
+ except Exception as e:
+ st.error(f"An error occurred while processing your query: {e}")
+ return📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| loop = asyncio.new_event_loop() | |
| asyncio.set_event_loop(loop) | |
| result = loop.run_until_complete(self.run_query(query)) | |
| with st.chat_message("assistant"): | |
| st.write(result) | |
| loop = asyncio.new_event_loop() | |
| asyncio.set_event_loop(loop) | |
| try: | |
| result = loop.run_until_complete(self.run_query(query)) | |
| except Exception as e: | |
| st.error(f"An error occurred while processing your query: {e}") | |
| return | |
| with st.chat_message("assistant"): | |
| st.write(result) |
This is task1 of technical writer assessment.
This project implements a custom agent for RAG and Text-to-SQL in a single query interface.
Find the X thread here : https://x.com/senasafni/status/1900779811116634298
Find the Typefully draft here (without video) : https://typefully.com/t/jddC0rK
Summary by CodeRabbit
New Features
Documentation