{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6e35563a",
   "metadata": {},
   "source": [
    "# Basic Completion and Embedding Examples\n",
    "\n",
    "## Completion\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "aa03e40d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The capital of France is Paris.\n",
      "The capital of France is Paris.\n",
      "Full Response:\n",
      "{\n",
      "  \"id\": \"chatcmpl-CyPuxOjKPmvuCvJwTJiLRH1lwO77J\",\n",
      "  \"choices\": [\n",
      "    {\n",
      "      \"finish_reason\": \"stop\",\n",
      "      \"index\": 0,\n",
      "      \"logprobs\": null,\n",
      "      \"message\": {\n",
      "        \"content\": \"The capital of France is Paris.\",\n",
      "        \"refusal\": null,\n",
      "        \"role\": \"assistant\",\n",
      "        \"annotations\": [],\n",
      "        \"audio\": null,\n",
      "        \"function_call\": null,\n",
      "        \"tool_calls\": null\n",
      "      },\n",
      "      \"provider_specific_fields\": {}\n",
      "    }\n",
      "  ],\n",
      "  \"created\": 1768515343,\n",
      "  \"model\": \"gpt-4o-2024-05-13\",\n",
      "  \"object\": \"chat.completion\",\n",
      "  \"service_tier\": null,\n",
      "  \"system_fingerprint\": \"fp_3eed281ddb\",\n",
      "  \"usage\": {\n",
      "    \"completion_tokens\": 8,\n",
      "    \"prompt_tokens\": 14,\n",
      "    \"total_tokens\": 22,\n",
      "    \"completion_tokens_details\": {\n",
      "      \"accepted_prediction_tokens\": 0,\n",
      "      \"audio_tokens\": 0,\n",
      "      \"reasoning_tokens\": 0,\n",
      "      \"rejected_prediction_tokens\": 0,\n",
      "      \"text_tokens\": null\n",
      "    },\n",
      "    \"prompt_tokens_details\": {\n",
      "      \"audio_tokens\": 0,\n",
      "      \"cached_tokens\": 0,\n",
      "      \"text_tokens\": null,\n",
      "      \"image_tokens\": null\n",
      "    }\n",
      "  },\n",
      "  \"formatted_response\": null,\n",
      "  \"content\": \"The capital of France is Paris.\"\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "# Copyright (c) 2024 Microsoft Corporation.\n",
    "# Licensed under the MIT License\n",
    "\n",
    "import os\n",
    "from collections.abc import AsyncIterator, Iterator\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from graphrag_llm.completion import LLMCompletion, create_completion\n",
    "from graphrag_llm.config import AuthMethod, ModelConfig\n",
    "from graphrag_llm.types import LLMCompletionChunk, LLMCompletionResponse\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "api_key = os.getenv(\"GRAPHRAG_API_KEY\")\n",
    "model_config = ModelConfig(\n",
    "    model_provider=\"azure\",\n",
    "    model=os.getenv(\"GRAPHRAG_MODEL\", \"gpt-4o\"),\n",
    "    azure_deployment_name=os.getenv(\"GRAPHRAG_MODEL\", \"gpt-4o\"),\n",
    "    api_base=os.getenv(\"GRAPHRAG_API_BASE\"),\n",
    "    api_version=os.getenv(\"GRAPHRAG_API_VERSION\", \"2025-04-01-preview\"),\n",
    "    api_key=api_key,\n",
    "    auth_method=AuthMethod.AzureManagedIdentity if not api_key else AuthMethod.ApiKey,\n",
    ")\n",
    "llm_completion: LLMCompletion = create_completion(model_config)\n",
    "\n",
    "response: LLMCompletionResponse | Iterator[LLMCompletionChunk] = (\n",
    "    llm_completion.completion(\n",
    "        messages=\"What is the capital of France?\",\n",
    "    )\n",
    ")\n",
    "\n",
    "if isinstance(response, Iterator):\n",
    "    # Streaming response\n",
    "    for chunk in response:\n",
    "        print(chunk.choices[0].delta.content or \"\", end=\"\", flush=True)\n",
    "else:\n",
    "    # Non-streaming response\n",
    "    print(response.choices[0].message.content)\n",
    "    # Or alternatively, access via the content property\n",
    "    # This is equivalent to the above line, getting the content of the first choice\n",
    "    print(response.content)\n",
    "\n",
    "print(\"Full Response:\")\n",
    "print(response.model_dump_json(indent=2))  # type: ignore"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "558392ce",
   "metadata": {},
   "source": [
    "## Async Completion\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8405fcb7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The capital of France is Paris.\n"
     ]
    }
   ],
   "source": [
    "response: LLMCompletionResponse = await llm_completion.completion_async(\n",
    "    messages=\"What is the capital of France?\",\n",
    ")  # type: ignore\n",
    "print(response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e70fc49a",
   "metadata": {},
   "source": [
    "## Streaming Completion\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9f60c4e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The capital of France is Paris."
     ]
    }
   ],
   "source": [
    "response = llm_completion.completion(\n",
    "    messages=\"What is the capital of France?\",\n",
    "    stream=True,\n",
    ")\n",
    "\n",
    "if isinstance(response, Iterator):\n",
    "    # Streaming response\n",
    "    for chunk in response:\n",
    "        print(chunk.choices[0].delta.content or \"\", end=\"\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe8c2e35",
   "metadata": {},
   "source": [
    "## Async Streaming Completion\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0be849ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The capital of France is Paris."
     ]
    }
   ],
   "source": [
    "response = await llm_completion.completion_async(\n",
    "    messages=\"What is the capital of France?\",\n",
    "    stream=True,\n",
    ")\n",
    "\n",
    "if isinstance(response, AsyncIterator):\n",
    "    # Streaming response\n",
    "    async for chunk in response:\n",
    "        print(chunk.choices[0].delta.content or \"\", end=\"\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c32070ad",
   "metadata": {},
   "source": [
    "## Completion Arguments\n",
    "\n",
    "The completion API adheres to litellm completion API and thus the OpanAI SDK API. The `messages` parameter can be one of the following:\n",
    "\n",
    "- `str`: Raw string for the prompt.\n",
    "- `list[dict[str, Any]]`: A list of dicts in the form `{\"role\": \"user|system|...\", \"content\": \"...\"}`\n",
    "- `list[ChatCompletionMessageParam]`: A list of OpenAI `ChatCompletionMessageParam`. `graphrag_llm.utils` provides a `ChatCompletionMessageParamBuilder` to help construct these objects. See the message builder notebook for more details on using `ChatCompletionMessageParamBuilder`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8fe480cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The capital of France is Paris.\n",
      "The capital of France is Paris.\n",
      "Arrr, ye got me there, matey! Truth be, back in 2006, them fancy scallywags at the International Astronomical Union be sayin' Pluto ain't a full-fledged planet no more. They be callin' it a \"dwarf planet\" now. So, officially, she be a dwarf planet, savvy?\n"
     ]
    }
   ],
   "source": [
    "from graphrag_llm.utils import (\n",
    "    CompletionMessagesBuilder,\n",
    ")\n",
    "\n",
    "# raw string input\n",
    "response1: LLMCompletionResponse = llm_completion.completion(\n",
    "    messages=\"What is the capital of France?\"\n",
    ")  # type: ignore\n",
    "print(response1.content)\n",
    "\n",
    "# list of message dicts input\n",
    "response2: LLMCompletionResponse = llm_completion.completion(\n",
    "    messages=[{\"role\": \"user\", \"content\": \"What is the capital of France?\"}]\n",
    ")  # type: ignore\n",
    "print(response2.content)\n",
    "\n",
    "# using the builder to create complex message\n",
    "messages = (\n",
    "    CompletionMessagesBuilder()\n",
    "    .add_system_message(\n",
    "        \"You are a helpful assistant that likes to talk like a pirate. Respond as if you are a pirate using pirate speak.\"\n",
    "    )\n",
    "    .add_user_message(\"Is pluto a planet? Respond with a yes or no.\")\n",
    "    .add_assistant_message(\"Aye, matey! Pluto be a planet in me book.\")\n",
    "    .add_user_message(\"Are you sure? I want the truth. Can you elaborate?\")\n",
    "    .build()\n",
    ")\n",
    "\n",
    "response3: LLMCompletionResponse = llm_completion.completion(messages=messages)  # type: ignore\n",
    "print(response3.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dda66594",
   "metadata": {},
   "source": [
    "## Embedding\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "51fe336b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.002078542485833168, -0.04908587411046028, 0.020946789532899857]\n",
      "[0.027567066252231598, -0.026544300839304924, -0.027091361582279205]\n"
     ]
    }
   ],
   "source": [
    "from graphrag_llm.embedding import LLMEmbedding, create_embedding\n",
    "from graphrag_llm.types import LLMEmbeddingResponse\n",
    "\n",
    "embedding_config = ModelConfig(\n",
    "    model_provider=\"azure\",\n",
    "    model=os.getenv(\"GRAPHRAG_EMBEDDING_MODEL\", \"text-embedding-3-small\"),\n",
    "    azure_deployment_name=os.getenv(\n",
    "        \"GRAPHRAG_LLM_EMBEDDING_MODEL\", \"text-embedding-3-small\"\n",
    "    ),\n",
    "    api_base=os.getenv(\"GRAPHRAG_API_BASE\"),\n",
    "    api_version=os.getenv(\"GRAPHRAG_API_VERSION\", \"2025-04-01-preview\"),\n",
    "    api_key=api_key,\n",
    "    auth_method=AuthMethod.AzureManagedIdentity if not api_key else AuthMethod.ApiKey,\n",
    ")\n",
    "\n",
    "llm_embedding: LLMEmbedding = create_embedding(embedding_config)\n",
    "\n",
    "embeddings_batch: LLMEmbeddingResponse = llm_embedding.embedding(\n",
    "    input=[\"Hello world\", \"How are you?\"]\n",
    ")\n",
    "for embedding in embeddings_batch.embeddings:\n",
    "    print(embedding[0:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3b7bedf",
   "metadata": {},
   "source": [
    "### First Embedding\n",
    "\n",
    "`.embedding` batches by default, it takes a list of strings to embed. If embedding a single string then you can use `.first_embedding` on the response to obtain the first embedding.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e428c64a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.05073608458042145, 0.003799507161602378, 0.019212841987609863]\n"
     ]
    }
   ],
   "source": [
    "embedding_response = llm_embedding.embedding(\n",
    "    input=[\"This is a single input string for embedding.\"]\n",
    ")\n",
    "\n",
    "print(embedding_response.first_embedding[0:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b4cf0fa",
   "metadata": {},
   "source": [
    "## Async Embedding\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c9519657",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.002078542485833168, -0.04908587411046028, 0.020946789532899857]\n",
      "[0.027567066252231598, -0.026544300839304924, -0.027091361582279205]\n"
     ]
    }
   ],
   "source": [
    "embeddings_batch = await llm_embedding.embedding_async(\n",
    "    input=[\"Hello world\", \"How are you?\"]\n",
    ")\n",
    "\n",
    "for embedding in embeddings_batch.embeddings:\n",
    "    print(embedding[0:3])"
   ]
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}