| from langchain_openai import ChatOpenAI | |
| from langchain_google_genai import ChatGoogleGenerativeAI | |
| from nemo.config import vector_search | |
| OLLAMA_BASE_URL = "http://localhost:11434/v1" | |
| GROQ_BASE_URL = "https://api.groq.com/openai/v1" | |
| def initialize_llm(model_api_key, provider, model): | |
| if provider == "gemini": | |
| return ChatGoogleGenerativeAI(google_api_key=model_api_key, model=model) | |
| elif provider == "openai": | |
| return ChatOpenAI(openai_api_key=model_api_key, model_name=model) | |
| elif provider == "groq": | |
| return ChatOpenAI(openai_api_key=model_api_key, openai_api_base=GROQ_BASE_URL, model_name=model) | |
| elif provider == "ollama": | |
| return ChatOpenAI(openai_api_key="", openai_api_base=OLLAMA_BASE_URL, model_name=model) | |
| else: | |
| return None | |
| 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. | |
| Use the following pieces of context to answer the user's question: | |
| {context} | |
| USER QUESTION: ```{question}``` | |
| Answer in markdown:""" | |
| def rag_chain(llm, message): | |
| context = vector_search(message) | |
| return llm.invoke(prompt_template(message, context)).content |