{"repo": "autogenai/easy-problems-that-llms-get-wrong", "files": {"llm_service.py": {"9e1f52b0dc5c79f8cef52b40aab9ffb0ceafbd5c": {"t": 1785674526.54705, "n": 16555, "h": "a04314cf436fef53", "c": "# %%\nimport os\nimport json\nimport time\nimport boto3\nimport asyncio\nimport litellm\nimport requests\nimport subprocess\nfrom decouple import AutoConfig\nfrom concurrent.futures import ThreadPoolExecutor\n\nconfig = AutoConfig(search_path=\".env\")\n\nos.environ[\"OPENAI_API_KEY\"] = config(\"OPENAI_API_KEY\")\nos.environ[\"COHERE_API_KEY\"] = config(\"COHERE_API_KEY\")\nos.environ[\"MISTRAL_API_KEY\"] = config(\"MISTRAL_API_KEY\")\nos.environ[\"ANTHROPIC_API_KEY\"] = config(\"ANTHROPIC_API_KEY\")\nos.environ[\"AWS_ACCESS_KEY_ID\"] = config(\"AWS_ACCESS_KEY_ID\")\nos.environ[\"AWS_SECRET_ACCESS_KEY\"] = config(\"AWS_SECRET_ACCESS_KEY\")\nos.environ[\"AWS_REGION_NAME\"] = \"us-west-2\"\nos.environ[\"AZURE_AI_KEYS\"] = config(\"AZURE_AI_KEYS\")\nos.environ[\"VERTEX_PROJECT\"] = config(\"VERTEX_PROJECT\")\nos.environ[\"VERTEX_LOCATION\"] = config(\"VERTEX_LOCATION\")\nos.environ[\"GOOGLE_AI_STUDIO\"] = config(\"GOOGLE_AI_STUDIO\")\n\nlitellm.vertex_project = config(\"VERTEX_PROJECT\")\nlitellm.vertex_location = config(\"VERTEX_LOCATION\")\nlitellm.set_verbose = True\n\n\n# %%\ndef openai_message_parse(response: dict):\n    try:\n        messages = [m[\"message\"][\"content\"] for m in response[\"choices\"]]\n    except Exception as e:\n        if response == {}:\n            return \"None\"\n        else:\n            raise ValueError(\n                f\"Response does not contain the expected key 'choices'. Error: {e}. Response: {response}\"\n            )\n    if len(messages) == 1:\n        messages = messages[0]\n    return messages\n\n\ndef anthropic_message_parse(response: dict):\n    try:\n        messages = response[\"content\"][0][\"text\"]\n    except Exception as e:\n        if response == {}:\n            return \"None\"\n        else:\n            raise ValueError(\n                f\"Response does not contain the expected key 'choices'. Error: {e}. Response: {response}\"\n            )\n    if len(messages) == 1:\n        messages = messages[0]\n    return messages\n\n\ndef google_message_parse(response: dict | list):\n    try:\n        if isinstance(response, dict):\n            text_response = response[\"candidates\"][0][\"content\"][\"parts\"][0][\"text\"]\n        else:\n            text_response = \"\".join(\n                [\n                    line[\"candidates\"][0][\"content\"][\"parts\"][0][\"text\"]\n                    for line in response\n                    if \"candidates\" in line and \"content\" in line[\"candidates\"][0]\n                ]\n            )\n        return text_response\n    except Exception as e:\n        if response in [{}, []]:\n            return \"None\"\n        else:\n            raise ValueError(\n                f\"Response does not contain the expected key 'candidates'. Error: {e}. Response: {response}\"\n            )\n\n\ndef message_parse(response: dict, model: str):\n    if response == {}:\n        return \"\"\n    if str(response)[:17] == \"ModelResponse(id=\":\n        messages = openai_message_parse(response)\n    elif model in [\"gpt-4-turbo-preview\", \"gpt-4-turbo\", \"gpt-4o\", \"o1-preview\"]:\n        messages = openai_message_parse(response)\n    elif \"azure\" in model or \"mistral\" in model or \"mixtral\" in model:\n        messages = openai_message_parse(response)\n    elif \"gemini\" in model:\n        messages = google_message_parse(response)\n    elif \"llama\" in model:\n        messages = response\n    elif \"claude\" in model:\n        messages = anthropic_message_parse(response)\n    else:\n        print(f\"!!!!!! Could not determine model in `message_parse`: '{model}'\")\n        messages = \"\"\n    return messages\n\n\ndef litellm_service():\n    return litellm\n\n\n# %%\n## Test LiteLLM Service\n# litellm_query = litellm_service()\n\n# messages = [{ \"content\": \"Write a sentence where every word starts with the letter A.\",\"role\": \"user\"}]\n\n# models = [\"gpt-4-turbo-preview\", \"meta.llama3-70b-instruct-v1:0\", \"command-r-plus\", \"mistral/mistral-large-latest\", \"mistral/open-mixtral-8x22b\", \"claude-3-opus-20240229\", \"claude-3-5-sonnet-20240620\", \"vertex_ai/gemini-1.5-pro\", \"vertex_ai/gemini-1.0-pro\"]\n\n# model = \"command-r-plus\"\n\n# response = litellm_query.completion(model=model, messages=messages, num_retries=2)\n# print(response)\n\n# message_parse(response, model)\n\n\n# %%\nclass custom_llm_service:\n    def __init__(self):\n        pass\n\n    def openai_query(\n        self,\n        messages: list,\n        model=\"gpt-4-turbo-preview\",\n        max_tokens=1000,\n        temperature=0,\n        n=1,\n        stream=False,\n    ):\n        \"\"\"\n        \"messages\": [\n            {\n                \"role\": \"user\",\n                \"content\": \"What is the meaning of life?\"\n            },\n        ]\n        \"\"\"\n        response = requests.post(\n            \"https://api.openai.com/v1/chat/completions\",\n            headers={\n                \"Authorization\": f\"Bearer {config('OPENAI_API_KEY')}\",\n                \"Content-Type\": \"application/json\",\n            },\n            json={\n                \"model\": model,\n                \"messages\": messages,\n                \"max_completion_tokens\": max_tokens,\n                \"temperature\": temperature if model not in [\"o1-preview\"] else 1,\n                \"n\": n,\n                \"stream\": stream,\n            },\n            timeout=300,\n        )\n        return response.json()\n\n    def anthropic_query(\n        self,\n        messages: list,\n        model=\"claude-3-5-sonnet-20240620\",\n        max_tokens=1000,\n        temperature=0,\n        n=1,\n        stream=False,\n    ):\n        \"\"\"\n        \"messages\": [\n            {\n                \"role\": \"user\",\n                \"content\": \"What is the meaning of life?\"\n            },\n        ]\n        \"\"\"\n        if n > 1:\n            print('The \"n\" variable is not supported by Anthropic')\n        response = requests.post(\n            \"https://api.anthropic.com/v1/messages\",\n            headers={\n                \"x-api-key\": f\"{config('ANTHROPIC_API_KEY')}\",\n                \"anthropic-version\": \"2023-06-01\",\n                \"Content-Type\": \"application/json\",\n            },\n            json={\n                \"model\": model,\n                \"messages\": messages,\n                \"max_tokens\": max_tokens,\n                \"temperature\": temperature,\n                \"stream\": stream,\n            },\n        )\n        return response.json()\n\n    def bedrock_query(\n        self,\n        messages: list,\n        model=\"meta.llama3-8b-instruct-v1:0\",\n        max_tokens=2048,\n        temperature=0,\n        n=1,\n        stream=False,\n    ):\n        # Initialize a session using your AWS credentials\n        session = boto3.Session(\n            aws_access_key_id=os.environ[\"AWS_ACCESS_KEY_ID\"],\n            aws_secret_access_key=os.environ[\"AWS_SECRET_ACCESS_KEY\"],\n            region_name=\"us-west-2\",\n        )\n\n        # Initialize the Bedrock client\n        bedrock_client = session.client(\"bedrock-runtime\", region_name=\"us-west-2\")\n\n        # Embed the prompt in Llama 3's instruction format.\n        formatted_prompt = f\"\"\"\n        <|begin_of_text|>\n        <|start_header_id|>user<|end_header_id|>\n        {messages[0]['content']}\n        <|eot_id|>\n        <|start_header_id|>assistant<|end_header_id|>\n        \"\"\"\n\n        # Format the request payload using the model's native structure.\n        native_request = {\n            \"prompt\": formatted_prompt,\n            \"max_gen_len\": max_tokens,\n            \"temperature\": temperature,\n            \"stream\": stream,\n        }\n\n        # Convert the native request to JSON.\n        request = json.dumps(native_request)\n\n        try:\n            # Invoke the model with the request.\n            response = bedrock_client.invoke_model(modelId=model, body=request)\n            # Decode the response body.\n            model_response = json.loads(response[\"body\"].read())\n\n            # Extract and print the response text.\n            response_text = model_response[\"generation\"]\n            # print(response_text)\n            return response_text\n\n        except Exception as e:\n            print(f\"ERROR: Can't invoke '{model}'. Reason: {e}\")\n\n    def azure_query(\n        self,\n        messages: list,\n        model=\"Meta-Llama-3-1-405B-Instruct-jjo.eastus.models.ai.azure.com\",\n        max_tokens=1000,\n        temperature=0,\n        n=1,\n        stream=False,\n    ):\n        ###\n        stream = False  ###!!!!\n        ###\n\n        azure_ai_keys = json.loads(os.environ[\"AZURE_AI_KEYS\"])\n        url = f\"https://{model}/v1/chat/completions\"\n        headers = {\n            \"Content-Type\": \"application/json\",\n            \"Authorization\": azure_ai_keys[model],\n        }\n        headers.update({\"Accept\": \"text/event-stream\"}) if stream else None\n\n        data = {\n            \"messages\": messages,\n            \"max_tokens\": max_tokens,\n            \"temperature\": temperature,\n            \"stream\": stream,\n        }\n\n        response = requests.post(url, headers=headers, json=data)\n        return response.json()\n\n    def google_vertex_query(\n        self,\n        messages: list,\n        model=\"gemini-experimental\",\n        max_tokens=8192,\n        temperature=0,\n        n=1,\n        stream=False,\n    ):\n        # Define the request payload\n        \"\"\"\n        \"messages\": [\n            {\n                \"role\": \"user\",\n                \"content\": \"What is the meaning of life?\"\n            },\n        ]\n        \"\"\"\n        messages_google_format = [\n            {\"role\": message[\"role\"], \"parts\": [{\"text\": message[\"content\"]}]}\n            for message in messages\n        ]\n        payload = {\n            \"contents\": messages_google_format,\n            \"generationConfig\": {\n                \"maxOutputTokens\": max_tokens,\n                \"temperature\": temperature,\n                \"topP\": 0.95,\n            },\n        }\n\n        # Project and API details\n        PROJECT_ID = config(\"VERTEX_PROJECT\")\n        LOCATION_ID = config(\"VERTEX_LOCATION\")\n        API_ENDPOINT = f'{config(\"VERTEX_LOCATION\")}-aiplatform.googleapis.com'\n        MODEL_ID = model\n\n        # Construct the URL\n        url = f\"https://{API_ENDPOINT}/v1/projects/{PROJECT_ID}/locations/{LOCATION_ID}/publishers/google/models/{MODEL_ID}:streamGenerateContent\"\n\n        # Get the access token\n        def get_access_token():\n            token = (\n                subprocess.check_output(\"gcloud auth print-access-token\", shell=True)\n                .decode()\n                .strip()\n            )\n            return token\n\n        headers = {\n            \"Content-Type\": \"application/json\",\n            \"Authorization\": f\"Bearer {get_access_token()}\",\n        }\n\n        response = requests.post(url, headers=headers, data=json.dumps(payload))\n        return response.json()\n\n    def google_ai_studio_query(\n        self,\n        messages: list,\n        model=\"gemini-1.5-pro-exp-0801\",\n        max_tokens=8192,\n        temperature=0,\n        n=1,\n        stream=False,\n    ):\n        \"\"\"\n        \"messages\": [\n            {\n                \"role\": \"user\",\n                \"content\": \"What is the meaning of life?\"\n            },\n        ]\n        \"\"\"\n        messages_google_format = [\n            {\"role\": message[\"role\"], \"parts\": [{\"text\": message[\"content\"]}]}\n            for message in messages\n        ]\n\n        API_KEY = config(\"GOOGLE_AI_STUDIO\")\n        url = f\"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={API_KEY}\"\n\n        headers = {\n            \"Content-Type\": \"application/json\",\n        }\n\n        data = {\n            \"contents\": messages_google_format,\n            \"generationConfig\": {\n                \"temperature\": temperature,\n                \"maxOutputTokens\": max_tokens,\n                \"responseMimeType\": \"text/plain\",\n            },\n        }\n\n        response = requests.post(url, headers=headers, json=data)\n        response_json = response.json()\n        if \"error\" in response_json:\n            raise ValueError(f\"Error in google ai studio response: {response_json}\")\n        return response_json\n\n    def mistral_query(\n        self,\n        messages: list,\n        model=\"mistral-large-latest\",\n        max_tokens=1000,\n        temperature=0,\n        n=1,\n        stream=False,\n    ):\n        payload = {\n            \"model\": model,\n            \"messages\": messages,\n            \"temperature\": temperature,\n            \"max_tokens\": max_tokens,\n            \"stream\": stream,\n        }\n        headers = {\n            \"Authorization\": f\"Bearer {os.environ['MISTRAL_API_KEY']}\",\n            \"Content-Type\": \"application/json\",\n        }\n\n        response = requests.post(\n            \"https://api.mistral.ai/v1/chat/completions\", json=payload, headers=headers\n        )\n        return response.json()\n\n    def completion(\n        self, messages: list, model=\"gpt-4-turbo-preview\", num_retries=2, **kwargs\n    ):\n        # Add in num_retry logic to match the LiteLLM service\n        if model in [\"gpt-4-turbo-preview\", \"gpt-4-turbo\", \"gpt-4o\", \"o1-preview\"]:\n            response = custom_llm_service.openai_query(\n                self, messages=messages, model=model, **kwargs\n            )\n        elif \"claude\" in model:\n            response = custom_llm_service.anthropic_query(\n                self, messages=messages, model=model, **kwargs\n            )\n        elif \"gemini\" in model:\n            response = custom_llm_service.google_ai_studio_query(\n                self, messages=messages, model=model, **kwargs\n            )\n        elif \"azure\" in model:\n            response = custom_llm_service.azure_query(\n                self, messages=messages, model=model, **kwargs\n            )\n        elif \"mistral\" in model or \"mixtral\" in model:\n            response = custom_llm_service.mistral_query(\n                self, messages=messages, model=model, **kwargs\n            )\n        elif \"llama\" in model:\n            response = custom_llm_service.bedrock_query(\n                self, messages=messages, model=model, **kwargs\n            )\n        else:\n            raise ValueError(f\"Model '{model}' not supported\")\n        return response\n\n\n# #%%\n# ##Test Custom LLM Service\n# model = \"gemini-1.5-pro-exp-0801\"\n# custom_llm_service_obj = custom_llm_service()\n# response = custom_llm_service_obj.completion(\n#     messages=[{\"role\": \"user\",\"content\": \"10 + 10 =\"}],\n#     model=model,\n#     max_tokens=100,\n#     temperature=0,\n#     n=1,\n# )\n# print(response)\n# print('\\n-----------\\n')\n# print(message_parse(response, model))\n\n\n# %%\nasync def run_in_executor(func, *args, **kwargs):\n    loop = asyncio.get_running_loop()\n    with ThreadPoolExecutor() as pool:\n        result = await loop.run_in_executor(pool, lambda: func(*args, **kwargs))\n    return result\n\n\nasync def runner(\n    func, messages: list, batch_size=1, validation_func=lambda x: True, **kwargs\n):\n    all_responses = [{}] * len(messages)\n    messages_copy = {i: message for i, message in enumerate(messages)}\n    for idx in range(0, len(messages), batch_size):\n        print(f\"> Processing batch {idx + 1}-{idx + batch_size} ex {len(messages)}\")\n        for _retry in range(1, 4):\n            batch_messages = {\n                i: message\n                for i, message in messages_copy.items()\n                if i in list(range(idx, idx + batch_size))\n            }\n            try:\n                _responses = await asyncio.gather(\n                    *(\n                        run_in_executor(func, messages=_messages, **kwargs)\n                        for _messages in batch_messages.values()\n                    )\n                )\n            except Exception as e:\n                print(f\"Error getting response: {e}. Retry #{_retry}\")\n                time.sleep(30)\n                continue\n            responses = dict(zip(batch_messages.keys(), _responses))\n            for _idx, response in responses.items():\n                if validation_func(message_parse(response, kwargs[\"model\"])):\n                    all_responses[_idx] = response\n                    del messages_copy[_idx]\n                else:\n                    print(f\"Validation failed on response {_idx}. Retry #{_retry}\")\n                    print(f\"Invalid Response: {response}\")\n\n            if len(messages_copy) == 0:\n                break\n            else:\n                time.sleep(30)\n\n    return all_responses\n\n\n# %%\n# messages = [{\"role\": \"user\", \"content\": \"What is the meaning of life?\"}]\n# hyperparams = {\n#     \"max_tokens\": 500,\n#     \"temperature\": 0,\n#     \"num_retries\": 1\n# }\n# responses = await runner(\n#     litellm.completion,\n#     messages=[messages] * 5,\n#     batch_size=5,\n#     model=\"vertex_ai/gemini-1.5-pro\",\n#     **hyperparams,\n# )\n\n# for response in responses:\n#     print(message_parse(response))\n#     print('\\n------------------\\n')\n\n# %%\n"}}}}