app.py

from gradio_client import Client
from trulens_eval.schema import Select
from trulens_eval.tru import Tru
from trulens_eval.feedback import Feedback
from trulens_eval.feedback import OpenAI as Feedback_OpenAI
from langchain.llms import HuggingFacePipeline
from langchain.prompts import PromptTemplate
from langchain.chains import ConversationChain
from langchain.chains.conversation.memory import ConversationBufferWindowMemory
import os

from langchain.schema.runnable import Runnable
# from langchain.llms import HuggingFaceInference


# Access environment variables
openai_api_key = os.environ.get("OPENAI_API_KEY")
huggingface_api_token = os.environ.get("HUGGINGFACE_API_TOKEN")


# Define a feedback function for query-statement relevance using OpenAI
feedback_openai = Feedback_OpenAI()
qa_relevance = Feedback(feedback_openai.relevance, name="Answer Relevance").on_input_output()

# Create a Tru object
tru = Tru()

# Set the window memory to go back 4 turns
window_memory = ConversationBufferWindowMemory(k=4)

# Define a custom function to interact with the Gradio client
def gradio_client_interaction(prompt):
    client = Client("https://tonic-stablemed-chat.hf.space/")
    result = client.predict(prompt, prompt, api_name="/predict")
    return result['data'][0]  # Assuming the response format

class GradioLLM(Runnable):
# class GradioLLM(HuggingFaceInference):
    def __init__(self, client_url):
        super().__init__()  # Call the constructor of Runnable if necessary
        self.client = Client(client_url)

    def invoke(self, prompt, **kwargs):
        # Implement the invoke method here
        result = self.client.predict(prompt, prompt, api_name="/predict")
        return result['data'][0]  # Adjust based on the actual response format    

    def generate(self, prompt, **kwargs):
        # Assuming the API expects 'prompt' and 'system_prompt' as inputs
        result = self.client.predict(prompt, prompt, api_name="/predict")
        return result['data'][0]  # Adjust based on the actual response format

# Initialize the GradioLLM with the URL
gradio_llm = GradioLLM("https://tonic-stablemed-chat.hf.space/")

# Create the ConversationChain with the GradioLLM
conversation = ConversationChain(
    llm=gradio_llm,
    verbose=True,
    memory=window_memory
)
# Update the conversation prompt template
conversation.prompt.template = '''The following is a consult between a clinical consultant and a public health and medical expert. The AI is an expert on medicine and public health and gives recommendations specific to location and conditions. If the AI does not know the answer to a question, it truthfully says it does not know.
Current conversation:
{history}
Human: {input}
AI:'''

# Wrap the conversation with TruChain to instrument it
tc_conversation = tru.Chain(conversation, app_id='Trulens-StableMed', feedbacks=[qa_relevance])

# Make a prediction using the wrapped conversation
user_input = "Howdy!"
system_prompt = "Howdy!"
result = tc_conversation.predict(user_input, system_prompt=system_prompt)

# Print the result
print(result)

# Run the TruLens dashboard
tru.run_dashboard()