File size: 2,029 Bytes
a909949 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import os
from nemoguardrails import LLMRails
from nemoguardrails.actions.actions import ActionResult
from qdrant_client import QdrantClient
from dotenv import load_dotenv
load_dotenv()
COLLECTION_NAME = "nemo-docs" # Name of the collection
qdrant_client = QdrantClient(
os.getenv("QDRANT_URL"),
api_key=os.getenv("QDRANT_API_KEY"),
)
def vector_search(query, limit=4):
documents = qdrant_client.query(collection_name=COLLECTION_NAME, query_text=query, limit=limit)
context = '\n\n'.join([f"PAGE_CONTENT: {doc.metadata['document']} SOURCE: {doc.metadata['source']}" for doc in documents])
return context
def add_vectors(chunks, metadata, ids):
qdrant_client.add(
collection_name=COLLECTION_NAME,
documents=chunks,
metadata=metadata,
ids=ids
)
def prompt_template(question, context):
return f"""You are an **GitDoc AI** Chatbot, a helpful assistant that assists users with their
**NVIDIA's NeMo Guardrails** related questions.
CONTEXT INFORMATION is below.
---------------------
{context}
---------------------
RULES:
1. Only Answer the USER QUESTION using the CONTEXT INFORMATION text above.
2. Keep your answer grounded in the facts of the CONTEXT.
3. If you don't know the answer, just say that you don't know politely.
4. Should not answer any out-of-context USER QUESTION.
5. Add references only if needed in markdown format.
USER QUESTION: ```{question}```
Answer in markdown:"""
def rag(context: dict, llm) -> ActionResult:
context_updates = {}
user_message = context.get("last_user_message")
relevant_chunks = vector_search(user_message)
context_updates["relevant_chunks"] = relevant_chunks
prompt = prompt_template(user_message, relevant_chunks)
answer = llm.invoke(prompt).content
context_updates["_last_bot_prompt"] = prompt
return ActionResult(return_value=answer, context_updates=context_updates)
def init(app: LLMRails):
app.register_action(rag, "rag") |