LLM.py

""" This file contains the code for calling all LLM APIs. """

import os
from functools import partial
import tiktoken
import requests
import time
from schema import TooLongPromptError, LLMError
import openai
enc = tiktoken.get_encoding("cl100k_base")

try:   
    import anthropic
    # setup anthropic API key
    anthropic_client = anthropic.Anthropic(api_key=open("claude_api_key.txt").read().strip())
except Exception as e:
    print(e)
    print("Could not load anthropic API key claude_api_key.txt.")
    
try:
    from openai import OpenAI
    # setup OpenAI API key
    openai.organization, openai.api_key  =  open("openai_api_key.txt").read().strip().split(":")    
except Exception as e:
    print(e)
    print("Could not load OpenAI API key openai_api_key.txt.")

try:
    import vertexai
    from vertexai.preview.generative_models import GenerativeModel, Part
    from google.cloud.aiplatform_v1beta1.types import SafetySetting, HarmCategory
    vertexai.init(project=PROJECT_ID, location="us-central1")
except Exception as e:
    print(e)
    print("Could not load VertexAI API.")

# try:
#     from openai import OpenAI
#     api_key = open("deepseek_api_key.txt").read().strip()
#     print(api_key)
#     client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com/")

# except Exception as e:
#     print(e)
#     print("Could not load DeepSeek API key deepseek_api_key.txt.")

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import StoppingCriteria, StoppingCriteriaList
import torch

loaded_hf_models = {}

class StopAtSpecificTokenCriteria(StoppingCriteria):
    def __init__(self, stop_sequence):
        super().__init__()
        self.stop_sequence = stop_sequence

    def __call__(self, input_ids, scores, **kwargs):
        # Create a tensor from the stop_sequence
        stop_sequence_tensor = torch.tensor(self.stop_sequence, device=input_ids.device, dtype=input_ids.dtype)

        # Check if the current sequence ends with the stop_sequence
        current_sequence = input_ids[:, -len(self.stop_sequence) :]
        return bool(torch.all(current_sequence == stop_sequence_tensor).item())

    
def log_to_file(log_file, prompt, completion, model, max_tokens_to_sample):
    """ Log the prompt and completion to a file."""
    with open(log_file, "a") as f:
        f.write("\n===================prompt=====================\n")
        f.write(f"{anthropic.HUMAN_PROMPT} {prompt} {anthropic.AI_PROMPT}")
        num_prompt_tokens = len(enc.encode(f"{anthropic.HUMAN_PROMPT} {prompt} {anthropic.AI_PROMPT}"))
        f.write(f"\n==================={model} response ({max_tokens_to_sample})=====================\n")
        f.write(completion)
        num_sample_tokens = len(enc.encode(completion))
        f.write("\n===================tokens=====================\n")
        f.write(f"Number of prompt tokens: {num_prompt_tokens}\n")
        f.write(f"Number of sampled tokens: {num_sample_tokens}\n")
        f.write("\n\n")

def complete_text_hf(prompt, stop_sequences=[], model="huggingface/codellama/CodeLlama-7b-hf", max_tokens_to_sample = 2000, temperature=0.5, log_file=None, **kwargs):
    model = model.split("/", 1)[1]
    if model in loaded_hf_models:
        hf_model, tokenizer = loaded_hf_models[model]
    else:
        hf_model = AutoModelForCausalLM.from_pretrained(model).to("cuda:0")
        tokenizer = AutoTokenizer.from_pretrained(model)
        loaded_hf_models[model] = (hf_model, tokenizer)
        
    encoded_input = tokenizer(prompt, return_tensors="pt", return_token_type_ids=False).to("cuda:0")
    stop_sequence_ids = tokenizer(stop_sequences, return_token_type_ids=False, add_special_tokens=False)
    stopping_criteria = StoppingCriteriaList()
    for stop_sequence_input_ids in stop_sequence_ids.input_ids:
        stopping_criteria.append(StopAtSpecificTokenCriteria(stop_sequence=stop_sequence_input_ids))

    output = hf_model.generate(
        **encoded_input,
        temperature=temperature,
        max_new_tokens=max_tokens_to_sample,
        do_sample=True,
        return_dict_in_generate=True,
        output_scores=True,
        stopping_criteria = stopping_criteria,
        **kwargs,
    )
    sequences = output.sequences
    sequences = [sequence[len(encoded_input.input_ids[0]) :] for sequence in sequences]
    all_decoded_text = tokenizer.batch_decode(sequences)
    completion = all_decoded_text[0]
    if log_file is not None:
        log_to_file(log_file, prompt, completion, model, max_tokens_to_sample)
    return completion


def complete_text_ollama(prompt, stop_sequences=[], model="llama3", max_tokens_to_sample=2000, temperature=0.5, log_file=None, **kwargs):
    """Call the Ollama API to complete a prompt."""
    url = "http://127.0.0.1:11434/api/generate"
    
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {
            "stop": stop_sequences,
            "temperature": temperature,
            "num_predict": max_tokens_to_sample,
        }
    }
    
    try:
        
        response = requests.post(url, json=payload)
        response.raise_for_status()
        result = response.json()
        completion = result['response']
        
        if log_file is not None:
            log_to_file(log_file, prompt, completion, model, max_tokens_to_sample)
        
        return completion
    except requests.exceptions.RequestException as e:
        print(f"Error calling Ollama API: {e}")
        raise LLMError(f"Ollama API error: {e}")
    
    
def complete_text_gemini(prompt, stop_sequences=[], model="gemini-pro", max_tokens_to_sample = 2000, temperature=0.5, log_file=None, **kwargs):
    """ Call the gemini API to complete a prompt."""
    # Load the model
    model = GenerativeModel("gemini-pro")
    # Query the model
    parameters = {
            "temperature": temperature,
            "max_output_tokens": max_tokens_to_sample,
            "stop_sequences": stop_sequences,
            **kwargs
        }
    safety_settings = {
            harm_category: SafetySetting.HarmBlockThreshold(SafetySetting.HarmBlockThreshold.BLOCK_NONE)
            for harm_category in iter(HarmCategory)
        }
    safety_settings = {
        }
    response = model.generate_content( [prompt], generation_config=parameters, safety_settings=safety_settings)
    completion = response.text
    if log_file is not None:
        log_to_file(log_file, prompt, completion, model, max_tokens_to_sample)
    return completion

def complete_text_claude(prompt, stop_sequences=[anthropic.HUMAN_PROMPT], model="claude-v1", max_tokens_to_sample = 2000, temperature=0.5, log_file=None, messages=None, **kwargs):
    """ Call the Claude API to complete a prompt."""

    ai_prompt = anthropic.AI_PROMPT
    if "ai_prompt" in kwargs is not None:
        ai_prompt = kwargs["ai_prompt"]

    
    try:
        if model in ["claude-3-opus-20240229", "claude-3-5-sonnet-20241022"]:
            while True:
                try:
                    message = anthropic_client.messages.create(
                        messages=[
                            {
                                "role": "user",
                                "content": prompt,
                            }
                        ] if messages is None else messages,
                        model=model,
                        stop_sequences=stop_sequences,
                        temperature=temperature,
                        max_tokens=max_tokens_to_sample,
                        **kwargs
                    )
                except anthropic.InternalServerError as e:
                    pass
                try:
                    completion = message.content[0].text
                    break
                except:
                    print("end_turn???")
                    pass
        else:
            rsp = anthropic_client.completions.create(
                prompt=f"{anthropic.HUMAN_PROMPT} {prompt} {ai_prompt}",
                stop_sequences=stop_sequences,
                model=model,
                temperature=temperature,
                max_tokens_to_sample=max_tokens_to_sample,
                **kwargs
            )
            completion = rsp.completion
        
    except anthropic.APIStatusError as e:
        print(e)
        raise TooLongPromptError()
    except Exception as e:
        raise LLMError(e)

    
    if log_file is not None:
        log_to_file(log_file, prompt, completion, model, max_tokens_to_sample)
    
    time.sleep(5)
    
    return completion

def complete_text_deppseek(prompt, stop_sequences=[], model="deepseek-chat", max_tokens_to_sample=2000, temperature=1.0, log_file=None, **kwargs):
    """ Call the DeepSeek API to complete a prompt."""
    client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com/")
    raw_request = {
        "model": model,
        "temperature": temperature,
        "max_tokens": max_tokens_to_sample,
        "stop": stop_sequences or None,  # API doesn't like empty list
        **kwargs
    }
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(**{"messages": messages,**raw_request})
    completion = response.choices[0].message.content
    if log_file is not None:
        log_to_file(log_file, prompt, completion, model, max_tokens_to_sample)
    return completion

def complete_text_openai(prompt, stop_sequences=[], model="gpt-3.5-turbo", max_tokens_to_sample=2000, temperature=0.2, log_file=None, **kwargs):
    client = OpenAI(api_key=openai.api_key)
    """ Call the OpenAI API to complete a prompt."""
    if model.startswith("o3"):
        raw_request = {
            "model": model,
            "reasoning_effort": "low",
            "max_completion_tokens": max_tokens_to_sample,
            "stop": stop_sequences or None,  # API doesn't like empty list
            **kwargs
        }
        messages = [{"role": "user", "content": prompt}]
        response = client.chat.completions.create(**{"messages": messages,**raw_request})
        completion = response["choices"][0]["message"]["content"]
    else:
        raw_request = {
            "model": model,
            "temperature": temperature,
            "max_tokens": max_tokens_to_sample,
            "stop": stop_sequences or None,  # API doesn't like empty list
            **kwargs
        }
        if model.startswith("gpt-3.5") or model.startswith("gpt-4"):
            messages = [{"role": "user", "content": prompt}]
            response = client.chat.completions.create(**{"messages": messages,**raw_request})
            completion = response.choices[0].message.content
        else:
            response = client.completions.create(**{"prompt": prompt,**raw_request})
            completion = response["choices"][0]["text"]
            
    if log_file is not None:
        log_to_file(log_file, prompt, completion, model, max_tokens_to_sample)
    return completion

def complete_text(prompt, log_file, model, **kwargs):
    """ Complete text using the specified model with appropriate API. """
    
    if model.startswith("claude"):
        completion = complete_text_claude(prompt, stop_sequences=[anthropic.HUMAN_PROMPT, "Observation:"], log_file=log_file, model=model, **kwargs)
    elif model.startswith("gemini"):
        completion = complete_text_gemini(prompt, stop_sequences=["Observation:"], log_file=log_file, model=model, **kwargs)
    elif model.startswith("huggingface"):
        completion = complete_text_hf(prompt, stop_sequences=["Observation:"], log_file=log_file, model=model, **kwargs)
    elif model.startswith("ollama:"):
        ollama_model = model.split(":", 1)[1]
        completion = complete_text_ollama(prompt, stop_sequences=["Observation:"], log_file=log_file, model=ollama_model, **kwargs)
    elif model.startswith("deepseek"):
        completion = complete_text_deppseek(prompt, stop_sequences=["Observation:"], log_file=log_file, model=model, **kwargs)
    else:
        completion = complete_text_openai(prompt, stop_sequences=["Observation:"], log_file=log_file, model=model, **kwargs)
    return completion

# specify fast models for summarization etc
def complete_text_fast(prompt, **kwargs):
    # 在函数内部获取环境变量的值
    fast_model = os.environ.get('FAST_LLM_NAME', 'NOT SET')
    return complete_text(prompt = prompt, model = fast_model, temperature =0.2, **kwargs)