1 line
16 KiB
JSON
1 line
16 KiB
JSON
{"repo": "bug-001/falcon-mia", "files": {"llm/query.py": {"11b235b64cb87bc44fdcd264af1e0ff727337096": {"t": 1785595499.8029358, "n": 15497, "h": "fece03988a3a4b8c", "c": "import argparse\nimport random\nimport yaml\nimport requests\nimport os\nfrom dotenv import load_dotenv\nfrom openai import OpenAI\nfrom colorama import Fore, init\nfrom abc import ABC, abstractmethod\nimport logging\nimport openai\nimport anthropic\nimport time\nfrom typing import List, Dict\n\nfrom .tools.utils import get_logger\n\nlogger = get_logger(\"query\", \"info\")\ninit()\nload_dotenv()\n\nclass ModelClient(ABC):\n @abstractmethod\n def chat_completion(self, messages, **kwargs):\n pass\n\nclass OpenAIClient(ModelClient):\n def __init__(self, api_key=None, base_url=None):\n # Allow setting base_url through parameters\n self.api_key = api_key\n self.base_url = base_url\n self.client_params = {\"api_key\": api_key}\n if base_url:\n self.client_params[\"base_url\"] = base_url\n self.client = OpenAI(**self.client_params)\n self.error_count = 0\n self.error_threshold = 5\n\n def chat_completion(self, messages, **kwargs):\n while True:\n try:\n completion = self.client.chat.completions.create(\n messages=messages,\n **kwargs\n )\n self.error_count = 0 # Reset error count after success\n ret = completion.choices[0].message.content\n if ret == None:\n logger.warning(\"Failed to get a response from the model. Details:\")\n logger.warning(completion)\n return ret\n except openai.RateLimitError as e:\n # Wait a short time and retry when rate limit is exceeded\n sleep_time = random.uniform(5, 10)\n logger.warning(f\"Credit limit exceeded, retrying in {sleep_time} seconds... Error: {e}\")\n time.sleep(sleep_time)\n except openai.APITimeoutError as e:\n # API timeout, wait and retry\n logger.warning(f\"API timeout, retrying in 3 seconds... Error: {e}\")\n time.sleep(3)\n except openai.InternalServerError as e:\n # Usually a custom error, wait and retry\n self.error_count += 1\n logger.warning(f\"Internal server error ({self.error_count}/{self.error_threshold}), retrying... Error: {e}\")\n \n if self.error_count >= self.error_threshold:\n logger.warning(\"Too many errors, recreating client...\")\n self.client.close()\n self.client = OpenAI(**self.client_params)\n self.error_count = 0\n \n time.sleep(random.randint(20, 40))\n except openai.APIStatusError as e:\n logger.warning(f\"Error during OpenAI API request: {e}\")\n return str(e)\n\nclass LocalClient(ModelClient):\n def __init__(self, base_url):\n self.base_url = base_url\n self.client = OpenAI(base_url=base_url)\n\n def chat_completion(self, messages, **kwargs):\n while True:\n try:\n completion = self.client.chat.completions.create(\n messages=messages,\n **kwargs\n )\n return completion.choices[0].message.content\n except openai.InternalServerError as e:\n if e.status_code == 502:\n print(\"Server is not connected. Retrying in 30 seconds...\")\n time.sleep(30)\n else:\n raise e\n except openai.BadRequestError as e:\n logger.warning(f\"Bad request error: {e}\")\n return str(e)\n except openai.APIConnectionError as e:\n print(\"Temporary network error. Disconnecting current client and recreating a new client...\")\n # Disconnect current client and recreate a new client\n time.sleep(5)\n self.client.close()\n self.client = OpenAI(base_url=self.base_url)\n\nclass InfinigenceClient(ModelClient):\n def __init__(self, api_key):\n self.api_key = api_key\n\n def chat_completion(self, messages, **kwargs):\n model = kwargs.get('model', 'qwen2.5-7b-instruct')\n url = f\"https://cloud.infini-ai.com/maas/{model}/nvidia/chat/completions\"\n \n payload = {\n \"model\": model,\n \"messages\": messages,\n **kwargs\n }\n headers = {\n \"Content-Type\": \"application/json\",\n \"Authorization\": f\"Bearer {self.api_key}\"\n }\n\n try:\n response = requests.post(url, json=payload, headers=headers)\n response.raise_for_status()\n return response.json()['choices'][0]['message']['content']\n except Exception as e:\n logger.warning(f\"Error during Infinigence API request: {e}\")\n return None\n\nclass AIMLClient(ModelClient):\n def __init__(self, api_key):\n self.client = OpenAI(api_key=api_key, base_url=\"https://api.aimlapi.com/v1\")\n\n def chat_completion(self, messages, **kwargs):\n while True:\n try:\n completion = self.client.chat.completions.create(\n messages=messages,\n **kwargs\n )\n return completion.choices[0].message.content\n except openai.InternalServerError as e:\n time.sleep(5)\n except Exception as e:\n logger.warning(f\"Error during AI/ML API request: {type(e)}\")\n raise e\n\nclass AnthropicClient(ModelClient):\n def __init__(self, api_key):\n self.client = anthropic.Anthropic(api_key=api_key)\n\n def chat_completion(self, messages: List[Dict], **kwargs):\n # Move the system prompt out of the messages list\n if len(messages) > 0 and messages[0]['role'] == 'system':\n system = messages.pop(0)['content']\n\n kwargs.pop('frequency_penalty') # Remove model from kwargs\n\n completion = self.client.messages.create(\n system=system,\n messages=messages,\n model=kwargs['model'],\n stop_sequences=kwargs['stop'],\n max_tokens=kwargs['max_tokens'],\n temperature=kwargs['temperature'],\n top_p=kwargs['top_p'],\n )\n return completion.content[0].text\n\nclass QueryProcessor:\n def __init__(self, query, full_output=False):\n self.query = query\n self.client = self._get_client()\n self.full_output = full_output\n\n # Set logger level based on full_output\n if not self.full_output:\n logger.setLevel(logging.WARNING) # Only output warnings and errors\n\n def _get_client(self):\n model_type = self.query[\"model_type\"]\n if model_type == \"openai\":\n # Get base_url from configuration\n base_url = self.query.get(\"base_url\")\n return OpenAIClient(\n api_key=get_api_key(\"openai\"),\n base_url=base_url\n )\n elif model_type == \"aiml\":\n return AIMLClient(api_key=get_api_key(\"aiml\"))\n elif model_type == \"local\":\n server = read_yaml(self.query['model'])\n base_url = f\"http://{server.get('host', '127.0.0.1')}:{server.get('port', 8000)}/v1\"\n self.query['model'] = server['model-tag']\n return LocalClient(base_url)\n elif model_type == \"infinigence\":\n return InfinigenceClient(api_key=get_api_key(\"infinigence\"))\n elif model_type == \"anthropic\":\n return AnthropicClient(api_key=get_api_key(\"anthropic\"))\n elif model_type == \"openrouter\":\n return OpenAIClient(api_key=get_api_key(\"openrouter\"), base_url=\"https://openrouter.ai/api/v1\")\n else:\n raise ValueError(f\"Unknown model_type: {model_type}\")\n\n def process_query(self):\n if self.query['query_type'] == 'chat':\n return self.process_chats()\n elif self.query['query_type'] == 'instruction':\n return self.process_instruction()\n else:\n raise ValueError(f\"Unsupported query type: {self.query['query_type']}\")\n\n def _get_stop_pattern(self, model_name):\n # Set default stop pattern based on model name\n model_name_lower = model_name.lower()\n if 'llama' in model_name_lower:\n return [\"[INST]\", \"[/INST]\", \"</s>\", \"[/s]\"]\n elif 'claude' in model_name_lower:\n return [\"\\n\\nHuman:\", \"\\n\\nAssistant:\"]\n elif 'mistral' in model_name_lower:\n return [\"[INST]\", \"[/INST]\"]\n elif 'vicuna' in model_name_lower:\n return [\"[INST]\", \"[/INST]\", \"</s>\", \"[/s]\"]\n else:\n # For other models, use a generic stop pattern\n return []\n\n def process_chats(self):\n chats = self.query.get('chats', [])\n\n all_responses = []\n\n for i, chat in enumerate(chats):\n logger.info(f\"Processing chat: {chat.get('name', 'Unnamed chat')}\")\n messages = chat.get('messages', [])\n updated_messages = []\n \n model_name = self.query.get('model', 'gpt-3.5-turbo')\n stop_pattern = self._get_stop_pattern(model_name)\n\n llm_responses = [] # To store LLM responses\n\n def send_to_llm(messages):\n response = self.client.chat_completion(\n messages,\n model=model_name,\n temperature=self.query.get('temperature', 0.7),\n max_tokens=self.query.get('max_tokens', 150),\n top_p=self.query.get('top_p', 1.0),\n frequency_penalty=self.query.get('frequency_penalty', 0.0),\n presence_penalty=self.query.get('presence_penalty', 0.0),\n n=self.query.get('n', 1),\n stop=self.query.get('stop', []) + stop_pattern,\n )\n if response == None:\n return \"[FAILED] The model returned None.\"\n return response\n\n for message in messages:\n if message['role'] == 'assistant' and message['content'] == '***TBA***':\n # If it's a TBA message from assistant, send previous context to LLM\n response = send_to_llm(updated_messages)\n updated_messages.append({\"role\": \"assistant\", \"content\": response})\n llm_responses.append(response)\n if self.full_output:\n print(f\"Assistant: \" + Fore.YELLOW + f\"{response}\" + Fore.RESET)\n else:\n updated_messages.append(message)\n if self.full_output:\n print(f\"{message['role'].capitalize()}: {message['content']}\")\n\n # Flush messages if the last message is not from the assistant\n if updated_messages[-1]['role'] != 'assistant':\n response = send_to_llm(updated_messages)\n updated_messages.append({\"role\": \"assistant\", \"content\": response})\n llm_responses.append(response)\n if self.full_output:\n print(f\"Assistant: \" + Fore.YELLOW + f\"{response}\" + Fore.RESET)\n\n all_responses.append(llm_responses)\n\n return all_responses\n\n def process_instruction(self) -> list:\n # TODO: Implement instruction processing\n raise NotImplementedError(\"Instruction query type is not yet supported.\")\n\ndef parse_arguments():\n parser = argparse.ArgumentParser(description=\"Large Language Model Query Tool\")\n parser.add_argument(\"-i\", \"--interactive\", action=\"store_true\", \n help=\"Run in interactive mode\")\n parser.add_argument(\"-f\", \"--full\", action=\"store_true\", \n help=\"Output full conversation (default: LLM responses only)\")\n parser.add_argument(\"-c\", \"--config\", default=os.path.join(os.path.dirname(__file__), \"examples\", \"query.yaml\"), \n help=\"Path to the YAML query file (default: query.yaml)\")\n return parser.parse_args()\n\ndef get_api_key(provider):\n key_name = f\"{provider.upper()}_API_KEY\"\n api_key = os.getenv(key_name)\n \n if not api_key:\n raise ValueError(f\"API key for {provider} not found. Please set the {key_name} environment variable in .env.\")\n \n return api_key\n\ndef read_yaml(file_path):\n with open(file_path, 'r') as file:\n return yaml.safe_load(file)\n\nclass ModelInterface:\n def __init__(self, query_config):\n self.query_config = query_config\n \n def query(self, prompt: List[Dict[str, str]], chat_name: str, query_hook_dict: dict = None) -> str:\n # Available hooks:\n # pre_hook: function(prompt: List[Dict[str, str]]) -> Dict\n # post_hook: function(llm_response: str) -> Dict\n\n config = self.query_config.copy()\n \n pre_hook_warning = None\n \n if query_hook_dict and 'pre_hook' in query_hook_dict:\n pre_hook_result = query_hook_dict['pre_hook'](prompt)\n \n # Check if pre_hook returned a dictionary with the expected format\n if isinstance(pre_hook_result, dict) and \"result\" in pre_hook_result:\n print(f\"Pre-hook information: {pre_hook_result['info']}\")\n\n # Store warning info if status is \"warning\"\n if pre_hook_result.get(\"status\") == \"warning\":\n pre_hook_warning = pre_hook_result.get(\"info\")\n \n # If status is \"filtered\", return the info message directly\n if pre_hook_result.get(\"status\") == \"filtered\":\n return [pre_hook_result.get(\"info\", \"Query rejected by security filter\")]\n \n prompt = pre_hook_result[\"result\"]\n else:\n # For backward compatibility\n prompt = pre_hook_result\n \n # If prompt is None, return early\n if prompt is None:\n return [\"Query rejected by security filter\"]\n \n chat_config = {\n \"name\": chat_name,\n \"messages\": prompt\n }\n config = dict(config)\n config['chats'] = [chat_config]\n llm_response = QueryProcessor(config).process_query()[0]\n \n # Append warning to the response if there was one\n if pre_hook_warning:\n llm_response = f\"[WARNING] {pre_hook_warning}\\n\\n{llm_response}\"\n \n if query_hook_dict and 'post_hook' in query_hook_dict:\n post_hook_result = query_hook_dict['post_hook'](llm_response)\n \n # Check if post_hook returned a dictionary with the expected format\n if isinstance(post_hook_result, dict) and \"result\" in post_hook_result:\n llm_response = post_hook_result[\"result\"]\n else:\n # For backward compatibility\n llm_response = post_hook_result\n\n return llm_response\n\nif __name__ == \"__main__\":\n args = parse_arguments()\n query = read_yaml(args.config)\n\n processor = QueryProcessor(query, args.full)\n \n if args.interactive:\n logger.info(\"Running in interactive mode.\")\n # TODO: Implement interactive mode logic\n else:\n logger.info(\"Running in non-interactive mode.\")\n print('----------')\n for s in processor.process_query()[0]:\n print(s, end='\\n----------\\n')"}}}} |