Snowflake Authentication Methods for TruLens¶

This notebook shows how to connect TruLens to Snowflake using key-pair authentication -- no password required.

Before starting, make sure you have:

1. Generated an RSA key pair and assigned the public key to your Snowflake user
2. Stored the private key file path in your environment

In [ ]:

# !pip install trulens trulens-connectors-snowflake trulens-providers-cortex

# !pip install trulens trulens-connectors-snowflake trulens-providers-cortex

In [ ]:

from dotenv import load_dotenv

load_dotenv()

from dotenv import load_dotenv load_dotenv()

Connect to Snowflake with Key-Pair Auth¶

Pass private_key_file directly to SnowflakeConnector. No password needed.

In [ ]:

import os

from trulens.connectors.snowflake import SnowflakeConnector
from trulens.core import TruSession

conn = SnowflakeConnector(
    account=os.environ["SNOWFLAKE_ACCOUNT"],
    user=os.environ["SNOWFLAKE_USER"],
    role=os.environ["SNOWFLAKE_ROLE"],
    database=os.environ["SNOWFLAKE_DATABASE"],
    schema=os.environ["SNOWFLAKE_SCHEMA"],
    warehouse=os.environ["SNOWFLAKE_WAREHOUSE"],
    private_key_file=os.environ["SNOWFLAKE_PRIVATE_KEY_FILE"],
)

session = TruSession(connector=conn)

import os from trulens.connectors.snowflake import SnowflakeConnector from trulens.core import TruSession conn = SnowflakeConnector( account=os.environ["SNOWFLAKE_ACCOUNT"], user=os.environ["SNOWFLAKE_USER"], role=os.environ["SNOWFLAKE_ROLE"], database=os.environ["SNOWFLAKE_DATABASE"], schema=os.environ["SNOWFLAKE_SCHEMA"], warehouse=os.environ["SNOWFLAKE_WAREHOUSE"], private_key_file=os.environ["SNOWFLAKE_PRIVATE_KEY_FILE"], ) session = TruSession(connector=conn)

Create simple LLM app¶

In [ ]:

from snowflake.cortex import Complete
from trulens.core.otel.instrument import instrument
from trulens.otel.semconv.trace import SpanAttributes


class LLM:
    def __init__(self, model="mistral-large2"):
        self.model = model

    @instrument(span_type=SpanAttributes.SpanType.GENERATION)
    def complete(self, prompt):
        return Complete(self.model, prompt, session=conn.snowpark_session)


llm = LLM()

from snowflake.cortex import Complete from trulens.core.otel.instrument import instrument from trulens.otel.semconv.trace import SpanAttributes class LLM: def __init__(self, model="mistral-large2"): self.model = model @instrument(span_type=SpanAttributes.SpanType.GENERATION) def complete(self, prompt): return Complete(self.model, prompt, session=conn.snowpark_session) llm = LLM()

Set up feedback functions¶

Here we'll test answer relevance and coherence.

In [ ]:

from trulens.core import Metric, Selector
from trulens.core import Select
from trulens.providers.cortex import Cortex

provider = Cortex(
    conn.snowpark_session,
    model_engine="mistral-large2",
)

f_answer_relevance = (
    Metric(
    implementation=provider.relevance_with_cot_reasons,
    name="Answer Relevance",
    selectors={
        "prompt": Selector.select_record_input(),
        "response": Selector.select_record_output(),
    },
))

f_coherence = Metric(
    implementation=provider.coherence_with_cot_reasons,
    name="Coherence",
    selectors={"response": Selector.select_record_output()},
)

from trulens.core import Metric, Selector from trulens.core import Select from trulens.providers.cortex import Cortex provider = Cortex( conn.snowpark_session, model_engine="mistral-large2", ) f_answer_relevance = ( Metric( implementation=provider.relevance_with_cot_reasons, name="Answer Relevance", selectors={ "prompt": Selector.select_record_input(), "response": Selector.select_record_output(), }, )) f_coherence = Metric( implementation=provider.coherence_with_cot_reasons, name="Coherence", selectors={"response": Selector.select_record_output()}, )

Construct and run the app¶

In [ ]:

from trulens.apps.app import TruApp

tru_llm = TruApp(
    llm,
    app_name="KeyPairAuthDemo",
    app_version="v1",
    feedbacks=[
        f_answer_relevance,
        f_coherence,
    ],
)

with tru_llm as recording:
    resp = llm.complete("What is TruLens?")

print(resp)

from trulens.apps.app import TruApp tru_llm = TruApp( llm, app_name="KeyPairAuthDemo", app_version="v1", feedbacks=[ f_answer_relevance, f_coherence, ], ) with tru_llm as recording: resp = llm.complete("What is TruLens?") print(resp)

In [ ]:

session.get_leaderboard()

session.get_leaderboard()