diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index cc108cb215c..4a924059410 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -1,5 +1,6 @@ import requests import json +from async_generator import async_generator, yield_ # ollama implementation def get_ollama_response_stream( @@ -32,4 +33,38 @@ def get_ollama_response_stream( yield {"choices": [{"delta": completion_obj}]} except Exception as e: print(f"Error decoding JSON: {e}") + session.close() + +# ollama implementation +@async_generator +async def async_get_ollama_response_stream( + api_base="http://localhost:11434", + model="llama2", + prompt="Why is the sky blue?" + ): + url = f"{api_base}/api/generate" + data = { + "model": model, + "prompt": prompt, + } + session = requests.Session() + + with session.post(url, json=data, stream=True) as resp: + for line in resp.iter_lines(): + 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 "response" in j: + completion_obj = { + "role": "assistant", + "content": "", + } + completion_obj["content"] = j["response"] + await yield_({"choices": [{"delta": completion_obj}]}) + except Exception as e: + print(f"Error decoding JSON: {e}") session.close() \ No newline at end of file diff --git a/litellm/main.py b/litellm/main.py index f27e07f6e9f..546d8a73456 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -75,7 +75,7 @@ async def acompletion(*args, **kwargs): loop = asyncio.get_event_loop() # Use a partial function to pass your keyword arguments - func = partial(completion, *args, **kwargs) + func = partial(completion, *args, **kwargs, acompletion=True) # Add the context to the function ctx = contextvars.copy_context() @@ -180,6 +180,7 @@ def completion( fallbacks=[], caching = False, cache_params = {}, # optional to specify metadata for caching + acompletion=False, ) -> ModelResponse: """ Perform a completion() using any of litellm supported llms (example gpt-4, gpt-3.5-turbo, claude-2, command-nightly) @@ -215,7 +216,7 @@ def completion( """ if mock_response: return mock_completion(model, messages, stream=stream, mock_response=mock_response) - + args = locals() try: logging = litellm_logging_obj @@ -928,6 +929,10 @@ def completion( logging.pre_call( input=prompt, api_key=None, additional_args={"endpoint": endpoint} ) + if acompletion == True: + async_generator = ollama.async_get_ollama_response_stream(endpoint, model, prompt) + return async_generator + generator = ollama.get_ollama_response_stream(endpoint, model, prompt) if optional_params.get("stream", False) == True: # assume all ollama responses are streamed