File tree Expand file tree Collapse file tree 2 files changed +57
-0
lines changed
e2e-tests/e2e_tests/llama_index Expand file tree Collapse file tree 2 files changed +57
-0
lines changed Original file line number Diff line number Diff line change
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 )
Original file line number Diff line number Diff line change @@ -21,6 +21,7 @@ llama-index = "0.10.31"
21
21
llama-index-vector-stores-astra-db = " 0.1.7"
22
22
llama-index-vector-stores-cassandra = " 0.1.3"
23
23
llama-index-embeddings-langchain = " 0.1.2"
24
+ llama-index-tools-cassandra = " 0.1.1"
24
25
llama-parse = " 0.4.1"
25
26
# optional integrations
26
27
# # azure
You can’t perform that action at this time.
0 commit comments