dsouzaJithesh/AIAlignmentChatBot

1from langchain_community.vectorstores import FAISS
2from langchain_core.documents import Document
3from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
4import os 
5from google import genai
6from google.genai import types
7
8# Set up the Gemini API key
9import os
10
11def index_text():
12
13    os.environ["NVIDIA_API_KEY"] =os.getenv("NVIDIA_API_KEY")
14
15    
16    nvidia_embeddings = NVIDIAEmbeddings(
17        model="nvidia/llama-3.2-nv-embedqa-1b-v2",
18        truncate="NONE"
19    )
20    vectorstore = FAISS.load_local("nvidia_faiss_index", embeddings=nvidia_embeddings,allow_dangerous_deserialization=True)
21    return vectorstore
22
23
24def answer_query(query, history,vectorstore):
25
26    os.environ["GEMINI_API_KEY"] = os.getenv("GEMINI_API_KEY")
27    client = genai.Client()
28
29    RAG_TEMPLATE = """
30#CONTEXT:
31{context}
32Use the provided context to answer the user query.
33"""
34    retriever = vectorstore.as_retriever()
35    search_results = retriever.invoke(query, k=2)
36    context = " ".join([doc.page_content for doc in search_results])
37    prompt = RAG_TEMPLATE.format(context=context, query=query)
38
39
40    gemini_history = []
41    for msg in history:
42        # The Gemini API uses 'model' for the assistant's role
43        
44        role = 'model' if msg['role'] == 'assistant' else 'user'
45        gemini_history.append(
46            types.Content(role=role, parts=[types.Part(text=msg['content'])])
47        )
48
49    chat = client.chats.create(
50        model="gemini-2.0-flash",
51        history=gemini_history,
52        config=types.GenerateContentConfig(
53            system_instruction=prompt)
54    )
55
56    response=chat.send_message(message=query)
57    return response.text