mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
107 lines
2.6 KiB
Python
107 lines
2.6 KiB
Python
import os
|
|
import sys, os
|
|
import traceback
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
import os, io
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../..")
|
|
) # Adds the parent directory to the system path
|
|
import pytest
|
|
import litellm
|
|
from litellm import embedding, completion, text_completion, completion_cost
|
|
|
|
from langchain.chat_models import ChatLiteLLM
|
|
from langchain.prompts.chat import (
|
|
ChatPromptTemplate,
|
|
SystemMessagePromptTemplate,
|
|
AIMessagePromptTemplate,
|
|
HumanMessagePromptTemplate,
|
|
)
|
|
from langchain.schema import AIMessage, HumanMessage, SystemMessage
|
|
|
|
def test_chat_gpt():
|
|
try:
|
|
chat = ChatLiteLLM(model="gpt-3.5-turbo", max_tokens=10)
|
|
messages = [
|
|
HumanMessage(
|
|
content="what model are you"
|
|
)
|
|
]
|
|
resp = chat(messages)
|
|
|
|
print(resp)
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
# test_chat_gpt()
|
|
|
|
|
|
def test_claude():
|
|
try:
|
|
chat = ChatLiteLLM(model="claude-2", max_tokens=10)
|
|
messages = [
|
|
HumanMessage(
|
|
content="what model are you"
|
|
)
|
|
]
|
|
resp = chat(messages)
|
|
|
|
print(resp)
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
# test_claude()
|
|
|
|
def test_palm():
|
|
try:
|
|
chat = ChatLiteLLM(model="palm/chat-bison", max_tokens=10)
|
|
messages = [
|
|
HumanMessage(
|
|
content="what model are you"
|
|
)
|
|
]
|
|
resp = chat(messages)
|
|
|
|
print(resp)
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
# test_palm()
|
|
|
|
|
|
# def test_openai_with_params():
|
|
# try:
|
|
# api_key = os.environ["OPENAI_API_KEY"]
|
|
# os.environ.pop("OPENAI_API_KEY")
|
|
# print("testing openai with params")
|
|
# llm = ChatLiteLLM(
|
|
# model="gpt-3.5-turbo",
|
|
# openai_api_key=api_key,
|
|
# # Prefer using None which is the default value, endpoint could be empty string
|
|
# openai_api_base= None,
|
|
# max_tokens=20,
|
|
# temperature=0.5,
|
|
# request_timeout=10,
|
|
# model_kwargs={
|
|
# "frequency_penalty": 0,
|
|
# "presence_penalty": 0,
|
|
# },
|
|
# verbose=True,
|
|
# max_retries=0,
|
|
# )
|
|
# messages = [
|
|
# HumanMessage(
|
|
# content="what model are you"
|
|
# )
|
|
# ]
|
|
# resp = llm(messages)
|
|
|
|
# print(resp)
|
|
# except Exception as e:
|
|
# pytest.fail(f"Error occurred: {e}")
|
|
|
|
# test_openai_with_params()
|
|
|