feat: add langchain cassandra tool (#440)

* feat: add langchain cassandra tool

* add

Author: nicoloboschi

libs/e2e-tests/e2e_tests/langchain/conftest.py

@@ -0,0 +1,25 @@
+import pytest
+
+from e2e_tests.conftest import (
+    get_vector_store_handler,
+)
+
+from e2e_tests.test_utils.vector_store_handler import (
+    VectorStoreImplementation,
+)
+
+
+@pytest.fixture
+def astra_db():
+    handler = get_vector_store_handler(VectorStoreImplementation.ASTRADB)
+    context = handler.before_test()
+    yield context
+    handler.after_test()
+
+
+@pytest.fixture
+def cassandra():
+    handler = get_vector_store_handler(VectorStoreImplementation.CASSANDRA)
+    context = handler.before_test()
+    yield context
+    handler.after_test()

libs/e2e-tests/e2e_tests/langchain/test_cassandra_tool.py

@@ -0,0 +1,63 @@
+import uuid
+
+import cassio
+
+from langchain.agents import AgentExecutor, create_openai_tools_agent
+from langchain import hub
+from langchain_community.tools.cassandra_database.tool import (
+    GetSchemaCassandraDatabaseTool,
+    GetTableDataCassandraDatabaseTool,
+    QueryCassandraDatabaseTool,
+)
+from langchain_community.utilities.cassandra_database import CassandraDatabase
+from langchain_openai import ChatOpenAI
+
+
+def test_tool_with_openai_tool(cassandra):
+    session = cassio.config.resolve_session()
+
+    session.execute(
+        """
+        CREATE TABLE IF NOT EXISTS default_keyspace.tool_table_users (
+        user_id UUID PRIMARY KEY ,
+        user_name TEXT ,
+        password TEXT
+    );
+    """
+    )
+    session.execute(
+        """
+    CREATE INDEX user_name
+   ON default_keyspace.tool_table_users (user_name);
+    """
+    )
+
+    user_id = uuid.uuid4()
+    session.execute(
+        f"""
+        INSERT INTO default_keyspace.tool_table_users (user_id, user_name) 
+        VALUES ({user_id}, 'my_user');
+    """
+    )
+    db = CassandraDatabase()
+
+    query_tool = QueryCassandraDatabaseTool(db=db)
+    schema_tool = GetSchemaCassandraDatabaseTool(db=db)
+    select_data_tool = GetTableDataCassandraDatabaseTool(db=db)
+
+    tools = [schema_tool, select_data_tool, query_tool]
+
+    model = ChatOpenAI(model="gpt-4o")
+
+    prompt = hub.pull("hwchase17/openai-tools-agent")
+
+    agent = create_openai_tools_agent(model, tools, prompt)
+    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+    response = agent_executor.invoke(
+        {
+            "input": "What is the user_id of the user named 'my_user' in table default_keyspace.tool_table_users?"
+        }
+    )
+    print(response)
+    assert response is not None
+    assert str(user_id) in str(response)

libs/e2e-tests/e2e_tests/langchain/test_unstructured.py

@@ -17,31 +17,13 @@
 from e2e_tests.conftest import (
     set_current_test_info,
     get_required_env,
-    get_vector_store_handler,
 )
 
 from e2e_tests.test_utils.vector_store_handler import (
-    VectorStoreImplementation,
     VectorStoreTestContext,
 )
 
 
-@pytest.fixture
-def astra_db():
-    handler = get_vector_store_handler(VectorStoreImplementation.ASTRADB)
-    context = handler.before_test()
-    yield context
-    handler.after_test()
-
-
-@pytest.fixture
-def cassandra():
-    handler = get_vector_store_handler(VectorStoreImplementation.CASSANDRA)
-    context = handler.before_test()
-    yield context
-    handler.after_test()
-
-
 @pytest.mark.parametrize("vector_store", ["cassandra", "astra_db"])
 @pytest.mark.parametrize("unstructured_mode", ["single", "elements"])
 def test_unstructured_api(vector_store, unstructured_mode, request):

libs/e2e-tests/pyproject.langchain.toml

@@ -19,6 +19,7 @@ pillow = "^10.2.0"
 python-dotenv = "^1.0.1"
 trulens-eval = "^0.21.0"
 nemoguardrails = "^0.8.0"
+langchainhub = "^0.1.15"
 
 # From LangChain optional deps, needed by WebBaseLoader
 beautifulsoup4 = "^4"

libs/e2e-tests/pyproject.llamaindex.toml

@@ -17,6 +17,7 @@ pillow = "^10.2.0"
 python-dotenv = "^1.0.1"
 trulens-eval = "^0.21.0"
 nemoguardrails = "^0.8.0"
+langchainhub = "^0.1.15"
 
 # From LangChain optional deps, needed by WebBaseLoader
 beautifulsoup4 = "^4"

libs/e2e-tests/pyproject.ragstack-ai.toml

@@ -20,6 +20,7 @@ pillow = "^10.2.0"
 python-dotenv = "^1.0.1"
 trulens-eval = "^0.21.0"
 nemoguardrails = "^0.8.0"
+langchainhub = "^0.1.15"
 
 # From LangChain optional deps, needed by WebBaseLoader
 beautifulsoup4 = "^4"

libs/e2e-tests/pyproject.toml

@@ -19,6 +19,7 @@ pillow = "^10.2.0"
 python-dotenv = "^1.0.1"
 trulens-eval = "^0.21.0"
 nemoguardrails = "^0.8.0"
+langchainhub = "^0.1.15"
 
 # From LangChain optional deps, needed by WebBaseLoader
 beautifulsoup4 = "^4"