Update app.py

app.py

@@ -3,12 +3,12 @@ import warnings
 import gradio as gr
 from dotenv import load_dotenv
 
+from langchain.chains import ConversationalRetrievalChain
 from langchain_community.vectorstores import FAISS
-from langchain_community.embeddings import OpenAIEmbeddings, AzureOpenAIEmbeddings
-# from langchain_community.embeddings.openai import OpenAIEmbeddings
-from openai import AzureOpenAI
+from langchain_community.embeddings import AzureOpenAIEmbeddings
+from langchain_community.chat_models import AzureChatOpenAI
 
-# Patch Gradio bug (schema parsing issue)
+# Patch Gradio bug
 import gradio_client.utils
 gradio_client.utils.json_schema_to_python_type = lambda schema, defs=None: "string"
 
@@ -17,63 +17,56 @@ load_dotenv()
 AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
 AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
 AZURE_OPENAI_LLM_DEPLOYMENT = os.getenv("AZURE_OPENAI_LLM_DEPLOYMENT")
+AZURE_OPENAI_EMBEDDING_DEPLOYMENT = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")
 
-embeddings = OpenAIEmbeddings(openai_api_key=AZURE_OPENAI_API_KEY)
-
-if not all([AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_LLM_DEPLOYMENT]):
-    raise ValueError("Azure OpenAI environment variables are missing.")
+if not all([AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_LLM_DEPLOYMENT, AZURE_OPENAI_EMBEDDING_DEPLOYMENT]):
+    raise ValueError("Missing one or more Azure OpenAI environment variables.")
 
 # Suppress warnings
 warnings.filterwarnings("ignore")
 
-# Initialize embedding model
-# embeddings = AzureOpenAIEmbeddings(
-#     azure_deployment=OPENAI_EMBEDDING,
-#     azure_endpoint=AZURE_OPENAI_ENDPOINT,
-#     openai_api_key=AZURE_OPENAI_API_KEY,
-#     openai_api_version="2024-08-01-preview",
-#     chunk_size=1000
-# )
+# Initialize Azure embedding model
+embeddings = AzureOpenAIEmbeddings(
+    azure_deployment=AZURE_OPENAI_EMBEDDING_DEPLOYMENT,
+    azure_endpoint=AZURE_OPENAI_ENDPOINT,
+    openai_api_key=AZURE_OPENAI_API_KEY,
+    openai_api_version="2024-08-01-preview",
+    chunk_size=1000
+)
 
 # Load FAISS vector store
 vectorstore = FAISS.load_local(
     "faiss_index_sysml", embeddings, allow_dangerous_deserialization=True
 )
 
-# Initialize Azure OpenAI client directly
-client = AzureOpenAI(
-    api_key=AZURE_OPENAI_API_KEY,
+# Initialize Azure chat model
+llm = AzureChatOpenAI(
+    deployment_name=AZURE_OPENAI_LLM_DEPLOYMENT,
     azure_endpoint=AZURE_OPENAI_ENDPOINT,
-    api_version="2024-08-01-preview"
+    openai_api_key=AZURE_OPENAI_API_KEY,
+    openai_api_version="2024-08-01-preview",
+    temperature=0.5
+)
+
+# Build conversational RAG chain
+qa = ConversationalRetrievalChain.from_llm(
+    llm=llm,
+    retriever=vectorstore.as_retriever(),
+    return_source_documents=False
 )
 
 history = []
 
-# Chatbot logic using AzureOpenAI directly
+# Chatbot logic
 def sysml_chatbot(message, history):
-    # Perform retrieval
-    retriever = vectorstore.as_retriever()
-    docs = retriever.get_relevant_documents(message)
-    context = "\n\n".join(doc.page_content for doc in docs[:4])
-
-    # Compose prompt with retrieved context
-    system_prompt = "You are a helpful assistant knowledgeable in SysML. Use the context below to answer the user's question.\n\nContext:\n" + context
-
-    response = client.chat.completions.create(
-        model=AZURE_OPENAI_LLM_DEPLOYMENT,
-        messages=[
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": message}
-        ]
-    )
-
-    answer = response.choices[0].message.content
+    result = qa({"question": message, "chat_history": history})
+    answer = result["answer"]
     history.append((message, answer))
     return "", history
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## SysML Chatbot")
+    gr.Markdown("## SysML Chatbot (Azure-Powered)")
     chatbot = gr.Chatbot()
     msg = gr.Textbox(placeholder="Ask me about SysML diagrams or concepts...")
     clear = gr.Button("Clear")