nvidia-nat-langchain

import os from langchain_core.prompts import ChatPromptTemplate from langchain_nvidia_ai_endpoints import ChatNVIDIA from langchain_core.output_parsers import StrOutputParser from nvidia_nat.langchain import add_profiling_callbacks # Ensure NVIDIA_API_KEY is set in your environment # You can get one from https://build.nvidia.com/ # os.environ["NVIDIA_API_KEY"] = "nvapi-xxxxxxxxxxxxxxxxxxxxxxxx" if not os.getenv("NVIDIA_API_KEY"): print("Please set the NVIDIA_API_KEY environment variable.") exit() # 1. Define your LangChain components prompt = ChatPromptTemplate.from_messages([ ("system", "You are a helpful AI assistant."), ("user", "{input}") ]) llm = ChatNVIDIA(model="meta/llama-3-8b-instruct") output_parser = StrOutputParser() # 2. Create your LangChain chain chain = prompt | llm | output_parser # 3. Add NeMo Agent Toolkit profiling callbacks to your chain profiled_chain = add_profiling_callbacks(chain) # 4. Invoke the profiled chain response = profiled_chain.invoke({"input": "What is the capital of France?"}) print(response) # In a real NAT setup, you would typically run this via the `nat` CLI # with a YAML configuration, which automatically applies these integrations.