Skip to content

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

Open
wants to merge 3 commits into
base: main
Choose a base branch
from

Conversation

dmohammadullah
Copy link

@dmohammadullah dmohammadullah commented Mar 14, 2025

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

    • Introduced a comprehensive guide with installation and setup instructions for the unified query interface.
  • New Features

    • Launched an interactive application that lets users perform natural language queries on city statistics.
    • Added a demonstration notebook showcasing advanced query capabilities that combine search and SQL interactions.

Copy link

coderabbitai bot commented Mar 14, 2025

Walkthrough

This 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 app.py) enables natural language queries on an in-memory SQL database of city statistics. Additionally, a Jupyter notebook integrates the RAG and Text-to-SQL workflows with custom agent logic to orchestrate calls between a SQL query engine and a LlamaCloud index, complete with event-driven processing steps.

Changes

File(s) Change Summary
RAG and Text-to-SQL/README.md Added a new README detailing project components, installation steps, API key configuration, dependency installation via pip, and instructions for running the Streamlit application.
RAG and Text-to-SQL/app.py Introduced a Streamlit application for querying an in-memory SQLite database containing city statistics, including new functions (main()) and classes (SQLDatabase, NLSQLTableQueryEngine, QueryEngineTool). The app configures OpenAI API for the "gpt-3.5-turbo" model and uses error handling and asynchronous support via nest_asyncio.
RAG and Text-to-SQL/llamacloud_sql_router.ipynb Added a Jupyter notebook that integrates RAG with a Text-to-SQL interface. New classes (RouterOutputAgentWorkflow, InputEvent, GatherToolsEvent, ToolCallEvent, ToolCallEventResult) and methods (reset, prepare_chat, chat, dispatch_calls, call_tool, gather) were created to implement an agent workflow orchestrating interactions between SQL and LlamaCloud query engines with event-driven processing.

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
Loading
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
Loading

Poem

I hopped through lines of code so neat,
With SQL stats and queries sweet,
Streamlit shone, and tools aligned,
In a dance where data's defined.
A rabbit’s cheer for every new feat—
Hopping on code, a joyful treat!

Tip

⚡🧪 Multi-step agentic review comment chat (experimental)
  • We're introducing multi-step agentic chat in review comments. This experimental feature enhances review discussions with the CodeRabbit agentic chat by enabling advanced interactions, including the ability to create pull requests directly from comments.
    - To enable this feature, set early_access to true under in the settings.

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Generate unit testing code for this file.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
    • @coderabbitai generate unit testing code for this file.
    • @coderabbitai modularize this function.
  • PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read src/utils.ts and generate unit testing code.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
    • @coderabbitai help me debug CodeRabbit configuration file.

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)

  • @coderabbitai pause to pause the reviews on a PR.
  • @coderabbitai resume to resume the paused reviews.
  • @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
  • @coderabbitai full review to do a full review from scratch and review all the files again.
  • @coderabbitai summary to regenerate the summary of the PR.
  • @coderabbitai generate docstrings to generate docstrings for this PR.
  • @coderabbitai resolve resolve all the CodeRabbit review comments.
  • @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
  • @coderabbitai help to get help.

Other keywords and placeholders

  • Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

  • You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
  • Please see the configuration documentation for more information.
  • If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

Copy link

@coderabbitai coderabbitai bot left a 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:

  1. Updating your environment or compiler to meet these requirements.
  2. 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

📥 Commits

Reviewing files that changed from the base of the PR and between 6140fc3 and b6806a9.

📒 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)

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>"
Copy link

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 Streamlit st.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.

Suggested change
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.

Copy link

@coderabbitai coderabbitai bot left a 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

📥 Commits

Reviewing files that changed from the base of the PR and between b6806a9 and f1c67a2.

📒 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.

Comment on lines +64 to +74
# 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)
Copy link

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.

Suggested change
# 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.")

Copy link

@coderabbitai coderabbitai bot left a 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 issue

Avoid 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 for tool_calls in GatherToolsEvent 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)
📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between f1c67a2 and d75d6f8.

📒 Files selected for processing (1)
  • RAG and Text-to-SQL/llamacloud_sql_router.ipynb (1 hunks)

Comment on lines +509 to +510
" self.llm: LLM = llm or OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n",
" self.chat_history: List[ChatMessage] = chat_history or []\n",
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

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.

Suggested change
" 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",

Comment on lines +567 to +591
" 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",
Copy link

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.

Suggested change
" 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)

Comment on lines +530 to +551
" 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",
Copy link

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.

Suggested change
" 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)

Comment on lines +388 to +391
" name=\"<Your Index Name>\", \n",
" project_name=\"<Your Project Name>\",\n",
" organization_id=\"<Your Org ID>\",\n",
" api_key=\"<Your API Key>\"\n",
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

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.

Suggested change
" 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
)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant