diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index 3ac9ec2a84..ceab2c7d37 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -1,14 +1,10 @@ import requests, types, time -import json +import json, uuid import traceback from typing import Optional import litellm import httpx, aiohttp, asyncio -try: - from async_generator import async_generator, yield_ # optional dependency - async_generator_imported = True -except ImportError: - async_generator_imported = False # this should not throw an error, it will impact the 'import litellm' statement +from .prompt_templates.factory import prompt_factory, custom_prompt class OllamaError(Exception): def __init__(self, status_code, message): @@ -106,9 +102,8 @@ class OllamaConfig(): and not isinstance(v, (types.FunctionType, types.BuiltinFunctionType, classmethod, staticmethod)) and v is not None} - # ollama implementation -def get_ollama_response_stream( +def get_ollama_response( api_base="http://localhost:11434", model="llama2", prompt="Why is the sky blue?", @@ -129,6 +124,7 @@ def get_ollama_response_stream( if k not in optional_params: # completion(top_k=3) > cohere_config(top_k=3) <- allows for dynamic variables to be passed in optional_params[k] = v + optional_params["stream"] = optional_params.get("stream", False) data = { "model": model, "prompt": prompt, @@ -146,9 +142,41 @@ def get_ollama_response_stream( else: response = ollama_acompletion(url=url, data=data, model_response=model_response, encoding=encoding, logging_obj=logging_obj) return response - - else: + elif optional_params.get("stream", False): return ollama_completion_stream(url=url, data=data, logging_obj=logging_obj) + response = requests.post( + url=f"{url}", + json=data, + ) + if response.status_code != 200: + raise OllamaError(status_code=response.status_code, message=response.text) + + ## LOGGING + logging_obj.post_call( + input=prompt, + api_key="", + original_response=response.text, + additional_args={ + "headers": None, + "api_base": api_base, + }, + ) + + response_json = response.json() + + ## RESPONSE OBJECT + model_response["choices"][0]["finish_reason"] = "stop" + if optional_params.get("format", "") == "json": + message = litellm.Message(content=None, tool_calls=[{"id": f"call_{str(uuid.uuid4())}", "function": {"arguments": response_json["response"], "name": ""}, "type": "function"}]) + model_response["choices"][0]["message"] = message + else: + model_response["choices"][0]["message"]["content"] = response_json["response"] + model_response["created"] = int(time.time()) + model_response["model"] = "ollama/" + model + prompt_tokens = response_json["prompt_eval_count"] # type: ignore + completion_tokens = response_json["eval_count"] + model_response["usage"] = litellm.Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens) + return model_response def ollama_completion_stream(url, data, logging_obj): with httpx.stream( @@ -157,13 +185,15 @@ def ollama_completion_stream(url, data, logging_obj): method="POST", timeout=litellm.request_timeout ) as response: - if response.status_code != 200: - raise OllamaError(status_code=response.status_code, message=response.text) - - streamwrapper = litellm.CustomStreamWrapper(completion_stream=response.iter_lines(), model=data['model'], custom_llm_provider="ollama",logging_obj=logging_obj) - for transformed_chunk in streamwrapper: - yield transformed_chunk - + try: + if response.status_code != 200: + raise OllamaError(status_code=response.status_code, message=response.text) + + streamwrapper = litellm.CustomStreamWrapper(completion_stream=response.iter_lines(), model=data['model'], custom_llm_provider="ollama",logging_obj=logging_obj) + for transformed_chunk in streamwrapper: + yield transformed_chunk + except Exception as e: + raise e async def ollama_async_streaming(url, data, model_response, encoding, logging_obj): try: @@ -194,38 +224,29 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj): text = await resp.text() raise OllamaError(status_code=resp.status, message=text) - completion_string = "" - async for line in resp.content: - if line: - try: - json_chunk = line.decode("utf-8") - chunks = json_chunk.split("\n") - for chunk in chunks: - if chunk.strip() != "": - j = json.loads(chunk) - if "error" in j: - completion_obj = { - "role": "assistant", - "content": "", - "error": j - } - raise Exception(f"OllamError - {chunk}") - if "response" in j: - completion_obj = { - "role": "assistant", - "content": j["response"], - } - completion_string = completion_string + completion_obj["content"] - except Exception as e: - traceback.print_exc() - + ## LOGGING + logging_obj.post_call( + input=data['prompt'], + api_key="", + original_response=resp.text, + additional_args={ + "headers": None, + "api_base": url, + }, + ) + + response_json = await resp.json() ## RESPONSE OBJECT model_response["choices"][0]["finish_reason"] = "stop" - model_response["choices"][0]["message"]["content"] = completion_string + if data.get("format", "") == "json": + message = litellm.Message(content=None, tool_calls=[{"id": f"call_{str(uuid.uuid4())}", "function": {"arguments": response_json["response"], "name": ""}, "type": "function"}]) + model_response["choices"][0]["message"] = message + else: + model_response["choices"][0]["message"]["content"] = response_json["response"] model_response["created"] = int(time.time()) model_response["model"] = "ollama/" + data['model'] - prompt_tokens = len(encoding.encode(data['prompt'])) # type: ignore - completion_tokens = len(encoding.encode(completion_string)) + prompt_tokens = response_json["prompt_eval_count"] # type: ignore + completion_tokens = response_json["eval_count"] model_response["usage"] = litellm.Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens) return model_response except Exception as e: diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index 4596e2b620..d908231cb6 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -348,7 +348,7 @@ def anthropic_pt(messages: list): # format - https://docs.anthropic.com/claude/r # Function call template def function_call_prompt(messages: list, functions: list): - function_prompt = "The following functions are available to you:" + function_prompt = "Produce JSON OUTPUT ONLY! The following functions are available to you:" for function in functions: function_prompt += f"""\n{function}\n""" @@ -425,6 +425,6 @@ def prompt_factory(model: str, messages: list, custom_llm_provider: Optional[str return alpaca_pt(messages=messages) else: return hf_chat_template(original_model_name, messages) - except: + except Exception as e: return default_pt(messages=messages) # default that covers Bloom, T-5, any non-chat tuned model (e.g. base Llama2) diff --git a/litellm/main.py b/litellm/main.py index 1e2a8323a4..0018844c34 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1329,23 +1329,11 @@ def completion( optional_params["images"] = images ## LOGGING - generator = ollama.get_ollama_response_stream(api_base, model, prompt, optional_params, logging_obj=logging, acompletion=acompletion, model_response=model_response, encoding=encoding) + generator = ollama.get_ollama_response(api_base, model, prompt, optional_params, logging_obj=logging, acompletion=acompletion, model_response=model_response, encoding=encoding) if acompletion is True or optional_params.get("stream", False) == True: return generator - else: - response_string = "" - for chunk in generator: - response_string+=chunk['content'] - - ## RESPONSE OBJECT - model_response["choices"][0]["finish_reason"] = "stop" - model_response["choices"][0]["message"]["content"] = response_string - model_response["created"] = int(time.time()) - model_response["model"] = "ollama/" + model - prompt_tokens = len(encoding.encode(prompt)) # type: ignore - completion_tokens = len(encoding.encode(response_string)) - model_response["usage"] = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens) - response = model_response + + response = generator elif ( custom_llm_provider == "baseten" or litellm.api_base == "https://app.baseten.co" diff --git a/litellm/tests/test_hf_prompt_templates.py b/litellm/tests/test_hf_prompt_templates.py index 9f1aedcff3..c67779f871 100644 --- a/litellm/tests/test_hf_prompt_templates.py +++ b/litellm/tests/test_hf_prompt_templates.py @@ -17,6 +17,14 @@ def test_prompt_formatting(): assert prompt == "[INST] Be a good bot [/INST] [INST] Hello world [/INST]" except Exception as e: pytest.fail(f"An exception occurred: {str(e)}") + +def test_prompt_formatting_custom_model(): + try: + prompt = prompt_factory(model="ehartford/dolphin-2.5-mixtral-8x7b", messages=[{"role": "system", "content": "Be a good bot"}, {"role": "user", "content": "Hello world"}], custom_llm_provider="huggingface") + print(f"prompt: {prompt}") + except Exception as e: + pytest.fail(f"An exception occurred: {str(e)}") +# test_prompt_formatting_custom_model() # def logger_fn(user_model_dict): # return # print(f"user_model_dict: {user_model_dict}") diff --git a/litellm/tests/test_ollama_local.py b/litellm/tests/test_ollama_local.py index 5beae7033e..35be975065 100644 --- a/litellm/tests/test_ollama_local.py +++ b/litellm/tests/test_ollama_local.py @@ -16,23 +16,61 @@ # user_message = "respond in 20 words. who are you?" # messages = [{ "content": user_message,"role": "user"}] +# def test_ollama_streaming(): +# try: +# litellm.set_verbose = False +# messages = [ +# {"role": "user", "content": "What is the weather like in Boston?"} +# ] +# functions = [ +# { +# "name": "get_current_weather", +# "description": "Get the current weather in a given location", +# "parameters": { +# "type": "object", +# "properties": { +# "location": { +# "type": "string", +# "description": "The city and state, e.g. San Francisco, CA" +# }, +# "unit": { +# "type": "string", +# "enum": ["celsius", "fahrenheit"] +# } +# }, +# "required": ["location"] +# } +# } +# ] +# response = litellm.completion(model="ollama/mistral", +# messages=messages, +# functions=functions, +# stream=True) +# for chunk in response: +# print(f"CHUNK: {chunk}") +# except Exception as e: +# print(e) + +# test_ollama_streaming() + # async def test_async_ollama_streaming(): # try: -# litellm.set_verbose = True +# litellm.set_verbose = False # response = await litellm.acompletion(model="ollama/mistral-openorca", # messages=[{"role": "user", "content": "Hey, how's it going?"}], # stream=True) # async for chunk in response: -# print(chunk) +# print(f"CHUNK: {chunk}") # except Exception as e: # print(e) -# asyncio.run(test_async_ollama_streaming()) +# # asyncio.run(test_async_ollama_streaming()) # def test_completion_ollama(): # try: +# litellm.set_verbose = True # response = completion( -# model="ollama/llama2", +# model="ollama/mistral", # messages=[{"role": "user", "content": "Hey, how's it going?"}], # max_tokens=200, # request_timeout = 10, @@ -44,7 +82,87 @@ # except Exception as e: # pytest.fail(f"Error occurred: {e}") -# test_completion_ollama() +# # test_completion_ollama() + +# def test_completion_ollama_function_calling(): +# try: +# litellm.set_verbose = True +# messages = [ +# {"role": "user", "content": "What is the weather like in Boston?"} +# ] +# functions = [ +# { +# "name": "get_current_weather", +# "description": "Get the current weather in a given location", +# "parameters": { +# "type": "object", +# "properties": { +# "location": { +# "type": "string", +# "description": "The city and state, e.g. San Francisco, CA" +# }, +# "unit": { +# "type": "string", +# "enum": ["celsius", "fahrenheit"] +# } +# }, +# "required": ["location"] +# } +# } +# ] +# response = completion( +# model="ollama/mistral", +# messages=messages, +# functions=functions, +# max_tokens=200, +# request_timeout = 10, +# ) +# for chunk in response: +# print(chunk) +# print(response) +# except Exception as e: +# pytest.fail(f"Error occurred: {e}") +# # test_completion_ollama_function_calling() + +# async def async_test_completion_ollama_function_calling(): +# try: +# litellm.set_verbose = True +# messages = [ +# {"role": "user", "content": "What is the weather like in Boston?"} +# ] +# functions = [ +# { +# "name": "get_current_weather", +# "description": "Get the current weather in a given location", +# "parameters": { +# "type": "object", +# "properties": { +# "location": { +# "type": "string", +# "description": "The city and state, e.g. San Francisco, CA" +# }, +# "unit": { +# "type": "string", +# "enum": ["celsius", "fahrenheit"] +# } +# }, +# "required": ["location"] +# } +# } +# ] +# response = await litellm.acompletion( +# model="ollama/mistral", +# messages=messages, +# functions=functions, +# max_tokens=200, +# request_timeout = 10, +# ) +# print(response) +# except Exception as e: +# pytest.fail(f"Error occurred: {e}") + +# # asyncio.run(async_test_completion_ollama_function_calling()) + # def test_completion_ollama_with_api_base(): # try: @@ -197,7 +315,7 @@ # ) # print("Response from ollama/llava") # print(response) -# test_ollama_llava() +# # test_ollama_llava() # # PROCESSED CHUNK PRE CHUNK CREATOR diff --git a/litellm/utils.py b/litellm/utils.py index 98806c101f..c449a239e2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2390,10 +2390,15 @@ def get_optional_params( # use the openai defaults non_default_params = {k: v for k, v in passed_params.items() if (k != "model" and k != "custom_llm_provider" and k in default_params and v != default_params[k])} optional_params = {} ## raise exception if function calling passed in for a provider that doesn't support it - if "functions" in non_default_params or "function_call" in non_default_params: + if "functions" in non_default_params or "function_call" in non_default_params or "tools" in non_default_params: if custom_llm_provider != "openai" and custom_llm_provider != "text-completion-openai" and custom_llm_provider != "azure": - if litellm.add_function_to_prompt: # if user opts to add it to prompt instead - optional_params["functions_unsupported_model"] = non_default_params.pop("functions") + if custom_llm_provider == "ollama": + # ollama actually supports json output + optional_params["format"] = "json" + litellm.add_function_to_prompt = True # so that main.py adds the function call to the prompt + optional_params["functions_unsupported_model"] = non_default_params.pop("tools", non_default_params.pop("functions")) + elif litellm.add_function_to_prompt: # if user opts to add it to prompt instead + optional_params["functions_unsupported_model"] = non_default_params.pop("tools", non_default_params.pop("functions")) else: raise UnsupportedParamsError(status_code=500, message=f"Function calling is not supported by {custom_llm_provider}. To add it to the prompt, set `litellm.add_function_to_prompt = True`.") @@ -5192,9 +5197,6 @@ def exception_type( raise original_exception raise original_exception elif custom_llm_provider == "ollama": - if "no attribute 'async_get_ollama_response_stream" in error_str: - exception_mapping_worked = True - raise ImportError("Import error - trying to use async for ollama. import async_generator failed. Try 'pip install async_generator'") if isinstance(original_exception, dict): error_str = original_exception.get("error", "") else: