Update app.py

app.py

@@ -9,6 +9,9 @@ from langchain.chains import ConversationChain
 from langchain.chains.conversation.memory import ConversationBufferWindowMemory
 import os
 
+from langchain.llms import Runnable
+
+
 # Access environment variables
 openai_api_key = os.environ.get("OPENAI_API_KEY")
 huggingface_api_token = os.environ.get("HUGGINGFACE_API_TOKEN")
@@ -30,13 +33,24 @@ def gradio_client_interaction(prompt):
     result = client.predict(prompt, prompt, api_name="/predict")
     return result['data'][0]  # Assuming the response format
 
-# Create the ConversationChain with the custom function
+class GradioLLM(Runnable):
+    def __init__(self, client_url):
+        self.client = Client(client_url)
+
+    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_client_interaction,
+    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: