AIAlignmentChatBot / indexer.py
dsouzaJithesh's picture
Upload indexer.py
20b0da7 verified
raw
history blame
1.71 kB
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
import os
from google import genai
from google.genai import types
# Set up the Gemini API key
import os
def index_text():
os.environ["NVIDIA_API_KEY"] =os.getenv("NVIDIA_API_KEY")
nvidia_embeddings = NVIDIAEmbeddings(
model="nvidia/llama-3.2-nv-embedqa-1b-v2",
truncate="NONE"
)
vectorstore = FAISS.load_local("nvidia_faiss_index", embeddings=nvidia_embeddings,allow_dangerous_deserialization=True)
return vectorstore
def answer_query(query, history,vectorstore):
os.environ["GEMINI_API_KEY"] = os.getenv("GEMINI_API_KEY")
client = genai.Client()
RAG_TEMPLATE = """
#CONTEXT:
{context}
Use the provided context to answer the user query.
"""
retriever = vectorstore.as_retriever()
search_results = retriever.invoke(query, k=2)
context = " ".join([doc.page_content for doc in search_results])
prompt = RAG_TEMPLATE.format(context=context, query=query)
gemini_history = []
for msg in history:
# The Gemini API uses 'model' for the assistant's role
role = 'model' if msg['role'] == 'assistant' else 'user'
gemini_history.append(
types.Content(role=role, parts=[types.Part(text=msg['content'])])
)
chat = client.chats.create(
model="gemini-2.0-flash",
history=gemini_history,
config=types.GenerateContentConfig(
system_instruction=prompt)
)
response=chat.send_message(message=query)
return response.text