from __future__ import annotations import difflib import traceback import concurrent.futures import os import concurrent.futures import time import urllib.parse import uuid from concurrent.futures import Future from datetime import timedelta from enum import Enum from pathlib import Path from typing import Callable, Generator, Any, Union, List import ast from packaging import version os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1" from huggingface_hub import SpaceStage from huggingface_hub.utils import ( build_hf_headers, ) from gradio_client import utils from importlib.metadata import distribution, PackageNotFoundError try: assert distribution('gradio_client') is not None have_gradio_client = True is_gradio_client_version7 = distribution('gradio_client').version.startswith('0.7.') except (PackageNotFoundError, AssertionError): have_gradio_client = False is_gradio_client_version7 = False from gradio_client.client import Job, DEFAULT_TEMP_DIR, Endpoint from gradio_client import Client def check_job(job, timeout=0.0, raise_exception=True, verbose=False): if timeout == 0: e = job.future._exception else: try: e = job.future.exception(timeout=timeout) except concurrent.futures.TimeoutError: # not enough time to determine if verbose: print("not enough time to determine job status: %s" % timeout) e = None if e: # raise before complain about empty response if some error hit if raise_exception: raise RuntimeError(e) else: return e # Local copy of minimal version from h2oGPT server class LangChainAction(Enum): """LangChain action""" QUERY = "Query" SUMMARIZE_MAP = "Summarize" EXTRACT = "Extract" pre_prompt_query0 = "Pay attention and remember the information below, which will help to answer the question or imperative after the context ends." prompt_query0 = "According to only the information in the document sources provided within the context above: " pre_prompt_summary0 = """""" prompt_summary0 = "Using only the information in the document sources above, write a condensed and concise summary of key results (preferably as bullet points)." pre_prompt_extraction0 = """In order to extract information, pay attention to the following text.""" prompt_extraction0 = "Using only the information in the document sources above, extract " hyde_llm_prompt0 = "Answer this question with vibrant details in order for some NLP embedding model to use that answer as better query than original question: " class GradioClient(Client): """ Parent class of gradio client To handle automatically refreshing client if detect gradio server changed """ def __init__( self, src: str, hf_token: str | None = None, max_workers: int = 40, serialize: bool = None, output_dir: str | Path | None = DEFAULT_TEMP_DIR, verbose: bool = False, auth: tuple[str, str] | None = None, h2ogpt_key: str = None, persist: bool = False, check_hash: bool = True, check_model_name: bool = False, ): """ Parameters: src: Either the name of the Hugging Face Space to load, (e.g. "abidlabs/whisper-large-v2") or the full URL (including "http" or "https") of the hosted Gradio app to load (e.g. "http://mydomain.com/app" or "https://bec81a83-5b5c-471e.gradio.live/"). hf_token: The Hugging Face token to use to access private Spaces. Automatically fetched if you are logged in via the Hugging Face Hub CLI. Obtain from: https://huggingface.co/settings/token max_workers: The maximum number of thread workers that can be used to make requests to the remote Gradio app simultaneously. serialize: Whether the client should serialize the inputs and deserialize the outputs of the remote API. If set to False, the client will pass the inputs and outputs as-is, without serializing/deserializing them. E.g. you if you set this to False, you'd submit an image in base64 format instead of a filepath, and you'd get back an image in base64 format from the remote API instead of a filepath. output_dir: The directory to save files that are downloaded from the remote API. If None, reads from the GRADIO_TEMP_DIR environment variable. Defaults to a temporary directory on your machine. verbose: Whether the client should print statements to the console. h2ogpt_key: h2oGPT key to gain access to the server persist: whether to persist the state, so repeated calls are aware of the prior user session This allows the scratch MyData to be reused, etc. This also maintains the chat_conversation history check_hash: whether to check git hash for consistency between server and client to ensure API always up to date check_model_name: whether to check the model name here (adds delays), or just let server fail (fater) """ if serialize is None: # else converts inputs arbitrarily and outputs mutate # False keeps as-is and is normal for h2oGPT serialize = False self.args = tuple([src]) self.kwargs = dict( hf_token=hf_token, max_workers=max_workers, serialize=serialize, output_dir=output_dir, verbose=verbose, h2ogpt_key=h2ogpt_key, persist=persist, check_hash=check_hash, check_model_name=check_model_name, ) if is_gradio_client_version7: self.kwargs.update(dict(auth=auth)) self.verbose = verbose self.hf_token = hf_token self.serialize = serialize self.space_id = None self.cookies: dict[str, str] = {} if is_gradio_client_version7: self.output_dir = ( str(output_dir) if isinstance(output_dir, Path) else output_dir ) else: self.output_dir = output_dir self.max_workers = max_workers self.src = src self.auth = auth self.config = None self.h2ogpt_key = h2ogpt_key self.persist = persist self.check_hash = check_hash self.check_model_name = check_model_name self.chat_conversation = [] # internal for persist=True self.server_hash = None # internal def __repr__(self): if self.config: return self.view_api(print_info=False, return_format="str") return "Not setup for %s" % self.src def __str__(self): if self.config: return self.view_api(print_info=False, return_format="str") return "Not setup for %s" % self.src def setup(self): src = self.src self.headers = build_hf_headers( token=self.hf_token, library_name="gradio_client", library_version=utils.__version__, ) # self.headers.pop('authorization', None) # else get illegal Bearer for old servers if src.startswith("http://") or src.startswith("https://"): _src = src if src.endswith("/") else src + "/" else: _src = self._space_name_to_src(src) if _src is None: raise ValueError( f"Could not find Space: {src}. If it is a private Space, please provide an hf_token." ) self.space_id = src self.src = _src state = self._get_space_state() if state == SpaceStage.BUILDING: if self.verbose: print("Space is still building. Please wait...") while self._get_space_state() == SpaceStage.BUILDING: time.sleep(2) # so we don't get rate limited by the API pass if state in utils.INVALID_RUNTIME: raise ValueError( f"The current space is in the invalid state: {state}. " "Please contact the owner to fix this." ) if self.verbose: print(f"Loaded as API: {self.src} ✔") if is_gradio_client_version7: if self.auth is not None: self._login(self.auth) self.api_url = urllib.parse.urljoin(self.src, utils.API_URL) if is_gradio_client_version7: self.sse_url = urllib.parse.urljoin(self.src, utils.SSE_URL) self.sse_data_url = urllib.parse.urljoin(self.src, utils.SSE_DATA_URL) self.ws_url = urllib.parse.urljoin( self.src.replace("http", "ws", 1), utils.WS_URL ) self.upload_url = urllib.parse.urljoin(self.src, utils.UPLOAD_URL) self.reset_url = urllib.parse.urljoin(self.src, utils.RESET_URL) self.config = self._get_config() if is_gradio_client_version7: self.protocol: str = self.config.get("protocol", "ws") self.app_version = version.parse(self.config.get("version", "2.0")) self._info = self._get_api_info() self.session_hash = str(uuid.uuid4()) if is_gradio_client_version7: from gradio_client.client import EndpointV3Compatibility endpoint_class = ( Endpoint if self.protocol.startswith("sse") else EndpointV3Compatibility ) else: endpoint_class = Endpoint if is_gradio_client_version7: self.endpoints = [ endpoint_class(self, fn_index, dependency, self.protocol) for fn_index, dependency in enumerate(self.config["dependencies"]) ] else: self.endpoints = [ endpoint_class(self, fn_index, dependency) for fn_index, dependency in enumerate(self.config["dependencies"]) ] # Create a pool of threads to handle the requests self.executor = concurrent.futures.ThreadPoolExecutor( max_workers=self.max_workers ) # Disable telemetry by setting the env variable HF_HUB_DISABLE_TELEMETRY=1 # threading.Thread(target=self._telemetry_thread).start() self.server_hash = self.get_server_hash() if is_gradio_client_version7: self.stream_open = False self.streaming_future: Future | None = None from gradio_client.utils import Message self.pending_messages_per_event: dict[str, list[Message | None]] = {} self.pending_event_ids: set[str] = set() return self def get_server_hash(self): if self.config is None: self.setup() """ Get server hash using super without any refresh action triggered Returns: git hash of gradio server """ if self.check_hash: return super().submit(api_name="/system_hash").result() else: return "GET_GITHASH" def refresh_client_if_should(self): if self.config is None: self.setup() # get current hash in order to update api_name -> fn_index map in case gradio server changed # FIXME: Could add cli api as hash server_hash = self.get_server_hash() if self.server_hash != server_hash: if self.verbose: print("server hash changed: %s %s" % (self.server_hash, server_hash), flush=True) if self.server_hash is not None and self.persist: if self.verbose: print("Failed to persist due to server hash change, only kept chat_conversation not user session hash", flush=True) # risky to persist if hash changed self.refresh_client() self.server_hash = server_hash def refresh_client(self): """ Ensure every client call is independent Also ensure map between api_name and fn_index is updated in case server changed (e.g. restarted with new code) Returns: """ if self.config is None: self.setup() kwargs = self.kwargs.copy() kwargs.pop('h2ogpt_key', None) kwargs.pop('persist', None) kwargs.pop('check_hash', None) kwargs.pop('check_model_name', None) ntrials = 3 client = None for trial in range(0, ntrials + 1): try: client = Client(*self.args, **kwargs) except ValueError as e: if trial >= ntrials: raise else: if self.verbose: print("Trying refresh %d/%d %s" % (trial, ntrials - 1, str(e))) trial += 1 time.sleep(10) if client is None: raise RuntimeError("Failed to get new client") session_hash0 = self.session_hash if self.persist else None for k, v in client.__dict__.items(): setattr(self, k, v) if session_hash0: # keep same system hash in case server API only changed and not restarted self.session_hash = session_hash0 if self.verbose: print("Hit refresh_client(): %s %s" % (self.session_hash, session_hash0)) # ensure server hash also updated self.server_hash = self.get_server_hash() def clone(self): if self.config is None: self.setup() client = GradioClient("") for k, v in self.__dict__.items(): setattr(client, k, v) client.reset_session() client.executor = concurrent.futures.ThreadPoolExecutor( max_workers=self.max_workers ) client.endpoints = [ Endpoint(client, fn_index, dependency) for fn_index, dependency in enumerate(client.config["dependencies"]) ] # transfer internals in case used client.server_hash = self.server_hash client.chat_conversation = self.chat_conversation return client def submit( self, *args, api_name: str | None = None, fn_index: int | None = None, result_callbacks: Callable | list[Callable] | None = None, ) -> Job: if self.config is None: self.setup() # Note predict calls submit try: self.refresh_client_if_should() job = super().submit(*args, api_name=api_name, fn_index=fn_index) except Exception as e: print("Hit e=%s\n\n%s" % (str(e), traceback.format_exc()), flush=True) # force reconfig in case only that self.refresh_client() job = super().submit(*args, api_name=api_name, fn_index=fn_index) # see if immediately failed e = check_job(job, timeout=0.01, raise_exception=False) if e is not None: print( "GR job failed: %s %s" % (str(e), "".join(traceback.format_tb(e.__traceback__))), flush=True, ) # force reconfig in case only that self.refresh_client() job = super().submit(*args, api_name=api_name, fn_index=fn_index) e2 = check_job(job, timeout=0.1, raise_exception=False) if e2 is not None: print( "GR job failed again: %s\n%s" % (str(e2), "".join(traceback.format_tb(e2.__traceback__))), flush=True, ) return job def question(self, instruction, *args, **kwargs) -> str: """ Prompt LLM (direct to LLM with instruct prompting required for instruct models) and get response """ kwargs["instruction"] = kwargs.get("instruction", instruction) kwargs["langchain_action"] = LangChainAction.QUERY.value kwargs["langchain_mode"] = 'LLM' ret = '' for response, texts_out in self.query_or_summarize_or_extract(*args, **kwargs): ret = response return ret def question_stream(self, instruction, *args, **kwargs) -> str: """ Prompt LLM (direct to LLM with instruct prompting required for instruct models) and get response """ kwargs["instruction"] = kwargs.get("instruction", instruction) kwargs["langchain_action"] = LangChainAction.QUERY.value kwargs["langchain_mode"] = 'LLM' ret = yield from self.query_or_summarize_or_extract(*args, **kwargs) return ret def query(self, query, *args, **kwargs) -> str: """ Search for documents matching a query, then ask that query to LLM with those documents """ kwargs["instruction"] = kwargs.get("instruction", query) kwargs["langchain_action"] = LangChainAction.QUERY.value ret = '' for response, texts_out in self.query_or_summarize_or_extract(*args, **kwargs): ret = response return ret def query_stream(self, query, *args, **kwargs) -> Generator[tuple[str | list[str], list[str]], None, None]: """ Search for documents matching a query, then ask that query to LLM with those documents """ kwargs["instruction"] = kwargs.get("instruction", query) kwargs["langchain_action"] = LangChainAction.QUERY.value ret = yield from self.query_or_summarize_or_extract(*args, **kwargs) return ret def summarize(self, *args, query=None, focus=None, **kwargs) -> str: """ Search for documents matching a focus, then ask a query to LLM with those documents If focus "" or None, no similarity search is done and all documents (up to top_k_docs) are used """ kwargs["prompt_summary"] = kwargs.get("prompt_summary", query or prompt_summary0) kwargs["instruction"] = kwargs.get('instruction', focus) kwargs["langchain_action"] = LangChainAction.SUMMARIZE_MAP.value ret = '' for response, texts_out in self.query_or_summarize_or_extract(*args, **kwargs): ret = response return ret def summarize_stream(self, *args, query=None, focus=None, **kwargs) -> str: """ Search for documents matching a focus, then ask a query to LLM with those documents If focus "" or None, no similarity search is done and all documents (up to top_k_docs) are used """ kwargs["prompt_summary"] = kwargs.get("prompt_summary", query or prompt_summary0) kwargs["instruction"] = kwargs.get('instruction', focus) kwargs["langchain_action"] = LangChainAction.SUMMARIZE_MAP.value ret = yield from self.query_or_summarize_or_extract(*args, **kwargs) return ret def extract(self, *args, query=None, focus=None, **kwargs) -> list[str]: """ Search for documents matching a focus, then ask a query to LLM with those documents If focus "" or None, no similarity search is done and all documents (up to top_k_docs) are used """ kwargs["prompt_extraction"] = kwargs.get("prompt_extraction", query or prompt_extraction0) kwargs["instruction"] = kwargs.get('instruction', focus) kwargs["langchain_action"] = LangChainAction.EXTRACT.value ret = '' for response, texts_out in self.query_or_summarize_or_extract(*args, **kwargs): ret = response return ret def extract_stream(self, *args, query=None, focus=None, **kwargs) -> list[str]: """ Search for documents matching a focus, then ask a query to LLM with those documents If focus "" or None, no similarity search is done and all documents (up to top_k_docs) are used """ kwargs["prompt_extraction"] = kwargs.get("prompt_extraction", query or prompt_extraction0) kwargs["instruction"] = kwargs.get('instruction', focus) kwargs["langchain_action"] = LangChainAction.EXTRACT.value ret = yield from self.query_or_summarize_or_extract(*args, **kwargs) return ret def query_or_summarize_or_extract(self, h2ogpt_key: str = None, instruction: str = "", text: list[str] | str | None = None, file: list[str] | str | None = None, url: list[str] | str | None = None, embed: bool = True, chunk: bool = True, chunk_size: int = 512, langchain_mode: str = None, langchain_action: str | None = None, langchain_agents: List[str] = [], top_k_docs: int = 10, document_choice: Union[str, List[str]] = "All", document_subset: str = "Relevant", document_source_substrings: Union[str, List[str]] = [], document_source_substrings_op: str = 'and', document_content_substrings: Union[str, List[str]] = [], document_content_substrings_op: str = 'and', system_prompt: str | None = '', pre_prompt_query: str | None = pre_prompt_query0, prompt_query: str | None = prompt_query0, pre_prompt_summary: str | None = pre_prompt_summary0, prompt_summary: str | None = prompt_summary0, pre_prompt_extraction: str | None = pre_prompt_extraction0, prompt_extraction: str | None = prompt_extraction0, hyde_llm_prompt: str | None = hyde_llm_prompt0, model: str | int | None = None, stream_output: bool = False, do_sample: bool = False, temperature: float = 0.0, top_p: float = 0.75, top_k: int = 40, repetition_penalty: float = 1.07, penalty_alpha: float = 0.0, max_time: int = 360, max_new_tokens: int = 1024, add_search_to_context: bool = False, chat_conversation: list[tuple[str, str]] | None = None, text_context_list: list[str] | None = None, docs_ordering_type: str | None = None, min_max_new_tokens: int = 512, max_input_tokens: int = -1, max_total_input_tokens: int = -1, docs_token_handling: str = "split_or_merge", docs_joiner: str = "\n\n", hyde_level: int = 0, hyde_template: str = None, hyde_show_only_final: bool = True, doc_json_mode: bool = False, asserts: bool = False, ) -> Generator[tuple[str | list[str], list[str]], None, None]: """ Query or Summarize or Extract using h2oGPT Args: instruction: Query for LLM chat. Used for similarity search For query, prompt template is: "{pre_prompt_query} \"\"\" {content} \"\"\" {prompt_query}{instruction}" If added to summarization, prompt template is "{pre_prompt_summary} \"\"\" {content} \"\"\" Focusing on {instruction}, {prompt_summary}" text: textual content or list of such contents file: a local file to upload or files to upload url: a url to give or urls to use embed: whether to embed content uploaded langchain_mode: "LLM" to talk to LLM with no docs, "MyData" for personal docs, "UserData" for shared docs, etc. langchain_action: Action to take, "Query" or "Summarize" or "Extract" langchain_agents: Which agents to use, if any top_k_docs: number of document parts. When doing query, number of chunks When doing summarization, not related to vectorDB chunks that are not used E.g. if PDF, then number of pages chunk: whether to chunk sources for document Q/A chunk_size: Size in characters of chunks document_choice: Which documents ("All" means all) -- need to use upload_api API call to get server's name if want to select document_subset: Type of query, see src/gen.py document_source_substrings: See gen.py document_source_substrings_op: See gen.py document_content_substrings: See gen.py document_content_substrings_op: See gen.py system_prompt: pass system prompt to models that support it. If 'auto' or None, then use automatic version If '', then use no system prompt (default) pre_prompt_query: Prompt that comes before document part prompt_query: Prompt that comes after document part pre_prompt_summary: Prompt that comes before document part None makes h2oGPT internally use its defaults E.g. "In order to write a concise single-paragraph or bulleted list summary, pay attention to the following text" prompt_summary: Prompt that comes after document part None makes h2oGPT internally use its defaults E.g. "Using only the text above, write a condensed and concise summary of key results (preferably as bullet points):\n" i.e. for some internal document part fstring, the template looks like: template = "%s \"\"\" %s \"\"\" %s" % (pre_prompt_summary, fstring, prompt_summary) hyde_llm_prompt: hyde prompt for first step when using LLM h2ogpt_key: Access Key to h2oGPT server (if not already set in client at init time) model: base_model name or integer index of model_lock on h2oGPT server None results in use of first (0th index) model in server to get list of models do client.list_models() pre_prompt_extraction: Same as pre_prompt_summary but for when doing extraction prompt_extraction: Same as prompt_summary but for when doing extraction do_sample: see src/gen.py temperature: see src/gen.py top_p: see src/gen.py top_k: see src/gen.py repetition_penalty: see src/gen.py penalty_alpha: see src/gen.py max_new_tokens: see src/gen.py min_max_new_tokens: see src/gen.py max_input_tokens: see src/gen.py max_total_input_tokens: see src/gen.py stream_output: Whether to stream output do_sample: whether to sample max_time: how long to take add_search_to_context: Whether to do web search and add results to context chat_conversation: List of tuples for (human, bot) conversation that will be pre-appended to an (instruction, None) case for a query text_context_list: List of strings to add to context for non-database version of document Q/A for faster handling via API etc. Forces LangChain code path and uses as many entries in list as possible given max_seq_len, with first assumed to be most relevant and to go near prompt. docs_ordering_type: By default uses 'reverse_ucurve_sort' for optimal retrieval max_input_tokens: Max input tokens to place into model context for each LLM call -1 means auto, fully fill context for query, and fill by original document chunk for summarization >=0 means use that to limit context filling to that many tokens max_total_input_tokens: like max_input_tokens but instead of per LLM call, applies across all LLM calls for single summarization/extraction action max_new_tokens: Maximum new tokens min_max_new_tokens: minimum value for max_new_tokens when auto-adjusting for content of prompt, docs, etc. docs_token_handling: 'chunk' means fill context with top_k_docs (limited by max_input_tokens or model_max_len) chunks for query or top_k_docs original document chunks summarization None or 'split_or_merge' means same as 'chunk' for query, while for summarization merges documents to fill up to max_input_tokens or model_max_len tokens docs_joiner: string to join lists of text when doing split_or_merge. None means '\n\n' hyde_level: 0-3 for HYDE. 0 uses just query to find similarity with docs 1 uses query + pure LLM response to find similarity with docs 2: uses query + LLM response using docs to find similarity with docs 3+: etc. hyde_template: see src/gen.py hyde_show_only_final: see src/gen.py doc_json_mode: see src/gen.py asserts: whether to do asserts to ensure handling is correct Returns: summary/answer: str or extraction List[str] """ if self.config is None: self.setup() if self.persist: client = self else: client = self.clone() h2ogpt_key = h2ogpt_key or self.h2ogpt_key client.h2ogpt_key = h2ogpt_key self.check_model(model) # chunking not used here # MyData specifies scratch space, only persisted for this individual client call langchain_mode = langchain_mode or "MyData" loaders = tuple([None, None, None, None, None, None]) doc_options = tuple([langchain_mode, chunk, chunk_size, embed]) asserts |= bool(os.getenv("HARD_ASSERTS", False)) if ( text and isinstance(text, list) and not file and not url and not text_context_list ): # then can do optimized text-only path text_context_list = text text = None res = [] if text: t0 = time.time() res = client.predict( text, *doc_options, *loaders, h2ogpt_key, api_name="/add_text" ) t1 = time.time() print("upload text: %s" % str(timedelta(seconds=t1 - t0)), flush=True) if asserts: assert res[0] is None assert res[1] == langchain_mode assert "user_paste" in res[2] assert res[3] == "" if file: # upload file(s). Can be list or single file # after below call, "file" replaced with remote location of file _, file = client.predict(file, api_name="/upload_api") res = client.predict( file, *doc_options, *loaders, h2ogpt_key, api_name="/add_file_api" ) if asserts: assert res[0] is None assert res[1] == langchain_mode assert os.path.basename(file) in res[2] assert res[3] == "" if url: res = client.predict( url, *doc_options, *loaders, h2ogpt_key, api_name="/add_url" ) if asserts: assert res[0] is None assert res[1] == langchain_mode assert url in res[2] assert res[3] == "" assert res[4] # should have file name or something similar if res and not res[4] and "Exception" in res[2]: print("Exception: %s" % res[2], flush=True) # ask for summary, need to use same client if using MyData api_name = "/submit_nochat_api" # NOTE: like submit_nochat but stable API for string dict passing pre_prompt_summary = pre_prompt_summary \ if langchain_action == LangChainAction.SUMMARIZE_MAP.value \ else pre_prompt_extraction prompt_summary = prompt_summary \ if langchain_action == LangChainAction.SUMMARIZE_MAP.value \ else prompt_extraction kwargs = dict( h2ogpt_key=h2ogpt_key, instruction=instruction, langchain_mode=langchain_mode, langchain_action=langchain_action, # uses full document, not vectorDB chunks langchain_agents=langchain_agents, top_k_docs=top_k_docs, document_choice=document_choice, document_subset=document_subset, document_source_substrings=document_source_substrings, document_source_substrings_op=document_source_substrings_op, document_content_substrings=document_content_substrings, document_content_substrings_op=document_content_substrings_op, system_prompt=system_prompt, pre_prompt_query=pre_prompt_query, prompt_query=prompt_query, pre_prompt_summary=pre_prompt_summary, prompt_summary=prompt_summary, hyde_llm_prompt=hyde_llm_prompt, visible_models=model, stream_output=stream_output, do_sample=do_sample, temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty, penalty_alpha=penalty_alpha, max_time=max_time, max_new_tokens=max_new_tokens, add_search_to_context=add_search_to_context, chat_conversation=chat_conversation if chat_conversation else self.chat_conversation, text_context_list=text_context_list, docs_ordering_type=docs_ordering_type, min_max_new_tokens=min_max_new_tokens, max_input_tokens=max_input_tokens, max_total_input_tokens=max_total_input_tokens, docs_token_handling=docs_token_handling, docs_joiner=docs_joiner, hyde_level=hyde_level, hyde_template=hyde_template, hyde_show_only_final=hyde_show_only_final, doc_json_mode=doc_json_mode, ) # in case server changed, update in case clone() self.server_hash = client.server_hash # ensure can fill conversation self.chat_conversation.append((instruction, None)) # get result trials = 3 for trial in range(trials): try: if not stream_output: res = client.predict( str(dict(kwargs)), api_name=api_name, ) # in case server changed, update in case clone() self.server_hash = client.server_hash res = ast.literal_eval(res) response = res["response"] if langchain_action != LangChainAction.EXTRACT.value: response = response.strip() else: response = [r.strip() for r in ast.literal_eval(response)] sources = res["sources"] scores_out = [x["score"] for x in sources] texts_out = [x["content"] for x in sources] if asserts: if text and not file and not url: assert any( text[:cutoff] == texts_out for cutoff in range(len(text)) ) assert len(texts_out) == len(scores_out) yield response, texts_out self.chat_conversation[-1] = (instruction, response) else: job = client.submit(str(dict(kwargs)), api_name=api_name) text0 = "" response = "" texts_out = [] while not job.done(): if job.communicator.job.latest_status.code.name == "FINISHED": break e = check_job(job, timeout=0, raise_exception=False) if e is not None: break outputs_list = job.communicator.job.outputs if outputs_list: res = job.communicator.job.outputs[-1] res_dict = ast.literal_eval(res) response = res_dict["response"] # keeps growing text_chunk = response[len(text0):] # only keep new stuff if not text_chunk: time.sleep(0.001) continue text0 = response assert text_chunk, "must yield non-empty string" yield text_chunk, texts_out time.sleep( 0.1 ) # let LLM deliver larger chunks, don't need to get every token output immediately # Get final response (if anything left), but also get the actual references (texts_out), above is empty. res_all = job.outputs() if len(res_all) > 0: # 0.1 slightly longer than 0.02 in open source check_job(job, timeout=0.1, raise_exception=True) res = res_all[-1] res_dict = ast.literal_eval(res) response = res_dict["response"] sources = res_dict["sources"] texts_out = [x["content"] for x in sources] yield response[len(text0):], texts_out self.chat_conversation[-1] = (instruction, response[len(text0):]) else: # 1.0 slightly longer than 0.3 in open source check_job(job, timeout=1.0, raise_exception=True) yield response[len(text0):], texts_out self.chat_conversation[-1] = (instruction, response[len(text0):]) break except Exception as e: print( "h2oGPT predict failed: %s %s" % (str(e), "".join(traceback.format_tb(e.__traceback__))), flush=True, ) if trial == trials - 1: raise else: print("trying again: %s" % trial, flush=True) time.sleep(1 * trial) finally: # in case server changed, update in case clone() self.server_hash = client.server_hash def check_model(self, model): if model != 0 and self.check_model_name: valid_llms = self.list_models() if ( isinstance(model, int) and model >= len(valid_llms) or isinstance(model, str) and model not in valid_llms ): did_you_mean = "" if isinstance(model, str): alt = difflib.get_close_matches(model, valid_llms, 1) if alt: did_you_mean = f"\nDid you mean {repr(alt[0])}?" raise RuntimeError( f"Invalid llm: {repr(model)}, must be either an integer between " f"0 and {len(valid_llms) - 1} or one of the following values: {valid_llms}.{did_you_mean}" ) def get_models_full(self) -> list[dict[str, Any]]: """ Full model info in list if dict """ if self.config is None: self.setup() return ast.literal_eval(self.predict(api_name="/model_names")) def list_models(self) -> list[str]: """ Model names available from endpoint """ if self.config is None: self.setup() return [x['base_model'] for x in ast.literal_eval(self.predict(api_name="/model_names"))]