import os

import cassio
import langchain
from langchain.cache import CassandraCache
from langchain_community.chat_models import ChatOpenAI
from langchain_core.messages import BaseMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda

use_cassandra = int(os.environ.get("USE_CASSANDRA_CLUSTER", "0"))
if use_cassandra:
    from .cassandra_cluster_init import get_cassandra_connection

    session, keyspace = get_cassandra_connection()
    cassio.init(
        session=session,
        keyspace=keyspace,
    )
else:
    cassio.init(
        token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
        database_id=os.environ["ASTRA_DB_ID"],
        keyspace=os.environ.get("ASTRA_DB_KEYSPACE"),
    )

# inits
langchain.llm_cache = CassandraCache(session=None, keyspace=None)
llm = ChatOpenAI()


# custom runnables
def msg_splitter(msg: BaseMessage):
    return [w.strip() for w in msg.content.split(",") if w.strip()]


# synonym-route preparation
synonym_prompt = ChatPromptTemplate.from_template(
    "List up to five comma-separated synonyms of this word: {word}"
)

chain = synonym_prompt | llm | RunnableLambda(msg_splitter)

langchain_core.prompts.ChatPromptTemplate.from_template langchain_community.chat_models.ChatOpenAI langchain_core.runnables.RunnableLambda langchain.cache.CassandraCache
[((788, 831), 'langchain.cache.CassandraCache', 'CassandraCache', ([], {'session': 'None', 'keyspace': 'None'}), '(session=None, keyspace=None)\n', (802, 831), False, 'from langchain.cache import CassandraCache\n'), ((838, 850), 'langchain_community.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {}), '()\n', (848, 850), False, 'from langchain_community.chat_models import ChatOpenAI\n'), ((1023, 1125), 'langchain_core.prompts.ChatPromptTemplate.from_template', 'ChatPromptTemplate.from_template', (['"""List up to five comma-separated synonyms of this word: {word}"""'], {}), "(\n    'List up to five comma-separated synonyms of this word: {word}')\n", (1055, 1125), False, 'from langchain_core.prompts import ChatPromptTemplate\n'), ((315, 359), 'os.environ.get', 'os.environ.get', (['"""USE_CASSANDRA_CLUSTER"""', '"""0"""'], {}), "('USE_CASSANDRA_CLUSTER', '0')\n", (329, 359), False, 'import os\n'), ((500, 547), 'cassio.init', 'cassio.init', ([], {'session': 'session', 'keyspace': 'keyspace'}), '(session=session, keyspace=keyspace)\n', (511, 547), False, 'import cassio\n'), ((1159, 1187), 'langchain_core.runnables.RunnableLambda', 'RunnableLambda', (['msg_splitter'], {}), '(msg_splitter)\n', (1173, 1187), False, 'from langchain_core.runnables import RunnableLambda\n'), ((714, 749), 'os.environ.get', 'os.environ.get', (['"""ASTRA_DB_KEYSPACE"""'], {}), "('ASTRA_DB_KEYSPACE')\n", (728, 749), False, 'import os\n')]