Query-rewriter / app.py
Daemontatox's picture
Update app.py
b4fe479 verified
Raw
History Blame Contribute Delete
11.5 kB
import subprocess
import os
import torch
from dotenv import load_dotenv
from langchain_community.vectorstores import Qdrant
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.prompts import ChatPromptTemplate
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser
from qdrant_client import QdrantClient, models
from langchain_openai import ChatOpenAI
import gradio as gr
import logging
from typing import List, Tuple, Generator
from dataclasses import dataclass
from datetime import datetime
from queue import Queue
from threading import Thread
from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@dataclass
class Message:
role: str
content: str
timestamp: str
class ChatHistory:
def __init__(self):
self.messages: List[Message] = []
def add_message(self, role: str, content: str):
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
self.messages.append(Message(role=role, content=content, timestamp=timestamp))
def get_formatted_history(self, max_messages: int = 10) -> str:
recent_messages = self.messages[-max_messages:] if len(self.messages) > max_messages else self.messages
formatted_history = "\n".join([
f"{msg.role}: {msg.content}" for msg in recent_messages
])
return formatted_history
def clear(self):
self.messages = []
# Load environment variables and setup
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")
QDRANT_URL = os.getenv("QDRANT_URL")
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
if not HF_TOKEN:
logger.error("HF_TOKEN is not set in the environment variables.")
exit(1)
# Query rewriting prompt template
query_rewrite_template = """
Given the chat history and the current question, rewrite the question to be more specific and include relevant context. The rewritten query should help retrieve more accurate information from the Mawared HR knowledge base.
Chat History:
{chat_history}
Current Question: {question}
Instructions:
1. Analyze the chat history and current question
2. Identify key concepts and context from previous messages
3. Incorporate relevant context into the question
4. Make the question more specific and detailed
5. Ensure the rewritten question focuses on Mawared HR functionality
Rewritten question:
{{Rewriten question}}
"""
query_rewrite_prompt = PromptTemplate(
input_variables=["chat_history", "question"],
template=query_rewrite_template
)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
try:
client = QdrantClient(
url=QDRANT_URL,
api_key=QDRANT_API_KEY,
prefer_grpc=False
)
except Exception as e:
logger.error("Failed to connect to Qdrant.")
exit(1)
collection_name = "mawared"
try:
client.create_collection(
collection_name=collection_name,
vectors_config=models.VectorParams(
size=384,
distance=models.Distance.COSINE
)
)
except Exception as e:
if "already exists" not in str(e):
logger.error(f"Error creating collection: {e}")
exit(1)
db = Qdrant(
client=client,
collection_name=collection_name,
embeddings=embeddings,
)
retriever = db.as_retriever(
search_type="similarity",
search_kwargs={"k": 5}
)
llm = ChatOpenAI(
model="meta-llama/Llama-3.3-70B-Instruct",
temperature=0,
max_tokens=None,
timeout=None,
max_retries=2,
api_key=HF_TOKEN,
base_url="https://api-inference.huggingface.co/v1/",
stream=True,
)
# Create query rewriting chain
query_rewrite_chain = LLMChain(
llm=llm,
prompt=query_rewrite_prompt,
verbose=True
)
# Updated Prompt Template with Chain-of-Thought Reasoning
reasoning_prompt_template = """
You are M.H.AI ,You are a highly specialized AI assistant for the Mawared HR System.
Your only function is to provide accurate, detailed, and contextually relevant support based strictly on the information within the provided context and the current chat history.
Your answer should include step-by-step reasoning (chain of thought) to explain your logic clearly and guide the user effectively.
If the question is not about Mawared HR System, don't answer it.
Your top priority is user experience and to make sure the user is satisfied.
After each answer ,keep the conversation flow by asking suitable questions about the user question and chat history.
When you receive a question, follow these steps to provide an accurate and relevant response:
1-♧Understand the Question♧: Read carefully to fully comprehend the context and details.
2-♧Break Down the Question♧: Break down the question into more specific sub-questions.
3-♧Identify Key Elements♧: Identify important points and potential sub-questions.
4-♧Formulate a Hypothesis♧: Propose a preliminary idea based on your understanding.
5-♧Gather Evidence♧: Verify information from your knowledge base.
6-♧Analyze Consequences♧: Consider the potential consequences of your response.
7-♧Question Your Hypotheses♧: Consider alternative perspectives to your initial hypothesis.
8-♧Consider Alternative Scenarios♧: Think of innovative solutions or alternative scenarios to solve the problem posed.
9-♧Use Analogies and Metaphors♧: Illustrate your points with analogies and metaphors to make the explanation more intuitive.
10-♧Provide Real Examples and Case Studies♧: Use concrete examples and case studies to make the response more concrete.
11-♧Acknowledge Limitations♧: Recognize the limits of your knowledge and indicate when you cannot provide an accurate response.
12-♧Use an Engaging and Accessible Tone♧: Use an engaging tone and accessible language to make the response more enjoyable to read and understand.
13-♧Explain Your Responses♧: Clearly explain your responses to reinforce transparency and user trust.
14-♧Cite Sources♧: If possible, cite reliable sources from your knowledge base to support your assertions.
15-♧Refine Your Response♧: Clarify your thoughts and improve your reasoning.
16-♧Review Your Response♧: Read through to ensure it is clear, concise, and error-free.
Then provide the final response.
Chat History:
{chat_history}
Retrieved Context:
{context}
Current Question:
{question}
Answer with detailed reasoning:
{{response}}
"""
reasoning_prompt = ChatPromptTemplate.from_template(reasoning_prompt_template)
def create_rag_chain(chat_history: str):
chain = (
{
"context": retriever,
"question": RunnablePassthrough(),
"chat_history": lambda x: chat_history
}
| reasoning_prompt
| llm
| StrOutputParser()
)
return chain
chat_history = ChatHistory()
def process_stream(stream_queue: Queue, history: List[List[str]]) -> Generator[List[List[str]], None, None]:
current_response = ""
while True:
chunk = stream_queue.get()
if chunk is None:
break
current_response += chunk
new_history = history.copy()
new_history[-1][1] = current_response
yield new_history
def rewrite_query_sync(question: str, formatted_history: str) -> str:
"""
Synchronous wrapper for the asynchronous `query_rewrite_chain` execution.
Collects stream output if returned and converts it into a full response.
"""
try:
# Run the query rewrite chain
result = query_rewrite_chain.run(
question=question,
chat_history=formatted_history
)
# If the result is a stream, concatenate its parts
if isinstance(result, Generator):
result = "".join(result)
return result.strip()
except Exception as e:
logger.error(f"Error in query rewriting: {e}")
return question
def ask_question_gradio(question: str, history: List[List[str]]) -> Generator[tuple, None, None]:
try:
if history is None:
history = []
chat_history.add_message("user", question)
formatted_history = chat_history.get_formatted_history()
# Rewrite the query (synchronously)
rewritten_question = rewrite_query_sync(question, formatted_history)
logger.info(f"Rewritten question: {rewritten_question}")
rag_chain = create_rag_chain(formatted_history)
history.append([question, ""])
stream_queue = Queue()
def stream_processor():
try:
for chunk in rag_chain.stream(rewritten_question):
stream_queue.put(chunk)
stream_queue.put(None)
except Exception as e:
logger.error(f"Streaming error: {e}")
stream_queue.put(None)
Thread(target=stream_processor).start()
response = ""
for updated_history in process_stream(stream_queue, history):
response = updated_history[-1][1]
yield "", updated_history
chat_history.add_message("assistant", response)
except Exception as e:
logger.error(f"Error during question processing: {e}")
if not history:
history = []
history.append([question, "An error occurred. Please try again later."])
yield "", history
def clear_chat():
chat_history.clear()
return [], ""
# Gradio Interface
with gr.Blocks(theme='earneleh/paris') as iface:
gr.Image("Image.webp", width=750, height=300, show_label=False, show_download_button=False)
gr.Markdown("# Mawared HR Assistant 2.7.6 Open-Beta")
gr.Markdown('### Patch Notes')
gr.Markdown("""
1-Changed Reasoning process to use an intricate and detailed 16 steps process!!
2-Fixed issues with Query/Question rewriting (i hope).
3-Fixed History bug where model would secretly Cache previous responses to use as Information even after clearing the chat.
4-Included better Error handling and Callbacks.
**WARNING:** MODEL IS HIGHLY INTELLIGENT AND IS PRONE TO OVERTHINKING, ANXIETY, TEXT Deadlock AND SELF DOUBT.
IF THAT HAPPENS PLEASE REPORT IT.
Lastly, happy bugging :] .
""")
chatbot = gr.Chatbot(
height=750,
show_label=False,
bubble_full_width=False,
)
with gr.Row():
with gr.Column(scale=20):
question_input = gr.Textbox(
label="Ask a question:",
placeholder="Type your question here...",
show_label=False
)
with gr.Column(scale=4):
with gr.Row():
with gr.Column():
send_button = gr.Button("Send", variant="primary", size="sm")
clear_button = gr.Button("Clear Chat", size="sm")
submit_events = [question_input.submit, send_button.click]
for submit_event in submit_events:
submit_event(
ask_question_gradio,
inputs=[question_input, chatbot],
outputs=[question_input, chatbot]
)
clear_button.click(
clear_chat,
outputs=[chatbot, question_input]
)
if __name__ == "__main__":
iface.launch()