Fix typing

This commit is contained in:
yujong-lee 2023-08-18 10:30:47 +09:00
parent 88fdef0c30
commit 1b78202fc0
3 changed files with 5 additions and 3 deletions

View file

@ -4,6 +4,7 @@ import requests
from litellm import logging from litellm import logging
import time import time
from typing import Callable from typing import Callable
from litellm.utils import ModelResponse
class AnthropicConstants(Enum): class AnthropicConstants(Enum):
HUMAN_PROMPT = "\n\nHuman:" HUMAN_PROMPT = "\n\nHuman:"
@ -36,7 +37,7 @@ class AnthropicLLM:
"x-api-key": self.api_key "x-api-key": self.api_key
} }
def completion(self, model: str, messages: list, model_response: dict, print_verbose: Callable, optional_params=None, litellm_params=None, logger_fn=None): # logic for parsing in - calling - parsing out model completion calls def completion(self, model: str, messages: list, model_response: ModelResponse, print_verbose: Callable, optional_params=None, litellm_params=None, logger_fn=None): # logic for parsing in - calling - parsing out model completion calls
model = model model = model
prompt = f"{AnthropicConstants.HUMAN_PROMPT.value}" prompt = f"{AnthropicConstants.HUMAN_PROMPT.value}"
for message in messages: for message in messages:

View file

@ -5,6 +5,7 @@ import requests
from litellm import logging from litellm import logging
import time import time
from typing import Callable from typing import Callable
from litellm.utils import ModelResponse
class HuggingfaceError(Exception): class HuggingfaceError(Exception):
def __init__(self, status_code, message): def __init__(self, status_code, message):
@ -26,7 +27,7 @@ class HuggingfaceRestAPILLM():
if self.api_key != None: if self.api_key != None:
self.headers["Authorization"] = f"Bearer {self.api_key}" self.headers["Authorization"] = f"Bearer {self.api_key}"
def completion(self, model: str, messages: list, custom_api_base: str, model_response: dict, print_verbose: Callable, optional_params=None, litellm_params=None, logger_fn=None): # logic for parsing in - calling - parsing out model completion calls def completion(self, model: str, messages: list, custom_api_base: str, model_response: ModelResponse, print_verbose: Callable, optional_params=None, litellm_params=None, logger_fn=None): # logic for parsing in - calling - parsing out model completion calls
if custom_api_base: if custom_api_base:
completion_url = custom_api_base completion_url = custom_api_base
elif "HF_API_BASE" in os.environ: elif "HF_API_BASE" in os.environ:

View file

@ -40,7 +40,7 @@ def completion(
# model specific optional params # model specific optional params
# used by text-bison only # used by text-bison only
top_k=40, request_timeout=0, # unused var for old version of OpenAI API top_k=40, request_timeout=0, # unused var for old version of OpenAI API
): ) -> ModelResponse:
try: try:
model_response = ModelResponse() model_response = ModelResponse()
if azure: # this flag is deprecated, remove once notebooks are also updated. if azure: # this flag is deprecated, remove once notebooks are also updated.