galileo-protect-demo.py

import os
from typing import List
from uuid import UUID
import streamlit as st

from galileo.handlers.langchain import GalileoCallback
from galileo.handlers.langchain.tool import ProtectTool, ProtectParser
from galileo_core.schemas.protect.ruleset import Ruleset
from galileo_core.schemas.protect.action import OverrideAction

from langchain.schema.document import Document
from langchain_community.callbacks import StreamlitCallbackHandler
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore
from dotenv import load_dotenv


load_dotenv()

# Create a callback handler
galileo_callback = GalileoCallback()

# A hack to "clear" the previous result when submitting a new prompt. This avoids
# the "previous run's text is grayed-out but visible during rerun" Streamlit behavior.
class DirtyState:
    NOT_DIRTY = "NOT_DIRTY"
    DIRTY = "DIRTY"
    UNHANDLED_SUBMIT = "UNHANDLED_SUBMIT"


def get_dirty_state() -> str:
    return st.session_state.get("dirty_state", DirtyState.NOT_DIRTY)


def set_dirty_state(state: str) -> None:
    st.session_state["dirty_state"] = state


def with_clear_container(submit_clicked):
    if get_dirty_state() == DirtyState.DIRTY:
        if submit_clicked:
            set_dirty_state(DirtyState.UNHANDLED_SUBMIT)
            st.rerun()
        else:
            set_dirty_state(DirtyState.NOT_DIRTY)

    if submit_clicked or get_dirty_state() == DirtyState.UNHANDLED_SUBMIT:
        set_dirty_state(DirtyState.DIRTY)
        return True

    return False


st.set_page_config(
    page_title="Finance Agent",
    page_icon="🔭",
    layout="centered",
    initial_sidebar_state="collapsed",
)

# Custom CSS to reduce width and left-align
st.markdown("""
<style>
    .main .block-container {
        max-width: 800px;
        padding-left: 2rem;
        padding-right: 2rem;
    }
    .stForm {
        max-width: 800px;
    }
    .stChatMessage {
        max-width: 800px;
    }
</style>
""", unsafe_allow_html=True)

"# 🔭 Finance Agent"

user_openai_api_key = os.getenv("OPENAI_API_KEY") or st.secrets["OPENAI_API_KEY"]

if user_openai_api_key:
    openai_api_key = user_openai_api_key
    enable_custom = True
else:
    openai_api_key = "not_supplied"
    enable_custom = False


input_rulesets = [
    Ruleset(
        rules=[
            {
                "metric": "prompt_injection",
                "operator": "eq",
                "target_value": "impersonation",
            },
        ],
        action=OverrideAction(
            choices=["Sorry, prompt injection detected in the user input. I cannot answer that question."]
        ),
    ),
    Ruleset(
        rules=[
            {
                "metric": "prompt_injection",
                "operator": "eq",
                "target_value": "new_context",
            },
        ],
        action=OverrideAction(
            choices=["Sorry, prompt injection detected in the user input. I cannot answer that question."]
        ),
    ),
    Ruleset(
        rules=[
            {
                "metric": "input_pii",
                "operator": "contains",
                "target_value": "date_of_birth",
            },
        ],
        action=OverrideAction(
            choices=["Sorry, PII detected in the user input. I cannot answer that question."]
        ),
    ),
]

output_rulesets = [
    Ruleset(
        rules=[
            {
                "metric": "context_adherence_luna",
                "operator": "gt",
                "target_value": 0.7,
            },
        ],
        action=OverrideAction(
            choices=["Sorry, hallucination detected in the model output. I cannot answer that question."]
        ),
    ),
    Ruleset(
        rules=[
            {
                "metric": "pii",
                "operator": "contains",
                "target_value": "address",
            },
        ],
        action=OverrideAction(
            choices=["Sorry, personal address detected in the model output. I cannot answer that question."]
        ),
    ),
]

galileo_api_key = os.getenv("GALILEO_API_KEY") or st.secrets["GALILEO_API_KEY"]

gp_tool = ProtectTool(
    stage_name="My first stage",
    project_name=os.getenv("GALILEO_PROJECT"),
    prioritized_rulesets=input_rulesets,
    api_key=galileo_api_key,
    timeout=10
)

gp_tool_output = ProtectTool(
    stage_name="My second stage",
    project_name=os.getenv("GALILEO_PROJECT"),
    timeout=15,
    prioritized_rulesets=output_rulesets,
    api_key=galileo_api_key,
)


embeddings = OpenAIEmbeddings()
vectordb = PineconeVectorStore(index_name="galileo-demo", embedding=embeddings, namespace="sp500-qa-demo")
llm = ChatOpenAI(temperature=0, api_key=user_openai_api_key)


def format_docs(docs: List[Document]) -> str:
    return "\n\n".join([d.page_content for d in docs])


retriever = vectordb.as_retriever()

template = """You are a helpful assistant. Given the context below, please answer the following questions:

    {context}

    Question: {question}
    """
prompt = ChatPromptTemplate.from_template(template)
model = ChatOpenAI(name="gpt-3.5-turbo", temperature=0)

rag_chain_original = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()} | prompt | model | StrOutputParser()
)

gp_output_parser = ProtectParser(chain=StrOutputParser())

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | {"output": model | StrOutputParser(), "input": lambda x: x.to_string()}
    | gp_tool_output
    | gp_output_parser.parser
)

gp_exec = ProtectParser(chain=rag_chain, echo_output=True)

gp_chain = gp_tool | gp_exec.parser

with st.form(key="form"):
    user_input = ""

    if enable_custom:
        user_input = st.text_input("")
    submit_clicked = st.form_submit_button("Submit")

output_container = st.empty()

if with_clear_container(submit_clicked):
    output_container = output_container.container()
    output_container.chat_message("user").write(user_input)

    answer_container = output_container.chat_message("assistant", avatar="🔭")
    st_callback = StreamlitCallbackHandler(answer_container)

    model_answer = rag_chain_original.invoke(user_input)
    gp_answer = gp_chain.invoke(user_input, config={"callbacks": [galileo_callback]})

    answer_container.write("**Response from the model:**")
    answer_container.write(model_answer)

    answer_container.write("**Response from Galileo Protect:**")
    answer_container.write(gp_answer)
    answer_container.write('Please go to Galileo Console  below to see the detailed trace.')
    st.link_button('Traces', 'https://console.demo-v2.galileocloud.io/project/61295445-fefb-4269-93c6-a37457d3480a/log-streams/035348b8-6448-4bfe-a135-f4c8e6ebabf2?tableLevel=traces&timeRange=%7B%22type%22%3A%22last6Months%22%7D')