from typing import List, Tuple, Dict, TypedDict, Optional, Any import os import gradio as gr from langchain_core.language_models.llms import LLM from langchain_openai.chat_models import ChatOpenAI from langchain_aws import ChatBedrock import boto3 from ask_candid.base.config.rest import OPENAI from ask_candid.base.config.models import Name2Endpoint from ask_candid.base.config.data import ALL_INDICES from ask_candid.utils import format_chat_ag_response from ask_candid.chat import run_chat ROOT = os.path.dirname(os.path.abspath(__file__)) BUCKET = "candid-data-science-reporting" PREFIX = "Assistant" class LoggedComponents(TypedDict): context: List[gr.components.Component] found_helpful: gr.components.Component will_recommend: gr.components.Component comments: gr.components.Component email: gr.components.Component def select_foundation_model(model_name: str, max_new_tokens: int) -> LLM: if model_name == "gpt-4o": llm = ChatOpenAI( model_name=Name2Endpoint[model_name], max_tokens=max_new_tokens, api_key=OPENAI["key"], temperature=0.0, streaming=True, ) elif model_name in {"claude-3.5-haiku", "llama-3.1-70b-instruct", "mistral-large", "mixtral-8x7B"}: llm = ChatBedrock( client=boto3.client("bedrock-runtime"), model=Name2Endpoint[model_name], max_tokens=max_new_tokens, temperature=0.0 ) else: raise gr.Error(f"Base model `{model_name}` is not supported") return llm def execute( thread_id: str, user_input: Dict[str, Any], history: List[Dict], model_name: str, max_new_tokens: int, indices: Optional[List[str]] = None, ): return run_chat( thread_id=thread_id, user_input=user_input, history=history, llm=select_foundation_model(model_name=model_name, max_new_tokens=max_new_tokens), indices=indices ) def build_rag_chat() -> Tuple[LoggedComponents, gr.Blocks]: with gr.Blocks(theme=gr.themes.Soft(), title="Chat") as demo: gr.Markdown( """
Please read the guide to get started.