Spaces:
Running
Running
| import os | |
| import gradio as gr | |
| from dotenv import load_dotenv | |
| from nemoguardrails import LLMRails, RailsConfig | |
| from langchain_openai import ChatOpenAI | |
| from langchain_google_genai import ChatGoogleGenerativeAI | |
| from chain import qa_chain | |
| from vectorstore import qdrant_client | |
| load_dotenv() | |
| os.environ["TOKENIZERS_PARALLELISM"] = "false" | |
| MODEL_API_KEY = os.getenv("OPENAI_API_KEY") or os.getenv("GOOGLE_API_KEY") or "" | |
| OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL") or "http://localhost:11434/v1" | |
| MODEL_LIST = { | |
| "openai": "gpt-4o-mini", | |
| "gemini": "gemini-1.5-pro-002" | |
| } | |
| DEFAULT_MODEL = "openai" | |
| def vector_search(message): | |
| documents = qdrant_client.query(collection_name="aws_faq", query_text=message, limit=4) | |
| context = '\n'.join([doc.metadata["document"] for doc in documents]) | |
| return context | |
| def initialize_app(llm): | |
| config = RailsConfig.from_path("config") | |
| app = LLMRails(config=config, llm=llm) | |
| return app | |
| def format_messages(message, relevant_chunks): | |
| messages = [{"role": "context", "content": {"relevant_chunks": relevant_chunks}}, {"role": "user", "content": message}] | |
| return messages | |
| async def predict(message, _, model_api_key, provider, is_guardrails): | |
| if not model_api_key: | |
| return "OpenAI/Gemini API Key is required to run this demo, please enter your OpenAI API key in the settings and configs section!" | |
| if provider == "gemini": | |
| llm = ChatGoogleGenerativeAI(google_api_key=model_api_key, model=MODEL_LIST[provider]) | |
| elif provider == "openai": | |
| llm = ChatOpenAI(openai_api_key=model_api_key, model_name=MODEL_LIST[provider]) | |
| elif provider == "ollama": | |
| llm = ChatOpenAI(openai_api_key="", openai_api_base=OLLAMA_BASE_URL, model_name=MODEL_LIST[provider]) | |
| else: | |
| return "Invalid provider selected, please select a valid provider from the dropdown!" | |
| context = vector_search(message) | |
| if not is_guardrails: | |
| return qa_chain(llm, message, context) | |
| app = initialize_app(llm) | |
| response = await app.generate_async(messages=format_messages(message, context)) | |
| return response["content"] | |
| with gr.Blocks() as demo: | |
| gr.HTML("""<div style='height: 10px'></div>""") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown( | |
| """ | |
| # AWS Chatbot | Guardrails | |
| Experiment on langchain with NeMo Guardrails. | |
| """ | |
| ) | |
| with gr.Column(scale=2): | |
| with gr.Group(): | |
| with gr.Row(): | |
| guardrail = gr.Checkbox(label="Guardrails", info="Enables NeMo Guardrails",value=True, scale=1) | |
| provider = gr.Dropdown(MODEL_LIST.keys(), value=DEFAULT_MODEL, show_label=False, scale=1) | |
| model_key = gr.Textbox(placeholder="Enter your OpenAI/Gemini API key", type="password", value=MODEL_API_KEY, show_label=False, scale=3) | |
| gr.ChatInterface( | |
| predict, | |
| chatbot=gr.Chatbot(height=600, type="messages", layout="panel"), | |
| theme="soft", | |
| examples=[["How reliable is Amazon S3 with data availability ?"], ["How do I get started with EC2 Capacity Blocks ?"]], | |
| type="messages", | |
| additional_inputs=[model_key, provider, guardrail] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |