Update rag.py

rag.py

@@ -1,66 +1,77 @@
-from langchain.document_loaders import DirectoryLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import FAISS
-from langchain.llms import Ollama  
-db = FAISS.load_local(
-    folder_path="got_embeddings",
-    embeddings=HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2"),
-    allow_dangerous_deserialization=True  
-)
-
-from langchain.retrievers import BM25Retriever, EnsembleRetriever
-
-# Vector Store Retriever
-vector_retriever = db.as_retriever(search_kwargs={"k": 3})
-
-# Keyword Retriever (BM25)
-bm25_retriever = BM25Retriever.from_documents(texts)
-bm25_retriever.k = 2
-
-# Combine both
-ensemble_retriever = EnsembleRetriever(
-    retrievers=[vector_retriever, bm25_retriever],
-    weights=[0.6, 0.4]  # Tune based on your tests
-)
-
-# Use in ask_question()
-
-
-from langchain_community.llms import Ollama
-from langchain_core.prompts import ChatPromptTemplate
-
-def ask_question(question,temparature=0.5,num_ctx=4096,top_k=40):
-    # 1. Retrieve relevant context from your vector DB
-    docs = ensemble_retriever.get_relevant_documents(question)
-    context = "\n\n".join([doc.page_content for doc in docs])
-    
-    # 2. Create optimized prompt template
-    prompt_template = ChatPromptTemplate.from_messages([
-        ("system", "You are a Game of Thrones expert. Answer strictly based on the context."),
-        ("human", """Context: {context}
-        
-        Question: {question}
-        
-        Rules:
-        - If answer isn't in context, say "I don't know"
-        - Keep answers under 5 sentences
-        - Include book/season references when possible""")
-    ])
-    
-
-    # 3. Configure Ollama with your specific model parameters
-    llm = Ollama(
-     model="llama3:8b-instruct-q4_0",
-        temperature=temparature,
-        num_ctx=num_ctx,
-        top_k=top_k,
-        repeat_penalty=1.1,
-        stop=["<|eot_id|>"],
-        
-                         # Disable GPU entirely
-)
-    
-    # 4. Generate response
-    chain = prompt_template | llm
-    return chain.invoke({"context": context, "question": question})
+from langchain.document_loaders import DirectoryLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.llms import Ollama  
+db = FAISS.load_local(
+    folder_path="got_embeddings",
+    embeddings=HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2"),
+    allow_dangerous_deserialization=True  
+)
+
+from langchain.retrievers import BM25Retriever, EnsembleRetriever
+
+# Vector Store Retriever
+vector_retriever = db.as_retriever(search_kwargs={"k": 3})
+
+# Keyword Retriever (BM25)
+bm25_retriever = BM25Retriever.from_documents(texts)
+bm25_retriever.k = 2
+
+# Combine both
+ensemble_retriever = EnsembleRetriever(
+    retrievers=[vector_retriever, bm25_retriever],
+    weights=[0.6, 0.4]  # Tune based on your tests
+)
+
+# Use in ask_question()
+
+
+from langchain_community.llms import Ollama
+from langchain_core.prompts import ChatPromptTemplate
+
+from langchain_community.llms import Ollama
+from langchain_core.prompts import ChatPromptTemplate
+
+def respond_rag_ollama(
+    message: str,
+    history: list[tuple[str, str]],
+    system_message: str,
+    num_ctx: int = 2048,
+    num_predict: int = 128,
+    temperature: float = 0.8,
+    top_k: int = 40,
+    repeat_penalty: float = 1.1,
+    stop: list[str] | None = None,
+):
+    # 1. Retrieve relevant context from your vector DB
+    docs = ensemble_retriever.get_relevant_documents(message)
+    context = "\n\n".join(doc.page_content for doc in docs)
+
+    # 2. Build a conversational prompt
+    prompt_template = ChatPromptTemplate.from_messages([
+        ("system", system_message),
+        ("human", f"""Context: {{context}}
+
+        Question: {{question}}
+
+        Rules:
+        - If the answer isn't in the context, respond with "I don't know"
+        - Keep answers under 5 sentences
+        - Include book/season references when possible""")
+    ])
+
+    # 3. Configure the Ollama LLM with adjustable parameters
+    llm = Ollama(
+        model="llama3:8b-instruct-q4_0",
+        temperature=temperature,
+        num_ctx=num_ctx,
+        num_predict=num_predict,
+        top_k=top_k,
+        repeat_penalty=repeat_penalty,
+        stop= ["<|eot_id|>"],
+    )
+
+
+    chain = prompt_template | llm
+    yield from chain.stream_invoke({"context": context, "question": message})