Update app.py

app.py

@@ -2,6 +2,10 @@ import os
 import json
 import re
 import hashlib
+import os
+import json
+import re
+import hashlib
 import gradio as gr
 from functools import partial
 import threading
@@ -104,14 +108,13 @@ knowledge_base = KnowledgeBase()
 # LLMs
 # repharser_llm = ChatNVIDIA(model="mistralai/mistral-7b-instruct-v0.3") | StrOutputParser()
 repharser_llm = ChatNVIDIA(model="microsoft/phi-3-mini-4k-instruct") | StrOutputParser()
-validation_llm = ChatNVIDIA(model="microsoft/phi-3-small-8k-instruct") | StrOutputParser()
-instruct_llm = ChatNVIDIA(model="mistralai/mistral-7b-instruct-v0.2") | StrOutputParser()
-#instruct_llm = ChatNVIDIA(model="mistralai/mixtral-8x22b-instruct-v0.1") | StrOutputParser()
+instruct_llm = ChatNVIDIA(model="mistralai/mixtral-8x22b-instruct-v0.1") | StrOutputParser()
 relevance_llm = ChatNVIDIA(model="nvidia/llama-3.1-nemotron-70b-instruct") | StrOutputParser()
 answer_llm = ChatOpenAI(
     model="gpt-4o",              
     temperature=0.3,             
-    openai_api_key=os.getenv("OPENAI_API_KEY")
+    openai_api_key=os.getenv("OPENAI_API_KEY"),
+    streaming=True
 ) | StrOutputParser()
 
 
@@ -395,7 +398,7 @@ select_and_prompt = RunnableLambda(lambda x:
 answer_chain = (
     prepare_answer_inputs
     | select_and_prompt
-    | answer_llm
+    | relevance_llm
 )
 
 def RExtract(pydantic_class: Type[BaseModel], llm, prompt):
@@ -460,7 +463,7 @@ def update_knowledge_base(user_input: str, assistant_response: str):
 full_pipeline = hybrid_chain | RunnableAssign({"validation": validation_chain}) | answer_chain
 
 
-def chat_interface(message, history):    
+def chat_interface(message, history):
     inputs = {
         "query": message,
         "all_queries": [message],
@@ -530,7 +533,7 @@ demo = gr.ChatInterface(
     description="💡 Ask anything about Krishna Vamsi Dhulipalla",
     examples=[
         "Give me an overview of Krishna Vamsi Dhulipalla’s work experience across different roles?",
-        "What programming languages and tools does Krishna use for data science and data engineering?",
+        "What programming languages and tools does Krishna use for data science?",
         "Can this chatbot tell me what Krishna's chatbot architecture looks like and how it works?"
     ],
 )