-
Notifications
You must be signed in to change notification settings - Fork 1.4k
Implementation of RAG and Text-to-SQL in a Single Query Interface #80
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
WalkthroughThis pull request introduces a unified query interface combining Retrieval-Augmented Generation and Text-to-SQL. A new README provides installation, setup, and API key instructions. A Streamlit application (in Changes
Sequence Diagram(s)sequenceDiagram
participant U as User
participant S as Streamlit UI
participant Q as Query Function (app.py)
participant DB as SQLDatabase
participant QE as NLSQLTableQueryEngine
U->>S: Enter query via Streamlit
S->>Q: Pass user query
Q->>DB: Query city statistics
DB-->>QE: Provide dataset for query processing
QE->>Q: Return query result
Q-->>S: Display result on UI
sequenceDiagram
participant U as User
participant A as RouterOutputAgentWorkflow
participant SQE as SQLQueryEngine
participant LCI as LlamaCloudIndex
U->>A: Submit natural language query
A->>A: prepare_chat() → Generate InputEvent
A->>A: chat() → Create GatherToolsEvent (or StopEvent)
A->>SQE: dispatch_calls() → Initiate ToolCallEvent for SQL query
SQE-->>A: Return ToolCallEventResult
A->>LCI: Optionally call LlamaCloud tool if needed
LCI-->>A: Return tool result
A->>A: gather() → Aggregate results into StopEvent
A-->>U: Deliver final answer
Poem
Tip ⚡🧪 Multi-step agentic review comment chat (experimental)
🪧 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 (
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 2
🧹 Nitpick comments (5)
RAG and Text-to-SQL/README.md (2)
14-14
: Fix malformed hyperlink for OpenAI.Line 14's link syntax is off, and "OpenAihttps://platform.openai.com/" should be changed to a properly formatted hyperlink, for example:
Get an API key from [OpenAI](https://platform.openai.com/) and set it in the `.env` file as follows:
1-42
: Add guidance on environment usage and highlight best practices.While the overall README is straightforward, consider emphasizing:
- The need for a recent GCC or equivalent on Windows (e.g. MSVC) to avoid potential NumPy installation errors.
- The importance of not committing any API keys to source control.
RAG and Text-to-SQL/app.py (1)
65-69
: Add error handling for query execution.If the SQL engine or query fails, the UI might hang or crash. Consider wrapping the query call in a try/except block to handle errors gracefully.
with st.spinner("Processing..."): + try: response = sql_query_engine.query(user_query) st.write("### Answer:") st.write(response.response) + except Exception as e: + st.error(f"Error executing query: {str(e)}")RAG and Text-to-SQL/llamacloud_sql_router.ipynb (2)
60-90
: Address installation errors with NumPy and ensure environment compatibility.The logs indicate that NumPy requires GCC >= 8.4, but the system has gcc 6.3.0. Consider:
- Updating your environment or compiler to meet these requirements.
- Documenting this requirement in the README to alert users ahead of time.
207-226
: Handle missing Arize Phoenix integration more gracefully.The import error suggests
llama-index-callbacks-arize-phoenix
isn't fully installed due to environment issues. Provide fallback logic or a clear error message if the user chooses not to enable observability, and instruct them on installing/updating dependencies.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
RAG and Text-to-SQL/README.md
(1 hunks)RAG and Text-to-SQL/app.py
(1 hunks)RAG and Text-to-SQL/llamacloud_sql_router.ipynb
(1 hunks)
RAG and Text-to-SQL/app.py
Outdated
from llama_index.core.tools import QueryEngineTool | ||
|
||
# Set up API key (Ensure to store securely in production) | ||
os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>" |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
Avoid setting API keys in plain code.
Using os.environ["OPENAI_API_KEY"]
directly in the code is insecure for production. Recommend:
- Loading secrets from a
.env
file or Streamlitst.secrets
. - Avoid committing secrets into version control.
- os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>"
+ # Instead, read from environment or st.secrets in production:
+ import dotenv
+ dotenv.load_dotenv()
+ # Now the OPENAI_API_KEY can be set in .env and read automatically.
📝 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.
os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>" | |
# Instead, read from environment or st.secrets in production: | |
import dotenv | |
dotenv.load_dotenv() | |
# Now the OPENAI_API_KEY can be set in .env and read automatically. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 1
🧹 Nitpick comments (5)
RAG and Text-to-SQL/app.py (5)
11-18
: Add API key validation to prevent runtime errors.While you're correctly loading the API key from environment variables, there's no validation to check if the key is actually present. If a user runs this application without setting the OPENAI_API_KEY environment variable, they'll encounter a cryptic error later in execution.
# Access the API key securely openai_api_key = os.getenv("OPENAI_API_KEY") + if not openai_api_key: + st.error("OPENAI_API_KEY not found in environment variables. Please set it and restart the application.") + st.stop()
20-21
: Consider making the OpenAI model configurable.The model is currently hardcoded to "gpt-3.5-turbo". Consider allowing users to select different models or making it configurable through environment variables for flexibility.
# Configure OpenAI LLM - Settings.llm = OpenAI("gpt-3.5-turbo") + model_name = os.getenv("OPENAI_MODEL", "gpt-3.5-turbo") + Settings.llm = OpenAI(model_name)
38-53
: Move the import statement to the top of the file.According to PEP 8 style guidelines, all imports should be grouped at the beginning of the file. Move the sqlite_insert import to join the other imports.
- from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer + from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer + from sqlalchemy.dialects.sqlite import insert as sqlite_insert# Insert sample data - from sqlalchemy.dialects.sqlite import insert as sqlite_insert
1-74
: Add a main function to improve code organization.The current code structure executes all statements at the module level. It's a best practice to wrap the application logic in a main function and use the if name == "main": pattern to make the code more maintainable and testable.
# Streamlit UI - st.title("RAG and Text-to-SQL") - st.write("Ask questions about city populations and states!") - - user_query = st.text_input("Enter your query:", "Which city has the highest population?") - - if st.button("Get Answer"): - with st.spinner("Processing..."): - response = sql_query_engine.query(user_query) - st.write("### Answer:") - st.write(response.response) + def main(): + st.title("RAG and Text-to-SQL") + st.write("Ask questions about city populations and states!") + + user_query = st.text_input("Enter your query:", "Which city has the highest population?") + + if st.button("Get Answer"): + with st.spinner("Processing..."): + try: + response = sql_query_engine.query(user_query) + st.write("### Answer:") + st.write(response.response) + except Exception as e: + st.error(f"An error occurred: {str(e)}") + st.info("Please try rephrasing your query or contact support if the issue persists.") + + if __name__ == "__main__": + main()
68-69
: Enhance the user experience with more query examples.Consider providing a few example queries as buttons or in an expander to help users understand the capabilities of the system.
user_query = st.text_input("Enter your query:", "Which city has the highest population?") + # Example queries to help users + st.expander("Example queries you can try:").write(""" + - Which city has the highest population? + - How many people live in Chicago? + - List all cities in California + - What is the average population of cities in the database? + - Which state has the most populous city? + """)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
RAG and Text-to-SQL/app.py
(1 hunks)
🔇 Additional comments (4)
RAG and Text-to-SQL/app.py (4)
1-10
: Good job on using the appropriate libraries and avoiding hardcoded API keys.You've correctly imported the necessary libraries for this application and used a proper approach for handling environment variables with dotenv. I also appreciate the use of nest_asyncio for Streamlit compatibility.
23-36
: Good database schema design for the demo.The table structure is well-defined with appropriate columns and types. Using an in-memory SQLite database is appropriate for this demonstration application.
48-53
: Well-implemented conflict handling for database inserts.The use of on_conflict_do_update is a good practice for handling potential duplicate entries, ensuring the application won't crash if it's restarted or if data is reloaded.
54-62
: Good job setting up the query engine with clear tool description.The SQL query engine setup is well-structured, and I appreciate the descriptive name and documentation for the query engine tool.
# Streamlit UI | ||
st.title("RAG and Text-to-SQL") | ||
st.write("Ask questions about city populations and states!") | ||
|
||
user_query = st.text_input("Enter your query:", "Which city has the highest population?") | ||
|
||
if st.button("Get Answer"): | ||
with st.spinner("Processing..."): | ||
response = sql_query_engine.query(user_query) | ||
st.write("### Answer:") | ||
st.write(response.response) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
Add error handling for query execution.
The current implementation lacks error handling for the query execution. If the API call fails or if there are issues with parsing the query, the application will crash without providing meaningful feedback to the user.
if st.button("Get Answer"):
with st.spinner("Processing..."):
- response = sql_query_engine.query(user_query)
- st.write("### Answer:")
- st.write(response.response)
+ try:
+ response = sql_query_engine.query(user_query)
+ st.write("### Answer:")
+ st.write(response.response)
+ except Exception as e:
+ st.error(f"An error occurred: {str(e)}")
+ st.info("Please try rephrasing your query or contact support if the issue persists.")
📝 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.
# Streamlit UI | |
st.title("RAG and Text-to-SQL") | |
st.write("Ask questions about city populations and states!") | |
user_query = st.text_input("Enter your query:", "Which city has the highest population?") | |
if st.button("Get Answer"): | |
with st.spinner("Processing..."): | |
response = sql_query_engine.query(user_query) | |
st.write("### Answer:") | |
st.write(response.response) | |
# Streamlit UI | |
st.title("RAG and Text-to-SQL") | |
st.write("Ask questions about city populations and states!") | |
user_query = st.text_input("Enter your query:", "Which city has the highest population?") | |
if st.button("Get Answer"): | |
with st.spinner("Processing..."): | |
try: | |
response = sql_query_engine.query(user_query) | |
st.write("### Answer:") | |
st.write(response.response) | |
except Exception as e: | |
st.error(f"An error occurred: {str(e)}") | |
st.info("Please try rephrasing your query or contact support if the issue persists.") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 4
♻️ Duplicate comments (1)
RAG and Text-to-SQL/llamacloud_sql_router.ipynb (1)
236-237
:⚠️ Potential issueAvoid hardcoding API keys in code.
API keys should be loaded from environment variables (.env file) rather than being hardcoded in the notebook, which presents a security risk.
-PHOENIX_API_KEY = "<phoenix_api_key>" -os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = f"api_key={PHOENIX_API_KEY}" +# Load the Phoenix API key from environment variables +PHOENIX_API_KEY = os.environ.get("PHOENIX_API_KEY", "") +if not PHOENIX_API_KEY: + print("Warning: PHOENIX_API_KEY not found in environment variables") +else: + os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = f"api_key={PHOENIX_API_KEY}"
🧹 Nitpick comments (8)
RAG and Text-to-SQL/llamacloud_sql_router.ipynb (8)
32-36
: API key loading approach is secure but needs validation.The code uses dotenv for loading environment variables, which is a secure practice for handling API keys. However, there's no validation to ensure the API key is actually loaded.
import os from dotenv import load_dotenv # Load environment variables from .env file load_dotenv() +# Validate that required API keys are present +if "OPENAI_API_KEY" not in os.environ or not os.environ["OPENAI_API_KEY"]: + raise ValueError("OPENAI_API_KEY environment variable is not set or empty")
286-286
: Explicitly include model selection in the Settings configuration.The code currently uses a hardcoded model name, but it should be configurable.
-Settings.llm = OpenAI("gpt-3.5-turbo") +# Allow for model configuration from environment variables with fallback +model_name = os.environ.get("OPENAI_MODEL", "gpt-3.5-turbo") +Settings.llm = OpenAI(model_name)
333-340
: Refactor database insertion for better error handling and efficiency.The current implementation inserts rows one by one, which is inefficient for larger datasets. Also, there's no error handling for database operations.
-# Insert rows with conflict handling -for row in rows: - stmt = sqlite_insert(city_stats_table).values(**row) - stmt = stmt.on_conflict_do_update( - index_elements=['city_name'], # Column(s) with the UNIQUE constraint - set_=row # Update the row with the new values if a conflict occurs - ) - with engine.begin() as connection: - connection.execute(stmt) +# Insert rows with conflict handling +try: + with engine.begin() as connection: + for row in rows: + stmt = sqlite_insert(city_stats_table).values(**row) + stmt = stmt.on_conflict_do_update( + index_elements=['city_name'], # Column(s) with the UNIQUE constraint + set_=row # Update the row with the new values if a conflict occurs + ) + connection.execute(stmt) + print("Database successfully populated.") +except Exception as e: + print(f"Error inserting data into database: {e}")
480-480
: Improve type hint fortool_calls
inGatherToolsEvent
class.The current type hint of
Any
doesn't provide clarity about the expected data structure.- tool_calls: Any + tool_calls: List[ToolSelection]
494-510
: Constructor needs improved parameter documentation.The constructor lacks docstrings for the parameters, making it difficult to understand their purpose and requirements.
def __init__(self, tools: List[BaseTool], timeout: Optional[float] = 10.0, disable_validation: bool = False, verbose: bool = False, llm: Optional[LLM] = None, chat_history: Optional[List[ChatMessage]] = None, ): - """Constructor.""" + """ + Constructor for RouterOutputAgentWorkflow. + + Args: + tools: List of tools available to the agent + timeout: Maximum time in seconds to wait for workflow steps to complete + disable_validation: If True, disables validation of workflow events + verbose: If True, prints detailed logs of operations + llm: Language model to use, defaults to OpenAI gpt-3.5-turbo if not provided + chat_history: Initial chat history, defaults to empty list if not provided + """
519-527
: Add error handling for missing input message.The current implementation raises an error when message is None, but it doesn't check for empty strings or invalid types.
@step() async def prepare_chat(self, ev: StartEvent) -> InputEvent: message = ev.get("message") - if message is None: - raise ValueError("'message' field is required.") + if not message or not isinstance(message, str): + raise ValueError("'message' field is required and must be a non-empty string.") # add msg to chat history chat_history = self.chat_history chat_history.append(ChatMessage(role="user", content=message)) return InputEvent()
594-607
: Improve gather method with validation.The gather method should validate the collected events before processing them.
@step(pass_context=True) async def gather(self, ctx: Context, ev: ToolCallEventResult) -> StopEvent | None: """Gathers tool calls.""" # wait for all tool call events to finish. tool_events = ctx.collect_events(ev, [ToolCallEventResult] * await ctx.get("num_tool_calls")) if not tool_events: return None + # Validate that all events have message content + invalid_events = [event for event in tool_events if not hasattr(event, 'msg') or not event.msg] + if invalid_events and self._verbose: + print(f"Warning: {len(invalid_events)} tool events returned without proper message content") + for tool_event in tool_events: + if not hasattr(tool_event, 'msg') or not tool_event.msg: + continue # append tool call chat messages to history self.chat_history.append(tool_event.msg) # # after all tool calls finish, pass input event back, restart agent loop return InputEvent()
622-622
: Consider making timeout configurable via environment variables.The timeout value is hardcoded, which limits flexibility in different environments.
-wf = RouterOutputAgentWorkflow(tools=[sql_tool, llama_cloud_tool], verbose=True, timeout=120) +# Get timeout from environment variables or use default +workflow_timeout = float(os.environ.get("WORKFLOW_TIMEOUT", "120")) +wf = RouterOutputAgentWorkflow(tools=[sql_tool, llama_cloud_tool], verbose=True, timeout=workflow_timeout)
" self.llm: LLM = llm or OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n", | ||
" self.chat_history: List[ChatMessage] = chat_history or []\n", |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Initialize LLM with proper import.
The constructor references OpenAI without proper importing, which could lead to runtime errors.
- self.llm: LLM = llm or OpenAI(temperature=0, model="gpt-3.5-turbo")
+ from llama_index.llms.openai import OpenAI
+ self.llm: LLM = llm or OpenAI(temperature=0, model="gpt-3.5-turbo")
📝 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.llm: LLM = llm or OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n", | |
" self.chat_history: List[ChatMessage] = chat_history or []\n", | |
" from llama_index.llms.openai import OpenAI\n", | |
" self.llm: LLM = llm or OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n", | |
" self.chat_history: List[ChatMessage] = chat_history or []\n", |
" async def call_tool(self, ev: ToolCallEvent) -> ToolCallEventResult:\n", | ||
" \"\"\"Calls tool.\"\"\"\n", | ||
"\n", | ||
" tool_call = ev.tool_call\n", | ||
"\n", | ||
" # get tool ID and function call\n", | ||
" id_ = tool_call.tool_id\n", | ||
"\n", | ||
" if self._verbose:\n", | ||
" print(f\"Calling function {tool_call.tool_name} with msg {tool_call.tool_kwargs}\")\n", | ||
"\n", | ||
" # call function and put result into a chat message\n", | ||
" tool = self.tools_dict[tool_call.tool_name]\n", | ||
" output = await tool.acall(**tool_call.tool_kwargs)\n", | ||
" msg = ChatMessage(\n", | ||
" name=tool_call.tool_name,\n", | ||
" content=str(output),\n", | ||
" role=\"tool\",\n", | ||
" additional_kwargs={\n", | ||
" \"tool_call_id\": id_,\n", | ||
" \"name\": tool_call.tool_name\n", | ||
" }\n", | ||
" )\n", | ||
"\n", | ||
" return ToolCallEventResult(msg=msg)\n", |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
Add error handling in call_tool method.
The current implementation doesn't handle exceptions that might occur during tool calls.
@step()
async def call_tool(self, ev: ToolCallEvent) -> ToolCallEventResult:
"""Calls tool."""
tool_call = ev.tool_call
# get tool ID and function call
id_ = tool_call.tool_id
if self._verbose:
print(f"Calling function {tool_call.tool_name} with msg {tool_call.tool_kwargs}")
# call function and put result into a chat message
tool = self.tools_dict[tool_call.tool_name]
- output = await tool.acall(**tool_call.tool_kwargs)
+ try:
+ output = await tool.acall(**tool_call.tool_kwargs)
+ except Exception as e:
+ error_msg = f"Error calling tool {tool_call.tool_name}: {str(e)}"
+ if self._verbose:
+ print(f"Tool call error: {error_msg}")
+ output = f"Failed to execute tool: {error_msg}"
📝 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.
" async def call_tool(self, ev: ToolCallEvent) -> ToolCallEventResult:\n", | |
" \"\"\"Calls tool.\"\"\"\n", | |
"\n", | |
" tool_call = ev.tool_call\n", | |
"\n", | |
" # get tool ID and function call\n", | |
" id_ = tool_call.tool_id\n", | |
"\n", | |
" if self._verbose:\n", | |
" print(f\"Calling function {tool_call.tool_name} with msg {tool_call.tool_kwargs}\")\n", | |
"\n", | |
" # call function and put result into a chat message\n", | |
" tool = self.tools_dict[tool_call.tool_name]\n", | |
" output = await tool.acall(**tool_call.tool_kwargs)\n", | |
" msg = ChatMessage(\n", | |
" name=tool_call.tool_name,\n", | |
" content=str(output),\n", | |
" role=\"tool\",\n", | |
" additional_kwargs={\n", | |
" \"tool_call_id\": id_,\n", | |
" \"name\": tool_call.tool_name\n", | |
" }\n", | |
" )\n", | |
"\n", | |
" return ToolCallEventResult(msg=msg)\n", | |
async def call_tool(self, ev: ToolCallEvent) -> ToolCallEventResult: | |
"""Calls tool.""" | |
tool_call = ev.tool_call | |
# get tool ID and function call | |
id_ = tool_call.tool_id | |
if self._verbose: | |
print(f"Calling function {tool_call.tool_name} with msg {tool_call.tool_kwargs}") | |
# call function and put result into a chat message | |
tool = self.tools_dict[tool_call.tool_name] | |
try: | |
output = await tool.acall(**tool_call.tool_kwargs) | |
except Exception as e: | |
error_msg = f"Error calling tool {tool_call.tool_name}: {str(e)}" | |
if self._verbose: | |
print(f"Tool call error: {error_msg}") | |
output = f"Failed to execute tool: {error_msg}" | |
msg = ChatMessage( | |
name=tool_call.tool_name, | |
content=str(output), | |
role="tool", | |
additional_kwargs={ | |
"tool_call_id": id_, | |
"name": tool_call.tool_name | |
} | |
) | |
return ToolCallEventResult(msg=msg) |
" async def chat(self, ev: InputEvent) -> GatherToolsEvent | StopEvent:\n", | ||
" \"\"\"Appends msg to chat history, then gets tool calls.\"\"\"\n", | ||
"\n", | ||
" # Put msg into LLM with tools included\n", | ||
" chat_res = await self.llm.achat_with_tools(\n", | ||
" self.tools,\n", | ||
" chat_history=self.chat_history,\n", | ||
" verbose=self._verbose,\n", | ||
" allow_parallel_tool_calls=True\n", | ||
" )\n", | ||
" tool_calls = self.llm.get_tool_calls_from_response(chat_res, error_on_no_tool_call=False)\n", | ||
" \n", | ||
" ai_message = chat_res.message\n", | ||
" self.chat_history.append(ai_message)\n", | ||
" if self._verbose:\n", | ||
" print(f\"Chat message: {ai_message.content}\")\n", | ||
"\n", | ||
" # no tool calls, return chat message.\n", | ||
" if not tool_calls:\n", | ||
" return StopEvent(result=ai_message.content)\n", | ||
"\n", | ||
" return GatherToolsEvent(tool_calls=tool_calls)\n", |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
The chat method should handle LLM API errors gracefully.
The method doesn't include error handling for failed API calls, which could cause the workflow to fail unexpectedly.
@step()
async def chat(self, ev: InputEvent) -> GatherToolsEvent | StopEvent:
"""Appends msg to chat history, then gets tool calls."""
# Put msg into LLM with tools included
- chat_res = await self.llm.achat_with_tools(
- self.tools,
- chat_history=self.chat_history,
- verbose=self._verbose,
- allow_parallel_tool_calls=True
- )
+ try:
+ chat_res = await self.llm.achat_with_tools(
+ self.tools,
+ chat_history=self.chat_history,
+ verbose=self._verbose,
+ allow_parallel_tool_calls=True
+ )
+ except Exception as e:
+ error_message = f"Error calling LLM API: {str(e)}"
+ if self._verbose:
+ print(error_message)
+ # Add system message about the error
+ self.chat_history.append(ChatMessage(role="system", content=error_message))
+ return StopEvent(result=f"I encountered an error: {error_message}")
📝 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.
" async def chat(self, ev: InputEvent) -> GatherToolsEvent | StopEvent:\n", | |
" \"\"\"Appends msg to chat history, then gets tool calls.\"\"\"\n", | |
"\n", | |
" # Put msg into LLM with tools included\n", | |
" chat_res = await self.llm.achat_with_tools(\n", | |
" self.tools,\n", | |
" chat_history=self.chat_history,\n", | |
" verbose=self._verbose,\n", | |
" allow_parallel_tool_calls=True\n", | |
" )\n", | |
" tool_calls = self.llm.get_tool_calls_from_response(chat_res, error_on_no_tool_call=False)\n", | |
" \n", | |
" ai_message = chat_res.message\n", | |
" self.chat_history.append(ai_message)\n", | |
" if self._verbose:\n", | |
" print(f\"Chat message: {ai_message.content}\")\n", | |
"\n", | |
" # no tool calls, return chat message.\n", | |
" if not tool_calls:\n", | |
" return StopEvent(result=ai_message.content)\n", | |
"\n", | |
" return GatherToolsEvent(tool_calls=tool_calls)\n", | |
async def chat(self, ev: InputEvent) -> GatherToolsEvent | StopEvent: | |
"""Appends msg to chat history, then gets tool calls.""" | |
# Put msg into LLM with tools included | |
try: | |
chat_res = await self.llm.achat_with_tools( | |
self.tools, | |
chat_history=self.chat_history, | |
verbose=self._verbose, | |
allow_parallel_tool_calls=True | |
) | |
except Exception as e: | |
error_message = f"Error calling LLM API: {str(e)}" | |
if self._verbose: | |
print(error_message) | |
# Add system message about the error | |
self.chat_history.append(ChatMessage(role="system", content=error_message)) | |
return StopEvent(result=f"I encountered an error: {error_message}") | |
tool_calls = self.llm.get_tool_calls_from_response(chat_res, error_on_no_tool_call=False) | |
ai_message = chat_res.message | |
self.chat_history.append(ai_message) | |
if self._verbose: | |
print(f"Chat message: {ai_message.content}") | |
# no tool calls, return chat message. | |
if not tool_calls: | |
return StopEvent(result=ai_message.content) | |
return GatherToolsEvent(tool_calls=tool_calls) |
" name=\"<Your Index Name>\", \n", | ||
" project_name=\"<Your Project Name>\",\n", | ||
" organization_id=\"<Your Org ID>\",\n", | ||
" api_key=\"<Your API Key>\"\n", |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Avoid hardcoding LlamaCloud credentials in code.
Similar to the API key issue earlier, hardcoded credentials are a security risk.
-index = LlamaCloudIndex(
- name="<Your Index Name>",
- project_name="<Your Project Name>",
- organization_id="<Your Org ID>",
- api_key="<Your API Key>"
-)
+# Load LlamaCloud configuration from environment variables
+llama_index_name = os.environ.get("LLAMA_INDEX_NAME", "")
+llama_project_name = os.environ.get("LLAMA_PROJECT_NAME", "")
+llama_org_id = os.environ.get("LLAMA_ORG_ID", "")
+llama_api_key = os.environ.get("LLAMA_API_KEY", "")
+
+# Validate required configuration
+if not all([llama_index_name, llama_project_name, llama_org_id, llama_api_key]):
+ raise ValueError("LlamaCloud environment variables are not properly configured")
+
+index = LlamaCloudIndex(
+ name=llama_index_name,
+ project_name=llama_project_name,
+ organization_id=llama_org_id,
+ api_key=llama_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.
" name=\"<Your Index Name>\", \n", | |
" project_name=\"<Your Project Name>\",\n", | |
" organization_id=\"<Your Org ID>\",\n", | |
" api_key=\"<Your API Key>\"\n", | |
# Load LlamaCloud configuration from environment variables | |
llama_index_name = os.environ.get("LLAMA_INDEX_NAME", "") | |
llama_project_name = os.environ.get("LLAMA_PROJECT_NAME", "") | |
llama_org_id = os.environ.get("LLAMA_ORG_ID", "") | |
llama_api_key = os.environ.get("LLAMA_API_KEY", "") | |
# Validate required configuration | |
if not all([llama_index_name, llama_project_name, llama_org_id, llama_api_key]): | |
raise ValueError("LlamaCloud environment variables are not properly configured") | |
index = LlamaCloudIndex( | |
name=llama_index_name, | |
project_name=llama_project_name, | |
organization_id=llama_org_id, | |
api_key=llama_api_key | |
) |
This PR implements a unified interface combining Retrieval-Augmented Generation (RAG) and Text-to-SQL for querying structured and unstructured data. Users can input natural language queries, and the system intelligently uses RAG for context-aware responses or Text-to-SQL for database queries. Key features include a retrieval mechanism, SQL query generation, and a single query interface. Added modules, updated documentation, and provided testing instructions. Future improvements include optimizing retrieval speed and expanding the knowledge base. Submitted as part of the DailyDoseofDS Technical Writer Task. Feedback is welcome! 🚀
Summary by CodeRabbit
Documentation
New Features