feat(ollama.py): add support for ollama function calling

2026-08-14 16:25:29 +00:00 · 2023-12-20 14:59:55 +05:30
parent bab8f3350d
commit f0df28362a
6 changed files with 211 additions and 74 deletions
@@ -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:
@@ -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)
    
@@ -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"
@@ -17,6 +17,14 @@ def test_prompt_formatting():
        assert prompt == "<s>[INST] Be a good bot [/INST]</s> [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}")
@@ -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
@@ -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: