{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9418b981",
   "metadata": {},
   "source": [
    "# Mocking\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d000d70",
   "metadata": {},
   "source": [
    "## Completions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "792c4fa3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Who cares?\n",
      "You tell me!\n",
      "{\"reports\":[{\"city\":\"New York\",\"temperature\":22.5,\"condition\":\"Sunny\"}]}\n",
      "Who cares?\n"
     ]
    }
   ],
   "source": [
    "# Copyright (c) 2024 Microsoft Corporation.\n",
    "# Licensed under the MIT License\n",
    "\n",
    "import os\n",
    "\n",
    "from graphrag_llm.completion import LLMCompletion, create_completion\n",
    "from graphrag_llm.config import LLMProviderType, ModelConfig\n",
    "from graphrag_llm.types import LLMCompletionResponse\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "\n",
    "class LocalWeather(BaseModel):\n",
    "    \"\"\"City weather information model.\"\"\"\n",
    "\n",
    "    city: str = Field(description=\"The name of the city\")\n",
    "    temperature: float = Field(description=\"The temperature in Celsius\")\n",
    "    condition: str = Field(description=\"The weather condition description\")\n",
    "\n",
    "\n",
    "class WeatherReports(BaseModel):\n",
    "    \"\"\"Weather information model.\"\"\"\n",
    "\n",
    "    reports: list[LocalWeather] = Field(\n",
    "        description=\"The weather reports for multiple cities\"\n",
    "    )\n",
    "\n",
    "\n",
    "weather_reports = WeatherReports(\n",
    "    reports=[\n",
    "        LocalWeather(city=\"New York\", temperature=22.5, condition=\"Sunny\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "api_key = os.getenv(\"GRAPHRAG_API_KEY\")\n",
    "model_config = ModelConfig(\n",
    "    type=LLMProviderType.MockLLM,\n",
    "    model_provider=\"openai\",\n",
    "    model=\"gpt-4o\",\n",
    "    mock_responses=[\"Who cares?\", \"You tell me!\", weather_reports.model_dump_json()],\n",
    ")\n",
    "llm_completion: LLMCompletion = create_completion(model_config)\n",
    "\n",
    "response: LLMCompletionResponse = llm_completion.completion(\n",
    "    messages=\"What is the capital of France?\",\n",
    ")  # type: ignore\n",
    "\n",
    "print(response.content)\n",
    "\n",
    "response: LLMCompletionResponse = llm_completion.completion(\n",
    "    messages=\"Should be second response\",\n",
    ")  # type: ignore\n",
    "print(response.content)\n",
    "\n",
    "response_formatted: LLMCompletionResponse[WeatherReports] = llm_completion.completion(\n",
    "    messages=\"Structured response.\",\n",
    "    response_format=WeatherReports,\n",
    ")  # type: ignore\n",
    "print(response_formatted.formatted_response.model_dump_json())  # type: ignore\n",
    "\n",
    "response: LLMCompletionResponse = llm_completion.completion(\n",
    "    messages=\"Should cycle back to first response\",\n",
    ")  # type: ignore\n",
    "print(response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c8f1b7a",
   "metadata": {},
   "source": [
    "## Embeddings\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eec6dc3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.0, 2.0, 3.0]\n",
      "[1.0, 2.0, 3.0]\n"
     ]
    }
   ],
   "source": [
    "from graphrag_llm.embedding import LLMEmbedding, create_embedding\n",
    "\n",
    "embedding_config = ModelConfig(\n",
    "    type=LLMProviderType.MockLLM,\n",
    "    model_provider=\"openai\",\n",
    "    model=\"text-embedding-3-small\",\n",
    "    mock_responses=[1.0, 2.0, 3.0],\n",
    ")\n",
    "\n",
    "llm_embedding: LLMEmbedding = create_embedding(embedding_config)\n",
    "\n",
    "embeddings_response = llm_embedding.embedding(input=[\"Hello world\", \"How are you?\"])\n",
    "for embedding in embeddings_response.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
}