Update app.py

app.py

@@ -0,0 +1,106 @@
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_milvus import Milvus
+from langchain.chat_models import init_chat_model
+from typing import List
+from langchain.agents.middleware import dynamic_prompt, ModelRequest
+from langchain.agents import create_agent
+from langchain_core.documents import Document
+from langchain_core.runnables import chain
+
+import gradio as gr
+import os
+import shutil # Import shutil for directory removal
+import tempfile # Import tempfile for temporary directory creation
+
+#loading data
+file_path = "PIE_Service_Rules_&_Policies.pdf"
+loader = PyPDFLoader(file_path)
+
+docs = loader.load()
+
+#splitting it
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size=1000, chunk_overlap=200, add_start_index=True
+)
+all_splits = text_splitter.split_documents(docs)
+
+#performing embeddings and storing in milvus
+embeddings = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")
+
+# Create a temporary directory for Milvus Lite
+temp_dir = tempfile.mkdtemp()
+URI = os.path.join(temp_dir, "milvus_data.db")
+
+# Explicitly remove the Milvus Lite data to ensure a clean start
+# This block is no longer needed as tempfile.mkdtemp() provides a clean directory
+# if os.path.exists(URI):
+#     if os.path.isdir(URI):
+#         shutil.rmtree(URI)
+#         print(f"Removed existing Milvus Lite data directory: {URI}")
+#     elif os.path.isfile(URI):
+#         os.remove(URI)
+#         print(f"Removed existing Milvus Lite data file: {URI}")
+
+vector_store = Milvus(
+    embedding_function=embeddings,
+    connection_args={"uri": URI},
+    index_params={"index_type": "FLAT", "metric_type": "L2"},
+    drop_old=True
+)
+
+ids = vector_store.add_documents(documents=all_splits)
+
+#Retriever
+@chain
+def retriever(query: str) -> List[Document]:
+    return vector_store.similarity_search(query, k=2)
+
+
+#model
+# from google.colab import userdata
+# key = userdata.get('Groq_Key')
+key = os.getenv('Groq_key2')
+os.environ["GROQ_API_KEY"] = key
+
+model = init_chat_model(
+    "llama-3.1-8b-instant",
+    model_provider="groq"
+)
+
+#using langchain middleware for dynamic prompts
+@dynamic_prompt
+def prompt_with_context(request: ModelRequest) -> str:
+    """Inject context into state messages."""
+    last_query = request.state["messages"][-1].text
+    retrieved_docs = vector_store.similarity_search(last_query)
+
+    docs_content = "\n\n".join(doc.page_content for doc in retrieved_docs)
+
+    system_message = (
+        "You are a helpful assistant who explain company policies to company employees. Use the following context in your response:"
+        f"\n\n{docs_content}"
+    )
+
+    return system_message
+
+
+agent = create_agent(model, tools=[], middleware=[prompt_with_context])
+
+def chat(message, history):
+
+    results = []
+    for step in agent.stream(
+        {"messages": [{"role": "user", "content": message}]},
+        stream_mode="values",
+    ):
+    # Grab the last message in the stream
+        last_message = step["messages"][-1]
+
+        # Append it to results instead of printing
+        results.append(last_message)
+    return results[1].content
+
+demo = gr.ChatInterface(fn=chat, title="PI Invent Help Assistant")
+demo.launch(debug = True)