Update app.py

app.py

@@ -3,12 +3,13 @@ import tempfile
 import gradio as gr
 import torch
 import logging
+import base64
 
 from operator import itemgetter
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
 from langchain_community.vectorstores.chroma import Chroma
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.schema import AIMessage, HumanMessage
@@ -16,6 +17,11 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain.globals import set_debug
 from dotenv import load_dotenv
 
+def image_to_base64(image_path):
+    with open(image_path, "rb") as image_file:
+        encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
+    return encoded_string
+
 # configure logging
 logging.basicConfig(level=logging.INFO)
 
@@ -92,10 +98,6 @@ llm = ChatOpenAI(
     model_name="gpt-4-0125-preview", openai_api_key=openai_api_key, temperature=0.1, streaming=True
 )
 
-llm_translate = ChatOpenAI(
-    model_name="gpt-3.5-turbo", openai_api_key=openai_api_key, temperature=0.0
-)
-
 retriever = configure_retriever(local_files)
 
 template = """Answer the question based only on the following context:
@@ -103,59 +105,78 @@ template = """Answer the question based only on the following context:
 
 Question: {question}
 
-Answer in German Language. If the question is not related to the context, answer with "I don't know" in German.
+Answer in German Language. If the question is not related to the context, answer with "I don't know".
 """
 
+
 prompt = ChatPromptTemplate.from_template(template)
 
-chain = (
-    {
-        "context": itemgetter("question") | retriever,
-        "question": itemgetter("question"),
-    }
-    | prompt
-    | llm
+
+chain_translate = (
+    llm
     | StrOutputParser()
 )
 
-chain_translate = (llm_translate
+chain_rag = (
+    { 
+    "context": itemgetter("question") | retriever, 
+    "question": itemgetter("question"),
+    "history": itemgetter("history")
+    }
+    | prompt
+    | llm
     | StrOutputParser()
 )
 
-
 def predict(message, history):
-    message = chain_translate.invoke(f"Translate this sentence to English: {message}")
+    message = chain_translate.invoke(f"Translate this query to English if it is in German otherwise return original contetn: {message}")
+
     history_langchain_format = []
+    partial_message = ""
+
     for human, ai in history:
         history_langchain_format.append(HumanMessage(content=human))
         history_langchain_format.append(AIMessage(content=ai))
     history_langchain_format.append(HumanMessage(content=message))
-    gpt_response = llm(history_langchain_format)
-    for chunk in chain.stream({"question": gpt_response.content}): # Stream the response
-        yield chunk
-
-
+    for response in chain_rag.stream({"question": message, "history": history_langchain_format}):
+        partial_message += response
+        yield partial_message 
+    
 image_path = "./ui/logo.png" if os.path.exists("./ui/logo.png") else "./logo.png"
+logo_base64 = image_to_base64(image_path)
+
+# CSS with the Base64-encoded image
+css = f"""
+body::before {{
+    content: '';
+    display: block;
+    height: 150px !important;  /* Adjust based on your logo's size */
+    background: url('data:image/png;base64,{logo_base64}') no-repeat center center !important;
+    background-size: contain !important; /* This makes sure the logo fits well in the header */
+}}
+
+#q-output {{
+    max-height: 60vh !important;
+    overflow: auto !important;
+}}
+"""
 
-with gr.Blocks() as demo:
-    gr.Image(image_path)
-    gr.ChatInterface(
-    predict,
-    chatbot=gr.Chatbot(height=500, show_share_button=True),
-    textbox=gr.Textbox(placeholder="stell mir Fragen", container=False, scale=7),
-    title="Beitrag Service",
-    description="Ich bin Ihr hilfreicher KI-Assistent",
-    theme="soft",
-    examples=[
-    "Generate auditing questions about Change Management",
-    "Generate auditing questions about Software Maintenance",
-    "Generate auditing questions about Data Protection",
-    "Generate auditing questions about IT",
-    "Generate auditing questions about control systems",
-    "Generate auditing questions about GDPR compliance",
-    ],
-    cache_examples=True,
-    ).launch(show_api= False)
-
-if __name__ == "__main__":
-    demo.launch()
+gr.ChatInterface(
+predict,
+chatbot=gr.Chatbot(likeable=True),
+textbox=gr.Textbox(placeholder="stell mir Fragen", scale=7),
+description="Ich bin Ihr hilfreicher KI-Assistent",
+theme="soft",
+submit_btn="Senden",
+retry_btn="🔄  Wiederholen",
+undo_btn="⏪  Rückgängig",
+clear_btn="🗑️  Löschen",
+examples=[
+"Generate auditing questions about Change Management",
+"Generate auditing questions about Software Maintenance",
+"Generate auditing questions about Data Protection"
+],
+#cache_examples=True,
+fill_height=True,
+css=css,
+).launch(show_api=False)