Update app.py

app.py

@@ -1,7 +1,73 @@
 import gradio as gr
+from langchain.embeddings import SentenceTransformerEmbeddings
+from langchain.vectorstores import FAISS
+from langchain_community.chat_models.huggingface import ChatHuggingFace
+from langchain.schema import SystemMessage, HumanMessage, AIMessage
+from langchain_community.llms import HuggingFaceEndpoint
 
-def greet(name):
-    return "Hello " + name + "!!"
+model_name = "sentence-transformers/all-mpnet-base-v2"
+embedding_llm = SentenceTransformerEmbeddings(model_name=model_name)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+db = FAISS.load_local("faiss_index", embedding_llm, allow_dangerous_deserialization=True)
+
+# Set up Hugging Face model
+llm = HuggingFaceEndpoint(
+    repo_id="HuggingFaceH4/starchat2-15b-v0.1",
+    task="text-generation",
+    max_new_tokens=4096,
+    temperature=0.6,
+    top_p=0.9,
+    top_k=40,
+    repetition_penalty=1.2,
+    do_sample=True,
+)
+chat_model = ChatHuggingFace(llm=llm)
+
+messages = [
+    SystemMessage(content="You are a helpful assistant."),
+    HumanMessage(content="Hi AI, how are you today?"),
+    AIMessage(content="I'm great thank you. How can I help you?")
+]
+
+def handle_query(mode: str, query: str):
+    if mode == "chat":
+        return chat_mode(query)
+    elif mode == "web-search":
+        return web_search(query)
+    else:
+        return "Invalid mode selected."
+
+def chat_mode(query: str):
+    prompt = HumanMessage(content=query)
+    messages.append(prompt)
+    response = chat_model.invoke(messages)
+    messages.append(response.content)
+    if len(messages) >= 6:
+        messages = messages[-6:]
+    return f"You: {query}\nIT-Assistant: {response.content}"
+
+def web_search(query: str):
+    similar_docs = db.similarity_search(query, k=3)
+    if similar_docs:
+        source_knowledge = "\n".join([x.page_content for x in similar_docs])
+    else:
+        source_knowledge = ""
+    augmented_prompt = f"""
+    Answer the next query using the provided Web Search.
+    If the answer is not contained in the Web Search, ignore the web search and respond independently with your own knowledge.
+
+    Query: {query}
+
+    Web Search:
+    {source_knowledge}
+    """
+    prompt = HumanMessage(content=augmented_prompt)
+    messages.append(prompt)
+    response = chat_model.invoke(messages)
+    messages.append(response.content)
+    if len(messages) >= 6:
+        messages = messages[-6:]
+    return f"You: {query}\nIT-Assistant: {response.content}"
+
+demo = gr.Interface(fn=handle_query, inputs=["text", "text"], outputs="text", title="IT Assistant", description="Choose a mode and enter your message, then click submit to interact.", inputs_layout="vertical", choices=["chat", "web-search"])
 demo.launch()