Spaces:
Sleeping
Sleeping
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 |