vectara-agentic 0.2.9__tar.gz → 0.2.11__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of vectara-agentic might be problematic. Click here for more details.

Files changed (36) hide show
  1. {vectara_agentic-0.2.9/vectara_agentic.egg-info → vectara_agentic-0.2.11}/PKG-INFO +12 -12
  2. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/README.md +3 -3
  3. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/requirements.txt +8 -8
  4. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/tests/test_tools.py +64 -13
  5. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/_callback.py +2 -2
  6. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/_observability.py +3 -2
  7. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/_prompts.py +5 -4
  8. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/_version.py +1 -1
  9. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/agent.py +1 -1
  10. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/db_tools.py +34 -14
  11. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/sub_query_workflow.py +82 -45
  12. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/tools.py +148 -60
  13. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11/vectara_agentic.egg-info}/PKG-INFO +12 -12
  14. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic.egg-info/requires.txt +8 -8
  15. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/LICENSE +0 -0
  16. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/MANIFEST.in +0 -0
  17. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/setup.cfg +0 -0
  18. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/setup.py +0 -0
  19. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/tests/__init__.py +0 -0
  20. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/tests/endpoint.py +0 -0
  21. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/tests/test_agent.py +0 -0
  22. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/tests/test_agent_planning.py +0 -0
  23. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/tests/test_agent_type.py +0 -0
  24. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/tests/test_fallback.py +0 -0
  25. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/tests/test_private_llm.py +0 -0
  26. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/tests/test_serialization.py +0 -0
  27. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/tests/test_workflow.py +0 -0
  28. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/__init__.py +0 -0
  29. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/agent_config.py +0 -0
  30. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/agent_endpoint.py +0 -0
  31. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/tools_catalog.py +0 -0
  32. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/types.py +0 -0
  33. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic/utils.py +0 -0
  34. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic.egg-info/SOURCES.txt +0 -0
  35. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic.egg-info/dependency_links.txt +0 -0
  36. {vectara_agentic-0.2.9 → vectara_agentic-0.2.11}/vectara_agentic.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: vectara_agentic
3
- Version: 0.2.9
3
+ Version: 0.2.11
4
4
  Summary: A Python package for creating AI Assistants and AI Agents with Vectara
5
5
  Home-page: https://github.com/vectara/py-vectara-agentic
6
6
  Author: Ofer Mendelevitch
@@ -16,31 +16,31 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
16
16
  Requires-Python: >=3.10
17
17
  Description-Content-Type: text/markdown
18
18
  License-File: LICENSE
19
- Requires-Dist: llama-index==0.12.26
19
+ Requires-Dist: llama-index==0.12.29
20
20
  Requires-Dist: llama-index-indices-managed-vectara==0.4.2
21
21
  Requires-Dist: llama-index-agent-llm-compiler==0.3.0
22
22
  Requires-Dist: llama-index-agent-lats==0.3.0
23
23
  Requires-Dist: llama-index-agent-openai==0.4.6
24
- Requires-Dist: llama-index-llms-openai==0.3.25
24
+ Requires-Dist: llama-index-llms-openai==0.3.32
25
25
  Requires-Dist: llama-index-llms-anthropic==0.6.10
26
26
  Requires-Dist: llama-index-llms-together==0.3.1
27
27
  Requires-Dist: llama-index-llms-groq==0.3.1
28
28
  Requires-Dist: llama-index-llms-fireworks==0.3.2
29
29
  Requires-Dist: llama-index-llms-cohere==0.4.0
30
- Requires-Dist: llama-index-llms-gemini==0.4.11
31
- Requires-Dist: llama-index-llms-bedrock==0.3.4
30
+ Requires-Dist: llama-index-llms-gemini==0.4.14
31
+ Requires-Dist: llama-index-llms-bedrock==0.3.8
32
32
  Requires-Dist: llama-index-tools-yahoo-finance==0.3.0
33
33
  Requires-Dist: llama-index-tools-arxiv==0.3.0
34
34
  Requires-Dist: llama-index-tools-database==0.3.0
35
35
  Requires-Dist: llama-index-tools-google==0.3.0
36
36
  Requires-Dist: llama-index-tools-tavily_research==0.3.0
37
+ Requires-Dist: llama_index.tools.brave_search==0.3.0
37
38
  Requires-Dist: llama-index-tools-neo4j==0.3.0
38
- Requires-Dist: llama-index-graph-stores-kuzu==0.6.0
39
+ Requires-Dist: llama-index-graph-stores-kuzu==0.7.0
39
40
  Requires-Dist: llama-index-tools-slack==0.3.0
40
41
  Requires-Dist: llama-index-tools-exa==0.3.0
41
- Requires-Dist: tavily-python==0.5.1
42
- Requires-Dist: exa-py==1.8.9
43
- Requires-Dist: yahoo-finance==1.4.0
42
+ Requires-Dist: tavily-python==0.5.4
43
+ Requires-Dist: exa-py==1.9.1
44
44
  Requires-Dist: openinference-instrumentation-llama-index==3.3.3
45
45
  Requires-Dist: opentelemetry-proto==1.31.0
46
46
  Requires-Dist: arize-phoenix==8.14.1
@@ -176,7 +176,7 @@ query_financial_reports_tool = vec_factory.create_rag_tool(
176
176
  )
177
177
  ```
178
178
 
179
- See the [docs](https://vectara.github.io/vectara-agentic-docs/) for additional arguments to customize your Vectara RAG tool.
179
+ See the [docs](https://vectara.github.io/py-vectara-agentic/latest/) for additional arguments to customize your Vectara RAG tool.
180
180
 
181
181
  ### 3. Create other tools (optional)
182
182
 
@@ -205,7 +205,7 @@ agent = Agent(
205
205
  )
206
206
  ```
207
207
 
208
- See the [docs](https://vectara.github.io/vectara-agentic-docs/) for additional arguments, including `agent_progress_callback` and `query_logging_callback`.
208
+ See the [docs](https://vectara.github.io/py-vectara-agentic/latest/) for additional arguments, including `agent_progress_callback` and `query_logging_callback`.
209
209
 
210
210
  ### 5. Run a chat interaction
211
211
 
@@ -376,7 +376,7 @@ specified in the Agent configuration.
376
376
  - `load_unique_values`: returns the top unique values for a given column
377
377
 
378
378
  In addition, we include various other tools from LlamaIndex ToolSpecs:
379
- * Tavily search and EXA.AI
379
+ * Tavily search, EXA.AI and Brave Search
380
380
  * arxiv
381
381
  * neo4j & Kuzu for Graph DB integration
382
382
  * Google tools (including gmail, calendar, and search)
@@ -108,7 +108,7 @@ query_financial_reports_tool = vec_factory.create_rag_tool(
108
108
  )
109
109
  ```
110
110
 
111
- See the [docs](https://vectara.github.io/vectara-agentic-docs/) for additional arguments to customize your Vectara RAG tool.
111
+ See the [docs](https://vectara.github.io/py-vectara-agentic/latest/) for additional arguments to customize your Vectara RAG tool.
112
112
 
113
113
  ### 3. Create other tools (optional)
114
114
 
@@ -137,7 +137,7 @@ agent = Agent(
137
137
  )
138
138
  ```
139
139
 
140
- See the [docs](https://vectara.github.io/vectara-agentic-docs/) for additional arguments, including `agent_progress_callback` and `query_logging_callback`.
140
+ See the [docs](https://vectara.github.io/py-vectara-agentic/latest/) for additional arguments, including `agent_progress_callback` and `query_logging_callback`.
141
141
 
142
142
  ### 5. Run a chat interaction
143
143
 
@@ -308,7 +308,7 @@ specified in the Agent configuration.
308
308
  - `load_unique_values`: returns the top unique values for a given column
309
309
 
310
310
  In addition, we include various other tools from LlamaIndex ToolSpecs:
311
- * Tavily search and EXA.AI
311
+ * Tavily search, EXA.AI and Brave Search
312
312
  * arxiv
313
313
  * neo4j & Kuzu for Graph DB integration
314
314
  * Google tools (including gmail, calendar, and search)
@@ -1,28 +1,28 @@
1
- llama-index==0.12.26
1
+ llama-index==0.12.29
2
2
  llama-index-indices-managed-vectara==0.4.2
3
3
  llama-index-agent-llm-compiler==0.3.0
4
4
  llama-index-agent-lats==0.3.0
5
5
  llama-index-agent-openai==0.4.6
6
- llama-index-llms-openai==0.3.25
6
+ llama-index-llms-openai==0.3.32
7
7
  llama-index-llms-anthropic==0.6.10
8
8
  llama-index-llms-together==0.3.1
9
9
  llama-index-llms-groq==0.3.1
10
10
  llama-index-llms-fireworks==0.3.2
11
11
  llama-index-llms-cohere==0.4.0
12
- llama-index-llms-gemini==0.4.11
13
- llama-index-llms-bedrock==0.3.4
12
+ llama-index-llms-gemini==0.4.14
13
+ llama-index-llms-bedrock==0.3.8
14
14
  llama-index-tools-yahoo-finance==0.3.0
15
15
  llama-index-tools-arxiv==0.3.0
16
16
  llama-index-tools-database==0.3.0
17
17
  llama-index-tools-google==0.3.0
18
18
  llama-index-tools-tavily_research==0.3.0
19
+ llama_index.tools.brave_search==0.3.0
19
20
  llama-index-tools-neo4j==0.3.0
20
- llama-index-graph-stores-kuzu==0.6.0
21
+ llama-index-graph-stores-kuzu==0.7.0
21
22
  llama-index-tools-slack==0.3.0
22
23
  llama-index-tools-exa==0.3.0
23
- tavily-python==0.5.1
24
- exa-py==1.8.9
25
- yahoo-finance==1.4.0
24
+ tavily-python==0.5.4
25
+ exa-py==1.9.1
26
26
  openinference-instrumentation-llama-index==3.3.3
27
27
  opentelemetry-proto==1.31.0
28
28
  arize-phoenix==8.14.1
@@ -2,7 +2,12 @@ import unittest
2
2
 
3
3
  from pydantic import Field, BaseModel
4
4
 
5
- from vectara_agentic.tools import VectaraTool, VectaraToolFactory, ToolsFactory, ToolType
5
+ from vectara_agentic.tools import (
6
+ VectaraTool,
7
+ VectaraToolFactory,
8
+ ToolsFactory,
9
+ ToolType,
10
+ )
6
11
  from vectara_agentic.agent import Agent
7
12
  from vectara_agentic.agent_config import AgentConfig
8
13
 
@@ -13,9 +18,7 @@ class TestToolsPackage(unittest.TestCase):
13
18
  def test_vectara_tool_factory(self):
14
19
  vectara_corpus_key = "corpus_key"
15
20
  vectara_api_key = "api_key"
16
- vec_factory = VectaraToolFactory(
17
- vectara_corpus_key, vectara_api_key
18
- )
21
+ vec_factory = VectaraToolFactory(vectara_corpus_key, vectara_api_key)
19
22
 
20
23
  self.assertEqual(vectara_corpus_key, vec_factory.vectara_corpus_key)
21
24
  self.assertEqual(vectara_api_key, vec_factory.vectara_api_key)
@@ -46,6 +49,48 @@ class TestToolsPackage(unittest.TestCase):
46
49
  self.assertIsInstance(search_tool, FunctionTool)
47
50
  self.assertEqual(search_tool.metadata.tool_type, ToolType.QUERY)
48
51
 
52
+ def test_vectara_tool_validation(self):
53
+ vectara_corpus_key = "corpus_key"
54
+ vectara_api_key = "api_key"
55
+ vec_factory = VectaraToolFactory(vectara_corpus_key, vectara_api_key)
56
+
57
+ class QueryToolArgs(BaseModel):
58
+ query: str = Field(description="The user query")
59
+ year: int = Field(
60
+ description="The year of the document",
61
+ example=2023,
62
+ )
63
+ ticker: str = Field(
64
+ description="The stock ticker",
65
+ example="AAPL",
66
+ )
67
+
68
+ query_tool = vec_factory.create_rag_tool(
69
+ tool_name="rag_tool",
70
+ tool_description="""
71
+ Returns a response (str) to the user query based on the data in this corpus.
72
+ """,
73
+ tool_args_schema=QueryToolArgs,
74
+ )
75
+ res = query_tool(
76
+ query="What is the stock price?",
77
+ the_year=2023,
78
+ )
79
+ self.assertIn("Malfunction", str(res))
80
+
81
+ search_tool = vec_factory.create_search_tool(
82
+ tool_name="search_tool",
83
+ tool_description="""
84
+ Returns a list of documents (str) that match the user query.
85
+ """,
86
+ tool_args_schema=QueryToolArgs,
87
+ )
88
+ res = search_tool(
89
+ query="What is the stock price?",
90
+ the_year=2023,
91
+ )
92
+ self.assertIn("Malfunction", str(res))
93
+
49
94
  def test_tool_factory(self):
50
95
  def mult(x: float, y: float) -> float:
51
96
  return x * y
@@ -59,17 +104,21 @@ class TestToolsPackage(unittest.TestCase):
59
104
  def test_llama_index_tools(self):
60
105
  tools_factory = ToolsFactory()
61
106
 
62
- llama_tools = tools_factory.get_llama_index_tools(
63
- tool_package_name="arxiv",
64
- tool_spec_name="ArxivToolSpec"
65
- )
66
-
67
- arxiv_tool = llama_tools[0]
107
+ arxiv_tool = tools_factory.get_llama_index_tools(
108
+ tool_package_name="arxiv", tool_spec_name="ArxivToolSpec"
109
+ )[0]
68
110
 
69
111
  self.assertIsInstance(arxiv_tool, VectaraTool)
70
112
  self.assertIsInstance(arxiv_tool, FunctionTool)
71
113
  self.assertEqual(arxiv_tool.metadata.tool_type, ToolType.QUERY)
72
114
 
115
+ yfinance_tool = tools_factory.get_llama_index_tools(
116
+ tool_package_name="yahoo_finance", tool_spec_name="YahooFinanceToolSpec"
117
+ )[0]
118
+ self.assertIsInstance(yfinance_tool, VectaraTool)
119
+ self.assertIsInstance(yfinance_tool, FunctionTool)
120
+ self.assertEqual(yfinance_tool.metadata.tool_type, ToolType.QUERY)
121
+
73
122
  def test_public_repo(self):
74
123
  vectara_corpus_key = "vectara-docs_1"
75
124
  vectara_api_key = "zqt_UXrBcnI2UXINZkrv4g1tQPhzj02vfdtqYJIDiA"
@@ -80,10 +129,12 @@ class TestToolsPackage(unittest.TestCase):
80
129
  tool_name="ask_vectara",
81
130
  data_description="data from Vectara website",
82
131
  assistant_specialty="RAG as a service",
83
- vectara_summarizer="mockingbird-1.0-2024-07-16"
132
+ vectara_summarizer="mockingbird-1.0-2024-07-16",
84
133
  )
85
134
 
86
- self.assertIn("Vectara is an end-to-end platform", str(agent.chat("What is Vectara?")))
135
+ self.assertIn(
136
+ "Vectara is an end-to-end platform", str(agent.chat("What is Vectara?"))
137
+ )
87
138
 
88
139
  def test_class_method_as_tool(self):
89
140
  class TestClass:
@@ -102,7 +153,7 @@ class TestToolsPackage(unittest.TestCase):
102
153
  tools=tools,
103
154
  topic=topic,
104
155
  custom_instructions=instructions,
105
- agent_config=config
156
+ agent_config=config,
106
157
  )
107
158
 
108
159
  self.assertEqual(
@@ -154,7 +154,7 @@ class AgentCallbackHandler(BaseCallbackHandler):
154
154
  elif event_type == CBEventType.AGENT_STEP:
155
155
  self._handle_agent_step(payload, event_id)
156
156
  elif event_type == CBEventType.EXCEPTION:
157
- print(f"Exception: {payload.get(EventPayload.EXCEPTION)}")
157
+ print(f"Exception in handle_event: {payload.get(EventPayload.EXCEPTION)}")
158
158
  else:
159
159
  print(f"Unknown event type: {event_type}, payload={payload}")
160
160
 
@@ -168,7 +168,7 @@ class AgentCallbackHandler(BaseCallbackHandler):
168
168
  elif event_type == CBEventType.AGENT_STEP:
169
169
  await self._ahandle_agent_step(payload, event_id)
170
170
  elif event_type == CBEventType.EXCEPTION:
171
- print(f"Exception: {payload.get(EventPayload.EXCEPTION)}")
171
+ print(f"Exception in ahandle_event: {payload.get(EventPayload.EXCEPTION)}")
172
172
  else:
173
173
  print(f"Unknown event type: {event_type}, payload={payload}")
174
174
 
@@ -8,7 +8,7 @@ import pandas as pd
8
8
  from .types import ObserverType
9
9
  from .agent_config import AgentConfig
10
10
 
11
- def setup_observer(config: AgentConfig) -> bool:
11
+ def setup_observer(config: AgentConfig, verbose: bool) -> bool:
12
12
  '''
13
13
  Setup the observer.
14
14
  '''
@@ -31,7 +31,8 @@ def setup_observer(config: AgentConfig) -> bool:
31
31
  tracer_provider = register(endpoint=phoenix_endpoint, project_name="vectara-agentic")
32
32
  LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
33
33
  return True
34
- print("No observer set.")
34
+ if verbose:
35
+ print("No observer set.")
35
36
  return False
36
37
 
37
38
 
@@ -37,10 +37,11 @@ GENERAL_INSTRUCTIONS = """
37
37
  - If including latex equations in the markdown response, make sure the equations are on a separate line and enclosed in double dollar signs.
38
38
  - Always respond in the language of the question, and in text (no images, videos or code).
39
39
  - If you are provided with database tools use them for analytical queries (such as counting, calculating max, min, average, sum, or other statistics).
40
- For each database, the database tools include: x_list_tables, x_load_data, x_describe_tables, and x_load_sample_data, where 'x' in the database name.
41
- The x_list_tables tool provides a list of available tables in the x database. Always use x_list_tables before using other database tools, to understand valid table names.
42
- Before using the x_load_data with a SQL query, always follow these steps:
43
- - Use the x_describe_tables tool to understand the schema of each table.
40
+ For each database, the database tools include: x_list_tables, x_load_data, x_describe_tables, x_load_unique_values, and x_load_sample_data, where 'x' in the database name.
41
+ for example, if the database name is "ev", the tools are: ev_list_tables, ev_load_data, ev_describe_tables, ev_load_unique_values, and ev_load_sample_data.
42
+ Before using the x_load_data with a SQL query, always follow these discovery steps:
43
+ - call the x_list_tables tool to list of available tables in the x database.
44
+ - Call the x_describe_tables tool to understand the schema of each table you want to query data from.
44
45
  - Use the x_load_unique_values tool to understand the unique values in each column.
45
46
  Sometimes the user may ask for a specific column value, but the actual value in the table may be different, and you will need to use the correct value.
46
47
  - Use the x_load_sample_data tool to understand the column names, and typical values in each column.
@@ -1,4 +1,4 @@
1
1
  """
2
2
  Define the version of the package.
3
3
  """
4
- __version__ = "0.2.9"
4
+ __version__ = "0.2.11"
@@ -252,7 +252,7 @@ class Agent:
252
252
 
253
253
  # Setup observability
254
254
  try:
255
- self.observability_enabled = setup_observer(self.agent_config)
255
+ self.observability_enabled = setup_observer(self.agent_config, self.verbose)
256
256
  except Exception as e:
257
257
  print(f"Failed to set up observer ({e}), ignoring")
258
258
  self.observability_enabled = False
@@ -15,7 +15,6 @@ from sqlalchemy.engine import Engine
15
15
  from sqlalchemy.exc import NoSuchTableError
16
16
  from sqlalchemy.schema import CreateTable
17
17
 
18
- from llama_index.core.readers.base import BaseReader
19
18
  from llama_index.core.utilities.sql_wrapper import SQLDatabase
20
19
  from llama_index.core.schema import Document
21
20
  from llama_index.core.tools.function_tool import FunctionTool
@@ -25,7 +24,7 @@ from llama_index.core.tools.utils import create_schema_from_function
25
24
 
26
25
  AsyncCallable = Callable[..., Awaitable[Any]]
27
26
 
28
- class DatabaseTools(BaseReader):
27
+ class DatabaseTools:
29
28
  """Database tools for vectara-agentic
30
29
  This class provides a set of tools to interact with a database.
31
30
  It allows you to load data, list tables, describe tables, and load unique values.
@@ -49,9 +48,11 @@ class DatabaseTools(BaseReader):
49
48
  user: Optional[str] = None,
50
49
  password: Optional[str] = None,
51
50
  dbname: Optional[str] = None,
51
+ tool_name_prefix: str = "db",
52
52
  **kwargs: Any,
53
53
  ) -> None:
54
54
  self.max_rows = max_rows
55
+ self.tool_name_prefix = tool_name_prefix
55
56
 
56
57
  if sql_database:
57
58
  self.sql_database = sql_database
@@ -86,35 +87,34 @@ class DatabaseTools(BaseReader):
86
87
  func = getattr(self, fn_name)
87
88
  except AttributeError:
88
89
  return None
89
- name = fn_name
90
+ name = self.tool_name_prefix + "_" + fn_name if self.tool_name_prefix else fn_name
90
91
  docstring = func.__doc__ or ""
91
92
  description = f"{name}{signature(func)}\n{docstring}"
92
93
  fn_schema = create_schema_from_function(fn_name, getattr(self, fn_name))
93
94
  return ToolMetadata(name=name, description=description, fn_schema=fn_schema)
94
95
 
95
- def _load_data(self, query: str) -> List[Document]:
96
+ def _load_data(self, sql_query: str) -> List[Document]:
96
97
  documents = []
97
98
  with self.sql_database.engine.connect() as connection:
98
- if query is None:
99
+ if sql_query is None:
99
100
  raise ValueError("A query parameter is necessary to filter the data")
100
- result = connection.execute(text(query))
101
+ result = connection.execute(text(sql_query))
101
102
  for item in result.fetchall():
102
103
  doc_str = ", ".join([str(entry) for entry in item])
103
104
  documents.append(Document(text=doc_str))
104
105
  return documents
105
106
 
106
- def load_data(self, *args: Any, **load_kwargs: Any) -> List[str]:
107
+ def load_data(self, sql_query: str) -> List[str]:
107
108
  """Query and load data from the Database, returning a list of Documents.
108
109
  Args:
109
- query (str): an SQL query to filter tables and rows.
110
+ sql_query (str): an SQL query to filter tables and rows.
110
111
  Returns:
111
- List[Document]: a list of Document objects from the database.
112
+ List[str]: a list of Document objects from the database.
112
113
  """
113
- query = args[0] if args else load_kwargs.get("args",{}).get("query")
114
- if query is None:
114
+ if sql_query is None:
115
115
  raise ValueError("A query parameter is necessary to filter the data")
116
116
 
117
- count_query = f"SELECT COUNT(*) FROM ({query})"
117
+ count_query = f"SELECT COUNT(*) FROM ({sql_query})"
118
118
  try:
119
119
  count_rows = self._load_data(count_query)
120
120
  except Exception as e:
@@ -126,9 +126,9 @@ class DatabaseTools(BaseReader):
126
126
  "Please refactor your query to make it return less rows. "
127
127
  ]
128
128
  try:
129
- res = self._load_data(query)
129
+ res = self._load_data(sql_query)
130
130
  except Exception as e:
131
- return [f"Error ({str(e)}) occurred while executing the query {query}"]
131
+ return [f"Error ({str(e)}) occurred while executing the query {sql_query}"]
132
132
  return [d.text for d in res]
133
133
 
134
134
  def load_sample_data(self, table_name: str, num_rows: int = 25) -> Any:
@@ -141,6 +141,11 @@ class DatabaseTools(BaseReader):
141
141
  Returns:
142
142
  Any: The result of the database query.
143
143
  """
144
+ if table_name not in self.list_tables():
145
+ return (
146
+ f"Table {table_name} does not exist in the database."
147
+ f"Valid table names are: {self.list_tables()}"
148
+ )
144
149
  try:
145
150
  res = self._load_data(f"SELECT * FROM {table_name} LIMIT {num_rows}")
146
151
  except Exception as e:
@@ -162,6 +167,15 @@ class DatabaseTools(BaseReader):
162
167
  str: A string representation of the table schemas.
163
168
  """
164
169
  table_names = tables or [table.name for table in self._metadata.sorted_tables]
170
+ if len(table_names) == 0:
171
+ return "You must specify at least one table name to describe."
172
+ for table_name in table_names:
173
+ if table_name not in self.list_tables():
174
+ return (
175
+ f"Table {table_name} does not exist in the database."
176
+ f"Valid table names are: {self.list_tables()}"
177
+ )
178
+
165
179
  table_schemas = []
166
180
  for table_name in table_names:
167
181
  table = next(
@@ -186,6 +200,12 @@ class DatabaseTools(BaseReader):
186
200
  Returns:
187
201
  Any: the result of the database query
188
202
  """
203
+ if table_name not in self.list_tables():
204
+ return (
205
+ f"Table {table_name} does not exist in the database."
206
+ f"Valid table names are: {self.list_tables()}"
207
+ )
208
+
189
209
  res = {}
190
210
  try:
191
211
  for column in columns: