Update app.py

app.py

@@ -12,8 +12,6 @@ PX_ENDPOINT_URL = "https://api.perplexity.ai"
 PX_API_KEY = os.getenv('PX_KEY')
 PASSWORD = os.getenv("PASSWD")  # Store the password in an environment variable
 
-DEPLOY_TO_HF = ["deepseek-ai/DeepSeek-V3"]
-
 px_client = OpenAI(base_url=PX_ENDPOINT_URL, api_key=PX_API_KEY)
 
 def respond(
@@ -34,17 +32,31 @@ def respond(
     messages.append({"role": "user", "content": message})
 
     response = ""
+    citations = []
     
-    for chunk in px_client.chat.completions.create(
-        model=AVAILABLE_MODELS[model_choice],  # Use the selected model
+    stream = px_client.chat.completions.create(
+        model=AVAILABLE_MODELS[model_choice],
         messages=messages,
         max_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
         stream=True,
-    ):
-        token = chunk.choices[0].delta.content or ""
-        response += token
+    )
+    
+    for chunk in stream:
+        if "choices" in chunk:
+            token = chunk.choices[0].delta.content or ""
+            response += token
+            yield response  # Stream response as it arrives
+        if "citations" in chunk:
+            citations = chunk["citations"]
+    
+    # Append citations as clickable links
+    if citations:
+        citation_text = "\n\nSources:\n" + "\n".join(
+            [f"[{i+1}] [{url}]({url})" for i, url in enumerate(citations)]
+        )
+        response += citation_text
         yield response
 
 def check_password(input_password):
@@ -62,24 +74,36 @@ with gr.Blocks() as demo:
         error_message = gr.Textbox(
             label="Error", visible=False, interactive=False
         )
-
+    
     with gr.Column(visible=False) as chat_interface:
+        system_prompt = gr.Textbox(
+            value="You are a helpful assistant.", label="System message"
+        )
         chat = gr.ChatInterface(
             respond,
-            additional_inputs=[
-                gr.Textbox(value="You are a helpful assistant.", label="System message"),
-                gr.Dropdown(
-                    choices=list(AVAILABLE_MODELS.keys()),
-                    value=list(AVAILABLE_MODELS.keys())[0],
-                    label="Select Model"
-                ),
-                gr.Slider(minimum=1, maximum=30000, value=2048, step=100, label="Max new tokens"),
-                gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-                gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
-            ],
+            additional_inputs=[],
+            height=600  # Increased height for better visibility
         )
-
-    submit_button.click(check_password, inputs=password_input, outputs=[password_input, chat_interface])
+        
+        with gr.Column():
+            model_choice = gr.Dropdown(
+                choices=list(AVAILABLE_MODELS.keys()),
+                value=list(AVAILABLE_MODELS.keys())[0],
+                label="Select Model"
+            )
+            max_tokens = gr.Slider(
+                minimum=1, maximum=30000, value=2048, step=100, label="Max new tokens"
+            )
+            temperature = gr.Slider(
+                minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"
+            )
+            top_p = gr.Slider(
+                minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"
+            )
+    
+    submit_button.click(
+        check_password, inputs=password_input, outputs=[password_input, chat_interface]
+    )
 
 if __name__ == "__main__":
     demo.launch(share=True)