Update rag.py

rag.py

@@ -14,22 +14,18 @@ from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain_community.retrievers import BM25Retriever
 
-# Load Zephyr model
-tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
-model = AutoModelForCausalLM.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
-
-# Create HF pipeline
-hf_pipeline = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=128,
-    temperature=0.8,
-    pad_token_id=tokenizer.eos_token_id,
-)
+import google.generativeai as genai
+import os
+import google.generativeai as genai
+
+# Initialize Gemini
+genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))  # Replace this with actual key or environment-safe config
+
+
+
+model = genai.GenerativeModel("gemini-1.5-flash")
+
 
-# Wrap in LangChain LLM
-llm = HuggingFacePipeline(pipeline=hf_pipeline)
 
 # Define your RAG response function
 
@@ -80,21 +76,25 @@ def respond_rag_huggingface(message: str):
     docs = ensemble_retriever(message)
     context = "\n\n".join(doc.page_content for doc in docs)
 
-    prompt_template = ChatPromptTemplate.from_messages([
-        ("system", "you are game of thrones measter answer the given question strictly based on the context provived.if u donot know the answer reply i dont know donot give gibberish answers"),
-        ("human", """Context: {context}
+    system_message = os.environ.get("SYSTEM_MESSAGE", 
+        "You are a Game of Thrones measter. Answer the given question strictly based on the provided context. If you don't know, say 'I don't know'. Do not guess.")
+
+    prompt = f"""{system_message}
+
+Context:
+{context}
+
+Question:
+{message}
 
-        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"""
 
-        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""")
-    ])
+    response = model.generate_content(prompt)
+    return response.text
 
-    chain = prompt_template | llm
-    response = chain.invoke({"context": context, "question": message})
-    return response.content
 
 
 __all__ = ["respond_rag_huggingface"]