{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install litellm==0.1.385"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from litellm import completion\n",
    "import asyncio"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup Messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_message = \"respond in 20 words. who are you?\"\n",
    "messages = [{ \"content\": user_message,\"role\": \"user\"}]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Call Ollama - llama2 with chatGPT Input/Output using litellm.completion() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<async_generator object get_ollama_response_stream at 0x1069198b0>\n"
     ]
    }
   ],
   "source": [
    "response = completion(model=\"llama2\", messages=messages, custom_api_base=\"http://localhost:11434\", custom_llm_provider=\"ollama\", stream=True)\n",
    "print(response)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Iterate through the generator - Streaming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I\n",
      "{'role': 'assistant', 'content': ' I'}\n",
      " am\n",
      "{'role': 'assistant', 'content': ' am'}\n",
      " L\n",
      "{'role': 'assistant', 'content': ' L'}\n",
      "La\n",
      "{'role': 'assistant', 'content': 'La'}\n",
      "MA\n",
      "{'role': 'assistant', 'content': 'MA'}\n",
      ",\n",
      "{'role': 'assistant', 'content': ','}\n",
      " an\n",
      "{'role': 'assistant', 'content': ' an'}\n",
      " A\n",
      "{'role': 'assistant', 'content': ' A'}\n",
      "I\n",
      "{'role': 'assistant', 'content': 'I'}\n",
      " assistant\n",
      "{'role': 'assistant', 'content': ' assistant'}\n",
      " developed\n",
      "{'role': 'assistant', 'content': ' developed'}\n",
      " by\n",
      "{'role': 'assistant', 'content': ' by'}\n",
      " Meta\n",
      "{'role': 'assistant', 'content': ' Meta'}\n",
      " A\n",
      "{'role': 'assistant', 'content': ' A'}\n",
      "I\n",
      "{'role': 'assistant', 'content': 'I'}\n",
      " that\n",
      "{'role': 'assistant', 'content': ' that'}\n",
      " can\n",
      "{'role': 'assistant', 'content': ' can'}\n",
      " understand\n",
      "{'role': 'assistant', 'content': ' understand'}\n",
      " and\n",
      "{'role': 'assistant', 'content': ' and'}\n",
      " respond\n",
      "{'role': 'assistant', 'content': ' respond'}\n",
      " to\n",
      "{'role': 'assistant', 'content': ' to'}\n",
      " human\n",
      "{'role': 'assistant', 'content': ' human'}\n",
      " input\n",
      "{'role': 'assistant', 'content': ' input'}\n",
      " in\n",
      "{'role': 'assistant', 'content': ' in'}\n",
      " a\n",
      "{'role': 'assistant', 'content': ' a'}\n",
      " convers\n",
      "{'role': 'assistant', 'content': ' convers'}\n",
      "ational\n",
      "{'role': 'assistant', 'content': 'ational'}\n",
      " manner\n",
      "{'role': 'assistant', 'content': ' manner'}\n",
      ".\n",
      "{'role': 'assistant', 'content': '.'}\n"
     ]
    }
   ],
   "source": [
    "\n",
    "async def get_response(generator):\n",
    "    response = \"\"\n",
    "    async for elem in generator:\n",
    "        print(elem)\n",
    "        response += elem['choices'][0]['delta'][\"content\"]\n",
    "    return response\n",
    "\n",
    "string_response = await get_response(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}