langchain/templates/cassandra-entomology-rag/cassandra_entomology_rag/cassandra_cluster_init.py
Erick Friis ebf998acb6
Templates (#12294)
Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
Co-authored-by: Lance Martin <lance@langchain.dev>
Co-authored-by: Jacob Lee <jacoblee93@gmail.com>
2023-10-25 18:47:42 -07:00

28 lines
849 B
Python

import os
from cassandra.cluster import Cluster
from cassandra.auth import PlainTextAuthProvider
def get_cassandra_connection():
contact_points = [
cp.strip()
for cp in os.environ.get("CASSANDRA_CONTACT_POINTS", "").split(',')
if cp.strip()
]
CASSANDRA_KEYSPACE = os.environ["CASSANDRA_KEYSPACE"]
CASSANDRA_USERNAME = os.environ.get("CASSANDRA_USERNAME")
CASSANDRA_PASSWORD = os.environ.get("CASSANDRA_PASSWORD")
#
if CASSANDRA_USERNAME and CASSANDRA_PASSWORD:
auth_provider = PlainTextAuthProvider(
CASSANDRA_USERNAME,
CASSANDRA_PASSWORD,
)
else:
auth_provider = None
c_cluster = Cluster(contact_points if contact_points else None, auth_provider=auth_provider)
session = c_cluster.connect()
return (session, CASSANDRA_KEYSPACE)