Skip to content

Commit fd2db17

Browse files
authored
feat: add llamaindex cassandra tool (#439)
1 parent 66d8689 commit fd2db17

File tree

2 files changed

+57
-0
lines changed

2 files changed

+57
-0
lines changed
Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
import uuid
2+
3+
import cassio
4+
5+
from llama_index.agent.openai import OpenAIAgent
6+
from llama_index.llms.openai import OpenAI
7+
8+
from llama_index.tools.cassandra.base import CassandraDatabaseToolSpec
9+
from llama_index.tools.cassandra.cassandra_database_wrapper import (
10+
CassandraDatabase,
11+
)
12+
13+
14+
def test_tool_with_openai_tool(cassandra):
15+
session = cassio.config.resolve_session()
16+
17+
session.execute(
18+
"""
19+
CREATE TABLE IF NOT EXISTS default_keyspace.tool_table_users (
20+
user_id UUID PRIMARY KEY ,
21+
user_name TEXT ,
22+
password TEXT
23+
);
24+
"""
25+
)
26+
session.execute(
27+
"""
28+
CREATE INDEX user_name
29+
ON default_keyspace.tool_table_users (user_name);
30+
"""
31+
)
32+
33+
user_id = uuid.uuid4()
34+
session.execute(
35+
f"""
36+
INSERT INTO default_keyspace.tool_table_users (user_id, user_name)
37+
VALUES ({user_id}, 'my_user');
38+
"""
39+
)
40+
db = CassandraDatabase()
41+
42+
spec = CassandraDatabaseToolSpec(db=db)
43+
44+
tools = spec.to_tool_list()
45+
for tool in tools:
46+
print(tool.metadata.name)
47+
48+
llm = OpenAI(model="gpt-4o")
49+
agent = OpenAIAgent.from_tools(tools, llm=llm, verbose=True)
50+
51+
response = agent.chat(
52+
"What is the user_id of the user named 'my_user' in table default_keyspace.tool_table_users?"
53+
)
54+
print(response)
55+
assert response is not None
56+
assert str(user_id) in str(response)

libs/llamaindex/pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ llama-index = "0.10.31"
2121
llama-index-vector-stores-astra-db = "0.1.7"
2222
llama-index-vector-stores-cassandra = "0.1.3"
2323
llama-index-embeddings-langchain = "0.1.2"
24+
llama-index-tools-cassandra = "0.1.1"
2425
llama-parse = "0.4.1"
2526
# optional integrations
2627
## azure

0 commit comments

Comments
 (0)