diff --git a/docs/docs/api_reference/embeddings/upstage.md b/docs/docs/api_reference/embeddings/upstage.md
new file mode 100644
index 0000000000000000000000000000000000000000..ab961b59ebc7d588aed8db130b6cdf93a5259d7b
--- /dev/null
+++ b/docs/docs/api_reference/embeddings/upstage.md
@@ -0,0 +1,4 @@
+::: llama_index.embeddings.upstage
+ options:
+ members:
+ - UpstageEmbedding
diff --git a/docs/docs/api_reference/llms/upstage.md b/docs/docs/api_reference/llms/upstage.md
new file mode 100644
index 0000000000000000000000000000000000000000..ea36317231ea6768aacbfd2a370830e34183cafd
--- /dev/null
+++ b/docs/docs/api_reference/llms/upstage.md
@@ -0,0 +1,4 @@
+::: llama_index.llms.upstage
+ options:
+ members:
+ - Upstage
diff --git a/docs/docs/examples/embeddings/upstage.ipynb b/docs/docs/examples/embeddings/upstage.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..e1252bee5acc2d1c05be96c8680fc596ca112deb
--- /dev/null
+++ b/docs/docs/examples/embeddings/upstage.ipynb
@@ -0,0 +1,241 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/embeddings/upstage.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+ ]
+ },
+ {
+ "attachments": {},
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Upstage Embeddings"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%pip install llama-index-embeddings-upstage==0.1.0"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "!pip install llama-index"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "\n",
+ "os.environ[\"UPSTAGE_API_KEY\"] = \"YOUR_API_KEY\""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.embeddings.upstage import UpstageEmbedding\n",
+ "from llama_index.core import Settings\n",
+ "\n",
+ "embed_model = UpstageEmbedding()\n",
+ "Settings.embed_model = embed_model"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Using Upstage `solar-1-mini-embedding`\n",
+ "\n",
+ "Note, you may have to update your openai client: `pip install -U openai`"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# get API key and create embeddings\n",
+ "from llama_index.embeddings.upstage import UpstageEmbedding\n",
+ "\n",
+ "embed_model = UpstageEmbedding()\n",
+ "\n",
+ "embeddings = embed_model.get_text_embedding(\n",
+ " \"Upstage new Embeddings models is great.\"\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "[0.02535058930516243, 0.007272760849446058, 0.015372460708022118, -0.007840132340788841, 0.0017625312320888042]\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(embeddings[:5])"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "4096\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(len(embeddings))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "embeddings = embed_model.get_query_embedding(\n",
+ " \"What are some great Embeddings model?\"\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "[0.03518765792250633, 0.01018011849373579, 0.013282101601362228, -0.008568626828491688, -0.005505830980837345]\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(embeddings[:5])"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "4096\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(len(embeddings))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# embed documents\n",
+ "embeddings = embed_model.get_text_embedding_batch(\n",
+ " [\n",
+ " \"Upstage new Embeddings models is awesome.\",\n",
+ " \"Upstage LLM is also awesome.\",\n",
+ " ]\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "2\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(len(embeddings))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "[0.028246860951185226, 0.008945596404373646, 0.01719627156853676, -0.005711239762604237, 0.0016300849383696914]\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(embeddings[0][:5])"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "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"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}
diff --git a/docs/docs/examples/llm/solar.ipynb b/docs/docs/examples/llm/solar.ipynb
index b0ea595833147fb06e6a86f92d6fe7ebffa4bb02..8519c63aed4a3d9715cd40ded3c67205fa6f3864 100644
--- a/docs/docs/examples/llm/solar.ipynb
+++ b/docs/docs/examples/llm/solar.ipynb
@@ -5,7 +5,10 @@
"id": "cae1b4a8",
"metadata": {},
"source": [
- "# Solar LLM"
+ "# Solar LLM\n",
+ "\n",
+ "Warning: Solar LLM is deprecated. Please use Upstage LLM instead. \n",
+ "See [Upstage LLM](https://docs.llamaindex.ai/en/stable/examples/llm/upstage/)"
]
},
{
diff --git a/docs/docs/examples/llm/upstage.ipynb b/docs/docs/examples/llm/upstage.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..3e4021c9769507237cc8c6732114f2cdc8ac1cd0
--- /dev/null
+++ b/docs/docs/examples/llm/upstage.ipynb
@@ -0,0 +1,578 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "9fd54a32",
+ "metadata": {},
+ "source": [
+ "<a href=\"https://colab.research.google.com/github/jerryjliu/llama_index/blob/main/docs/docs/examples/llm/upstage.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "9e3a8796-edc8-43f2-94ad-fe4fb20d70ed",
+ "metadata": {},
+ "source": [
+ "# Upstage"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "081d07d2",
+ "metadata": {},
+ "source": [
+ "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "3f6f8702",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%pip install llama-index-llms-upstage==0.1.0"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "83ea30ee",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "!pip install llama-index"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "b007403c-6b7a-420c-92f1-4171d05ed9bb",
+ "metadata": {},
+ "source": [
+ "## Basic Usage"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "8ead155e-b8bd-46f9-ab9b-28fc009361dd",
+ "metadata": {},
+ "source": [
+ "#### Call `complete` with a prompt"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "859f7447-bbc3-4036-b541-a90854a732f1",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "\n",
+ "os.environ[\"UPSTAGE_API_KEY\"] = \"YOUR_API_KEY\""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "60be18ae-c957-4ac2-a58a-0652e18ee6d6",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.llms.upstage import Upstage\n",
+ "\n",
+ "resp = Upstage().complete(\"Paul Graham is \")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "ac2cbebb-a444-4a46-9d85-b265a3483d68",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Paul Graham is a computer scientist, entrepreneur, and essayist. He is best known as the co-founder of the venture capital firm Y Combinator, which has funded and incubated many successful startups. He is also the author of several influential essays on entrepreneurship, startup culture, and technology.\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(resp)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "14831268-f90f-499d-9d86-925dbc88292b",
+ "metadata": {},
+ "source": [
+ "#### Call `chat` with a list of messages"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "bbe29574-4af1-48d5-9739-f60652b6ce6c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core.llms import ChatMessage\n",
+ "from llama_index.llms.upstage import Upstage\n",
+ "\n",
+ "messages = [\n",
+ " ChatMessage(\n",
+ " role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
+ " ),\n",
+ " ChatMessage(role=\"user\", content=\"What is your name\"),\n",
+ "]\n",
+ "resp = Upstage().chat(messages)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "9cbd550a-0264-4a11-9b2c-a08d8723a5ae",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "assistant: I am Captain Redbeard, the fearless pirate!\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(resp)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "2ed5e894-4597-4911-a623-591560f72b82",
+ "metadata": {},
+ "source": [
+ "## Streaming"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "4cb7986f-aaed-42e2-abdd-f274f6d4fc59",
+ "metadata": {},
+ "source": [
+ "Using `stream_complete` endpoint"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "d43f17a2-0aeb-464b-a7a7-732ba5e8ef24",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.llms.upstage import Upstage\n",
+ "\n",
+ "llm = Upstage()\n",
+ "resp = llm.stream_complete(\"Paul Graham is \")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "0214e911-cf0d-489c-bc48-9bb1d8bf65d8",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Paul Graham is a computer scientist, entrepreneur, and essayist. He is best known for co-founding the startup accelerator Y Combinator, which has helped launch some of the most successful tech companies in the world, including Airbnb, Dropbox, and Stripe. He is also the author of several influential essays on startup culture and entrepreneurship, including \"How to Start a Startup\" and \"Hackers & Painters.\""
+ ]
+ }
+ ],
+ "source": [
+ "for r in resp:\n",
+ " print(r.delta, end=\"\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "40350dd8-3f50-4a2f-8545-5723942039bb",
+ "metadata": {},
+ "source": [
+ "Using `stream_chat` endpoint"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "bc636e65-a67b-4dcd-ac60-b25abc9d8dbd",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core.llms import ChatMessage\n",
+ "from llama_index.llms.upstage import Upstage\n",
+ "\n",
+ "llm = Upstage()\n",
+ "messages = [\n",
+ " ChatMessage(\n",
+ " role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
+ " ),\n",
+ " ChatMessage(role=\"user\", content=\"What is your name\"),\n",
+ "]\n",
+ "resp = llm.stream_chat(messages)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "4475a6bc-1051-4287-abce-ba83324aeb9e",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "I am Captain Redbeard, the fearless pirate!"
+ ]
+ }
+ ],
+ "source": [
+ "for r in resp:\n",
+ " print(r.delta, end=\"\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "009d3f1c-ef35-4126-ae82-0b97adb746e3",
+ "metadata": {},
+ "source": [
+ "## Configure Model"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "e973e3d1-a3c9-43b9-bee1-af3e57946ac3",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.llms.upstage import Upstage\n",
+ "\n",
+ "llm = Upstage()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "e2c9bcf6-c950-4dfc-abdc-598d5bdedf40",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "resp = llm.complete(\"Paul Graham is \")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "2edc85ca-df17-4774-a3ea-e80109fa1811",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Paul Graham is a computer scientist, entrepreneur, and essayist. He is best known for co-founding the startup accelerator Y Combinator, which has helped launch some of the most successful tech companies in the world, including Airbnb, Dropbox, and Stripe. He is also the author of several influential essays on entrepreneurship, startups, and technology, including \"How to Start a Startup\" and \"Hackers & Painters.\"\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(resp)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "026fdb77-b61f-4571-8eaf-4a51e8415458",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "messages = [\n",
+ " ChatMessage(\n",
+ " role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
+ " ),\n",
+ " ChatMessage(role=\"user\", content=\"What is your name\"),\n",
+ "]\n",
+ "resp = llm.chat(messages)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "2286a16c-188b-437f-a1a3-4efe299b759d",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "assistant: I am Captain Redbeard, the fearless pirate!\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(resp)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "90f07f7e-927f-47a2-9797-de5a86d61e1f",
+ "metadata": {},
+ "source": [
+ "## Function Calling\n",
+ "\n",
+ "Upstage models have native support for function calling. This conveniently integrates with LlamaIndex tool abstractions, letting you plug in any arbitrary Python function to the LLM."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "212bb2d2-2bed-4188-85ad-3cd497d4b864",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from pydantic import BaseModel\n",
+ "from llama_index.core.tools import FunctionTool\n",
+ "\n",
+ "\n",
+ "class Song(BaseModel):\n",
+ " \"\"\"A song with name and artist\"\"\"\n",
+ "\n",
+ " name: str\n",
+ " artist: str\n",
+ "\n",
+ "\n",
+ "def generate_song(name: str, artist: str) -> Song:\n",
+ " \"\"\"Generates a song with provided name and artist.\"\"\"\n",
+ " return Song(name=name, artist=artist)\n",
+ "\n",
+ "\n",
+ "tool = FunctionTool.from_defaults(fn=generate_song)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "ba29ea39-39cd-42f4-934e-78fba4935a1c",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "name='My Song' artist='John Doe'\n"
+ ]
+ }
+ ],
+ "source": [
+ "from llama_index.llms.upstage import Upstage\n",
+ "\n",
+ "llm = Upstage()\n",
+ "response = llm.predict_and_call([tool], \"Generate a song\")\n",
+ "print(str(response))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "629bdb64-e5bd-4539-9537-54352dd8dbb4",
+ "metadata": {},
+ "source": [
+ "We can also do multiple function calling."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "c098cb38-bfed-4907-8109-6810756f2ae5",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Name: generate_song, Input: {'args': (), 'kwargs': {'name': 'Beatles', 'artist': 'Beatles'}}, Output: name='Beatles' artist='Beatles'\n"
+ ]
+ }
+ ],
+ "source": [
+ "llm = Upstage()\n",
+ "response = llm.predict_and_call(\n",
+ " [tool],\n",
+ " \"Generate five songs from the Beatles\",\n",
+ " allow_parallel_tool_calls=True,\n",
+ ")\n",
+ "for s in response.sources:\n",
+ " print(f\"Name: {s.tool_name}, Input: {s.raw_input}, Output: {str(s)}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "df5fa1ab-f598-46da-80f3-f6af5dd2fe83",
+ "metadata": {},
+ "source": [
+ "## Async"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "3e088b90-12b6-4211-a9ca-696e9897e9ae",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.llms.upstage import Upstage\n",
+ "\n",
+ "llm = Upstage()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "8f649683-896a-439e-b14b-e5df5d803b82",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "resp = await llm.acomplete(\"Paul Graham is \")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "d05d8b7c-07b3-4ce8-9a8c-fcd1e830316c",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Paul Graham is a computer scientist, entrepreneur, and essayist. He is best known as the co-founder of the startup accelerator Y Combinator, which has helped launch and fund many successful tech companies. He is also the author of several influential essays on startups, entrepreneurship, and technology, including \"How to Start a Startup\" and \"Hackers & Painters.\"\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(resp)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "abc72d09-3bcd-4d48-9bb7-0920c56310c1",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "resp = await llm.astream_complete(\"Paul Graham is \")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b28d7a06-1518-4ec0-b3cc-6364395b3561",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Paul Graham is a computer scientist, entrepreneur, and essayist. He is best known as the co-founder of the startup accelerator Y Combinator, which has helped launch some of the most successful tech companies in the world, including Airbnb, Dropbox, and Stripe. Graham is also a prolific writer, and his essays on topics such as startup advice, artificial intelligence, and the future of work have been widely read and influential in the tech industry."
+ ]
+ }
+ ],
+ "source": [
+ "async for delta in resp:\n",
+ " print(delta.delta, end=\"\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "3782e752-3c65-4b11-b2c6-8efe55fc33b1",
+ "metadata": {},
+ "source": [
+ "Async function calling is also supported."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "7dcfaa77-edfe-42f9-8138-884c99dd3e62",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "name='My Song' artist='Me'\n"
+ ]
+ }
+ ],
+ "source": [
+ "llm = Upstage()\n",
+ "response = await llm.apredict_and_call([tool], \"Generate a song\")\n",
+ "print(str(response))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "a2782f06",
+ "metadata": {},
+ "source": [
+ "## Set API Key at a per-instance level\n",
+ "If desired, you can have separate LLM instances using separate API keys."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "015c2d39",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Paul Graham is a computer scientist, entrepreneur, and essayist. He is best known as the co-founder of the startup accelerator Y Combinator, which has helped launch and fund many successful tech companies. He is also the author of several influential essays on startup culture and technology, including \"How to Start a Startup\" and \"Hackers & Painters.\"\n"
+ ]
+ }
+ ],
+ "source": [
+ "from llama_index.llms.upstage import Upstage\n",
+ "\n",
+ "llm = Upstage(api_key=\"BAD_KEY\")\n",
+ "resp = Upstage().complete(\"Paul Graham is \")\n",
+ "print(resp)"
+ ]
+ }
+ ],
+ "metadata": {
+ "colab": {
+ "provenance": []
+ },
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "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"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/.gitignore b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..990c18de229088f55c6c514fd0f2d49981d1b0e7
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/.gitignore
@@ -0,0 +1,153 @@
+llama_index/_static
+.DS_Store
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+bin/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+etc/
+include/
+lib/
+lib64/
+parts/
+sdist/
+share/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+.ruff_cache
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+notebooks/
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+# However, in case of collaboration, if having platform-specific dependencies or dependencies
+# having no cross-platform support, pipenv may install dependencies that don't work, or not
+# install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+pyvenv.cfg
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Jetbrains
+.idea
+modules/
+*.swp
+
+# VsCode
+.vscode
+
+# pipenv
+Pipfile
+Pipfile.lock
+
+# pyright
+pyrightconfig.json
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/BUILD b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/BUILD
new file mode 100644
index 0000000000000000000000000000000000000000..0896ca890d8bffd60a44fa824f8d57fecd73ee53
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/BUILD
@@ -0,0 +1,3 @@
+poetry_requirements(
+ name="poetry",
+)
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/Makefile b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/Makefile
new file mode 100644
index 0000000000000000000000000000000000000000..2b045fb919d9135c6566736c3334da58147e7de2
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/Makefile
@@ -0,0 +1,20 @@
+GIT_ROOT ?= $(shell git rev-parse --show-toplevel)
+
+help: ## Show all Makefile targets.
+ @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
+
+format: ## Run code autoformatters (black).
+ pre-commit install
+ git ls-files | xargs pre-commit run black --files
+
+lint: ## Run linters: pre-commit (black, ruff, codespell) and mypy
+ pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
+
+test: ## Run tests via pytest.
+ pytest tests/unit_tests
+
+integration_test: ## Run integration tests via pytest.
+ pytest tests/integration_tests
+
+watch-docs: ## Build and watch documentation.
+ sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/README.md b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..bab4d2639842119444dc87e4e66a3e60faf9cfa5
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/README.md
@@ -0,0 +1,9 @@
+# LlamaIndex Embeddings Integration: Upstage
+
+[Upstage](https://www.upstage.ai/) is a leading artificial intelligence (AI) company specializing in delivering above-human-grade performance LLM components.
+
+## Solar Embeddings
+
+The embeddings API converts text into numbers that computers can understand. Imagine converting a sentence into a list of numbers, each capturing a piece of the sentence's meaning. This makes it easier for machines to do tasks like finding similar texts, sorting information, or even answering questions.
+
+Check more details [here](https://developers.upstage.ai/docs/apis/embeddings).
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/llama_index/embeddings/upstage/BUILD b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/llama_index/embeddings/upstage/BUILD
new file mode 100644
index 0000000000000000000000000000000000000000..db46e8d6c978c67e301dd6c47bee08c1b3fd141c
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/llama_index/embeddings/upstage/BUILD
@@ -0,0 +1 @@
+python_sources()
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/llama_index/embeddings/upstage/__init__.py b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/llama_index/embeddings/upstage/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8086dac1ac1697a0cdd63e123639e0e5cc581837
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/llama_index/embeddings/upstage/__init__.py
@@ -0,0 +1,4 @@
+from llama_index.embeddings.upstage.base import UpstageEmbedding
+
+
+__all__ = ["UpstageEmbedding"]
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/llama_index/embeddings/upstage/base.py b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/llama_index/embeddings/upstage/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac2e32ecc84612c5c445355af778d7a32a4934eb
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/llama_index/embeddings/upstage/base.py
@@ -0,0 +1,219 @@
+import logging
+import warnings
+from typing import Dict, Any, Optional, List, Tuple
+
+import httpx
+from llama_index.core.callbacks import CallbackManager
+from llama_index.embeddings.openai import OpenAIEmbedding
+from pydantic import Field
+
+from llama_index.embeddings.upstage.utils import (
+ DEFAULT_UPSTAGE_API_BASE,
+ resolve_upstage_credentials,
+)
+
+logger = logging.getLogger(__name__)
+
+UPSTAGE_EMBEDDING_MODELS = {
+ "solar-1-mini-embedding": {
+ "query": "solar-1-mini-embedding-query",
+ "passage": "solar-1-mini-embedding-passage",
+ },
+ "solar-embedding-1-large": {
+ "query": "solar-embedding-1-large-query",
+ "passage": "solar-embedding-1-large-passage",
+ },
+}
+
+MAX_EMBED_BATCH_SIZE = 100
+
+
+def get_engine(model) -> Tuple[Any, Any]:
+ """
+ get query engine and passage engine for the model.
+ """
+ if model not in UPSTAGE_EMBEDDING_MODELS:
+ raise ValueError(
+ f"Unknown model: {model}. Please provide a valid Upstage model name in: {', '.join(UPSTAGE_EMBEDDING_MODELS.keys())}"
+ )
+ return (
+ UPSTAGE_EMBEDDING_MODELS[model]["query"],
+ UPSTAGE_EMBEDDING_MODELS[model]["passage"],
+ )
+
+
+class UpstageEmbedding(OpenAIEmbedding):
+ """
+ Class for Upstage embeddings.
+
+ Args:
+ model_name (str): Model for embedding.
+ Defaults to "solar-1-mini-embedding".
+ """
+
+ additional_kwargs: Dict[str, Any] = Field(
+ default_factory=dict, description="Additional kwargs for the Upstage API."
+ )
+
+ api_key: str = Field(description="The Upstage API key.")
+ api_base: Optional[str] = Field(
+ default=DEFAULT_UPSTAGE_API_BASE, description="The base URL for OpenAI API."
+ )
+ dimensions: Optional[int] = Field(
+ None,
+ description="Not supported yet. The number of dimensions the resulting output embeddings should have.",
+ )
+
+ def __init__(
+ self,
+ model: str = "solar-embedding-1-large",
+ embed_batch_size: int = 100,
+ dimensions: Optional[int] = None,
+ additional_kwargs: Dict[str, Any] = None,
+ api_key: Optional[str] = None,
+ api_base: Optional[str] = None,
+ max_retries: int = 10,
+ timeout: float = 60.0,
+ reuse_client: bool = True,
+ callback_manager: Optional[CallbackManager] = None,
+ default_headers: Optional[Dict[str, str]] = None,
+ http_client: Optional[httpx.Client] = None,
+ **kwargs: Any,
+ ) -> None:
+ additional_kwargs = additional_kwargs or {}
+ if dimensions is not None:
+ warnings.warn("Received dimensions argument. This is not supported yet.")
+ additional_kwargs["dimensions"] = dimensions
+
+ if embed_batch_size > MAX_EMBED_BATCH_SIZE:
+ raise ValueError(
+ f"embed_batch_size should be less than or equal to {MAX_EMBED_BATCH_SIZE}."
+ )
+
+ api_key, api_base = resolve_upstage_credentials(
+ api_key=api_key, api_base=api_base
+ )
+
+ if "model_name" in kwargs:
+ model = kwargs.pop("model_name")
+
+ # if model endswith with "-query" or "-passage", remove the suffix and print a warning
+ if model.endswith(("-query", "-passage")):
+ model = model.rsplit("-", 1)[0]
+ logger.warning(
+ f"Model name should not end with '-query' or '-passage'. The suffix has been removed. "
+ f"Model name: {model}"
+ )
+
+ super().__init__(
+ embed_batch_size=embed_batch_size,
+ dimensions=dimensions,
+ callback_manager=callback_manager,
+ model_name=model,
+ additional_kwargs=additional_kwargs,
+ api_key=api_key,
+ api_base=api_base,
+ max_retries=max_retries,
+ reuse_client=reuse_client,
+ timeout=timeout,
+ default_headers=default_headers,
+ **kwargs,
+ )
+ self._client = None
+ self._aclient = None
+ self._http_client = http_client
+
+ self._query_engine, self._text_engine = get_engine(model)
+
+ def class_name(cls) -> str:
+ return "UpstageEmbedding"
+
+ def _get_credential_kwargs(self) -> Dict[str, Any]:
+ return {
+ "api_key": self.api_key,
+ "base_url": self.api_base,
+ "max_retries": self.max_retries,
+ "timeout": self.timeout,
+ "default_headers": self.default_headers,
+ "http_client": self._http_client,
+ }
+
+ def _get_query_embedding(self, query: str) -> List[float]:
+ """Get query embedding."""
+ client = self._get_client()
+ text = query.replace("\n", " ")
+ return (
+ client.embeddings.create(
+ input=text, model=self._query_engine, **self.additional_kwargs
+ )
+ .data[0]
+ .embedding
+ )
+
+ async def _aget_query_embedding(self, query: str) -> List[float]:
+ """The asynchronous version of _get_query_embedding."""
+ client = self._get_aclient()
+ text = query.replace("\n", " ")
+ return (
+ (
+ await client.embeddings.create(
+ input=text, model=self._query_engine, **self.additional_kwargs
+ )
+ )
+ .data[0]
+ .embedding
+ )
+
+ def _get_text_embedding(self, text: str) -> List[float]:
+ """Get text embedding."""
+ client = self._get_client()
+ return (
+ client.embeddings.create(
+ input=text, model=self._text_engine, **self.additional_kwargs
+ )
+ .data[0]
+ .embedding
+ )
+
+ async def _aget_text_embedding(self, text: str) -> List[float]:
+ """Asynchronously get text embedding."""
+ client = self._get_aclient()
+ return (
+ (
+ await client.embeddings.create(
+ input=text, model=self._text_engine, **self.additional_kwargs
+ )
+ )
+ .data[0]
+ .embedding
+ )
+
+ def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+ """Get text embeddings."""
+ client = self._get_client()
+ batch_size = min(self.embed_batch_size, len(texts))
+ texts = [text.replace("\n", " ") for text in texts]
+
+ embeddings = []
+ for i in range(0, len(texts), batch_size):
+ batch = texts[i : i + batch_size]
+ response = client.embeddings.create(
+ input=batch, model=self._text_engine, **self.additional_kwargs
+ )
+ embeddings.extend([r.embedding for r in response.data])
+ return embeddings
+
+ async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+ """Asynchronously get text embeddings."""
+ client = self._get_aclient()
+ batch_size = min(self.embed_batch_size, len(texts))
+ texts = [text.replace("\n", " ") for text in texts]
+
+ embeddings = []
+ for i in range(0, len(texts), batch_size):
+ batch = texts[i : i + batch_size]
+ response = await client.embeddings.create(
+ input=batch, model=self._text_engine, **self.additional_kwargs
+ )
+ embeddings.extend([r.embedding for r in response.data])
+ return embeddings
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/llama_index/embeddings/upstage/utils.py b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/llama_index/embeddings/upstage/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..543aab2d98e3fe8b2a1961f02b3b5d0e8969f69e
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/llama_index/embeddings/upstage/utils.py
@@ -0,0 +1,74 @@
+import logging
+from typing import Optional, Tuple, Callable, Any
+
+import openai
+from llama_index.core.base.llms.generic_utils import get_from_param_or_env
+from tenacity import (
+ wait_random_exponential,
+ wait_exponential,
+ stop_after_attempt,
+ retry_if_exception_type,
+ retry,
+ before_sleep_log,
+)
+from tenacity.stop import stop_base, stop_after_delay
+
+DEFAULT_UPSTAGE_API_BASE = "https://api.upstage.ai/v1/solar"
+
+logger = logging.getLogger(__name__)
+
+
+def resolve_upstage_credentials(
+ api_key: Optional[str] = None, api_base: Optional[str] = None
+) -> Tuple[Optional[str], str]:
+ """Resolve Upstage credentials.
+
+ The order of precedence is:
+ 1. param
+ 2. env
+ 4. default
+ """
+ # resolve from param or env
+ api_key = get_from_param_or_env("api_key", api_key, "UPSTAGE_API_KEY", "")
+ api_base = get_from_param_or_env("api_base", api_base, "UPSTAGE_API_BASE", "")
+
+ final_api_key = api_key or ""
+ final_api_base = api_base or DEFAULT_UPSTAGE_API_BASE
+
+ return final_api_key, str(final_api_base)
+
+
+def create_retry_decorator(
+ max_retries: int,
+ random_exponential: bool = False,
+ stop_after_delay_seconds: Optional[float] = None,
+ min_seconds: float = 4,
+ max_seconds: float = 10,
+) -> Callable[[Any], Any]:
+ wait_strategy = (
+ wait_random_exponential(min=min_seconds, max=max_seconds)
+ if random_exponential
+ else wait_exponential(multiplier=1, min=min_seconds, max=max_seconds)
+ )
+
+ stop_strategy: stop_base = stop_after_attempt(max_retries)
+ if stop_after_delay_seconds is not None:
+ stop_strategy = stop_strategy | stop_after_delay(stop_after_delay_seconds)
+
+ return retry(
+ reraise=True,
+ stop=stop_strategy,
+ wait=wait_strategy,
+ retry=(
+ retry_if_exception_type(
+ (
+ openai.APITimeoutError,
+ openai.APIError,
+ openai.APIConnectionError,
+ openai.RateLimitError,
+ openai.APIStatusError,
+ )
+ )
+ ),
+ before_sleep=before_sleep_log(logger, logging.WARNING),
+ )
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/pyproject.toml b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/pyproject.toml
new file mode 100644
index 0000000000000000000000000000000000000000..348618780bd8bbbe1d59e8da4e6791ebd8f8b2e3
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/pyproject.toml
@@ -0,0 +1,76 @@
+[build-system]
+build-backend = "poetry.core.masonry.api"
+requires = ["poetry-core"]
+
+[tool.codespell]
+check-filenames = true
+check-hidden = true
+# Feel free to un-skip examples, and experimental, you will just need to
+# work through many typos (--write-changes and --interactive will help)
+skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb"
+
+[tool.llamahub]
+contains_example = false
+import_path = "llama_index.embeddings.upstage"
+
+[tool.llamahub.class_authors]
+UpstageEmbedding = "llama-index"
+
+[tool.mypy]
+disallow_untyped_defs = true
+# Remove venv skip when integrated with pre-commit
+exclude = ["_static", "build", "examples", "notebooks", "venv"]
+ignore_missing_imports = true
+python_version = "3.8"
+
+[tool.poetry]
+authors = ["Your Name <you@example.com>"]
+description = "llama-index embeddings upstage integration"
+license = "MIT"
+name = "llama-index-embeddings-upstage"
+packages = [{include = "llama_index/"}]
+readme = "README.md"
+version = "0.1.0"
+
+[tool.poetry.dependencies]
+python = ">=3.8.1,<4.0"
+llama-index-core = "^0.10.0"
+llama-index-embeddings-openai = "^0.1.9"
+
+[tool.poetry.group.dev.dependencies]
+black = {extras = ["jupyter"], version = "<=23.9.1,>=23.7.0"}
+codespell = {extras = ["toml"], version = ">=v2.2.6"}
+ipython = "8.10.0"
+jupyter = "^1.0.0"
+mypy = "0.991"
+pre-commit = "3.2.0"
+pylint = "2.15.10"
+pytest = "7.2.1"
+pytest-asyncio = "^0.23.6"
+pytest-mock = "3.11.1"
+ruff = "0.0.292"
+tree-sitter-languages = "^1.8.0"
+types-Deprecated = ">=0.1.0"
+types-PyYAML = "^6.0.12.12"
+types-protobuf = "^4.24.0.4"
+types-redis = "4.5.5.0"
+types-requests = "2.28.11.8" # TODO: unpin when mypy>0.991
+types-setuptools = "67.1.0.0"
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+# --strict-markers will raise errors on unknown marks.
+# https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks
+#
+# https://docs.pytest.org/en/7.1.x/reference/reference.html
+# --strict-config any warnings encountered while parsing the `pytest`
+# section of the configuration file raise errors.
+#
+# https://github.com/tophat/syrupy
+# --snapshot-warn-unused Prints a warning on unused snapshots rather than fail the test suite.
+# Registering custom markers.
+markers = [
+ "asyncio: mark tests as requiring asyncio",
+ "compile: mark placeholder test used to compile integration tests without running them",
+ "requires: mark tests as requiring a specific library",
+]
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/__init__.py b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/integration_tests/BUILD b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/integration_tests/BUILD
new file mode 100644
index 0000000000000000000000000000000000000000..57341b1358b565681df03bff7c9a81b9270fece2
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/integration_tests/BUILD
@@ -0,0 +1,3 @@
+python_tests(
+ name="tests",
+)
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/integration_tests/test_integrations.py b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/integration_tests/test_integrations.py
new file mode 100644
index 0000000000000000000000000000000000000000..ccf7e6cc3825dd3e1cc4d1288c2bcbe9fc36c812
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/integration_tests/test_integrations.py
@@ -0,0 +1,64 @@
+import os
+
+import pytest
+
+
+@pytest.fixture()
+def upstage_embedding():
+ UpstageEmbedding = pytest.importorskip(
+ "llama_index.embeddings.upstage", reason="Cannot import UpstageEmbedding"
+ ).UpstageEmbedding
+
+ if os.getenv("UPSTAGE_API_KEY") is None:
+ pytest.skip("UPSTAGE_API_KEY is not set.")
+ return UpstageEmbedding()
+
+
+def test_upstage_embedding_query_embedding(upstage_embedding):
+ query = "hello"
+ embedding = upstage_embedding.get_query_embedding(query)
+ assert isinstance(embedding, list)
+
+
+async def test_upstage_embedding_async_query_embedding(upstage_embedding):
+ query = "hello"
+ embedding = await upstage_embedding.aget_query_embedding(query)
+ assert isinstance(embedding, list)
+
+
+def test_upstage_embedding_text_embedding(upstage_embedding):
+ text = "hello"
+ embedding = upstage_embedding.get_text_embedding(text)
+ assert isinstance(embedding, list)
+
+
+async def test_upstage_embedding_async_text_embedding(upstage_embedding):
+ text = "hello"
+ embedding = await upstage_embedding.aget_text_embedding(text)
+ assert isinstance(embedding, list)
+
+
+def test_upstage_embedding_text_embeddings(upstage_embedding):
+ texts = ["hello", "world"]
+ embeddings = upstage_embedding.get_text_embedding_batch(texts)
+ assert isinstance(embeddings, list)
+ assert len(embeddings) == len(texts)
+ assert all(isinstance(embedding, list) for embedding in embeddings)
+
+
+def test_upstage_embedding_text_embeddings_fail_large_batch():
+ UpstageEmbedding = pytest.importorskip(
+ "llama_index.embeddings.upstage", reason="Cannot import UpstageEmbedding"
+ ).UpstageEmbedding
+ texts = ["hello"] * 2049
+ with pytest.raises(ValueError):
+ upstage_embedding = UpstageEmbedding(embed_batch_size=2049)
+ upstage_embedding.get_text_embedding_batch(texts)
+
+
+async def test_upstage_embedding_async_text_embeddings(upstage_embedding):
+ texts = ["hello", "world"]
+ embeddings = await upstage_embedding.aget_text_embedding_batch(texts)
+ assert isinstance(embeddings, list)
+ assert len(embeddings) == len(texts)
+ assert all(isinstance(embedding, list) for embedding in embeddings)
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/unit_tests/BUILD b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/unit_tests/BUILD
new file mode 100644
index 0000000000000000000000000000000000000000..57341b1358b565681df03bff7c9a81b9270fece2
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/unit_tests/BUILD
@@ -0,0 +1,3 @@
+python_tests(
+ name="tests",
+)
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/unit_tests/test_embeddings_upstage.py b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/unit_tests/test_embeddings_upstage.py
new file mode 100644
index 0000000000000000000000000000000000000000..04de64a45a7b1819e2f1b3f617eb7fd3a0c18930
--- /dev/null
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-upstage/tests/unit_tests/test_embeddings_upstage.py
@@ -0,0 +1,19 @@
+import pytest
+from llama_index.core.base.embeddings.base import BaseEmbedding
+
+
+@pytest.fixture()
+def upstage_embedding():
+ return pytest.importorskip(
+ "llama_index.embeddings.upstage", reason="Cannot import UpstageEmbedding"
+ ).UpstageEmbedding
+
+
+def test_upstage_embedding_class(upstage_embedding):
+ names_of_base_classes = [b.__name__ for b in upstage_embedding.__mro__]
+ assert BaseEmbedding.__name__ in names_of_base_classes
+
+
+def upstage_embedding_fail_wrong_model(upstage_embedding):
+ with pytest.raises(ValueError):
+ upstage_embedding(model="foo")
diff --git a/llama-index-integrations/llms/llama-index-llms-solar/llama_index/llms/solar/base.py b/llama-index-integrations/llms/llama-index-llms-solar/llama_index/llms/solar/base.py
index c954d051adbee3db77f168d8fd3a2c4df3c0ff03..5f97113565d15850dad0aed70dc28e3df06c0b85 100644
--- a/llama-index-integrations/llms/llama-index-llms-solar/llama_index/llms/solar/base.py
+++ b/llama-index-integrations/llms/llama-index-llms-solar/llama_index/llms/solar/base.py
@@ -1,3 +1,4 @@
+import warnings
from typing import (
Any,
Callable,
@@ -90,6 +91,12 @@ class Solar(OpenAI):
output_parser: Optional[BaseOutputParser] = None,
**kwargs: Any,
) -> None:
+ # add warning for this class is deprecated
+ warnings.warn(
+ """Solar LLM is deprecated. Please use Upstage LLM instead.
+ Install the package using `pip install llama-index-llms-upstage`
+ """,
+ )
api_key, api_base = resolve_solar_credentials(
api_key=api_key,
api_base=api_base,
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/.gitignore b/llama-index-integrations/llms/llama-index-llms-upstage/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..990c18de229088f55c6c514fd0f2d49981d1b0e7
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-upstage/.gitignore
@@ -0,0 +1,153 @@
+llama_index/_static
+.DS_Store
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+bin/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+etc/
+include/
+lib/
+lib64/
+parts/
+sdist/
+share/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+.ruff_cache
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+notebooks/
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+# However, in case of collaboration, if having platform-specific dependencies or dependencies
+# having no cross-platform support, pipenv may install dependencies that don't work, or not
+# install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+pyvenv.cfg
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Jetbrains
+.idea
+modules/
+*.swp
+
+# VsCode
+.vscode
+
+# pipenv
+Pipfile
+Pipfile.lock
+
+# pyright
+pyrightconfig.json
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/BUILD b/llama-index-integrations/llms/llama-index-llms-upstage/BUILD
new file mode 100644
index 0000000000000000000000000000000000000000..0896ca890d8bffd60a44fa824f8d57fecd73ee53
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-upstage/BUILD
@@ -0,0 +1,3 @@
+poetry_requirements(
+ name="poetry",
+)
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/Makefile b/llama-index-integrations/llms/llama-index-llms-upstage/Makefile
new file mode 100644
index 0000000000000000000000000000000000000000..b9eab05aa370629a4a3de75df3ff64cd53887b68
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-upstage/Makefile
@@ -0,0 +1,17 @@
+GIT_ROOT ?= $(shell git rev-parse --show-toplevel)
+
+help: ## Show all Makefile targets.
+ @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
+
+format: ## Run code autoformatters (black).
+ pre-commit install
+ git ls-files | xargs pre-commit run black --files
+
+lint: ## Run linters: pre-commit (black, ruff, codespell) and mypy
+ pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
+
+test: ## Run tests via pytest.
+ pytest tests
+
+watch-docs: ## Build and watch documentation.
+ sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/README.md b/llama-index-integrations/llms/llama-index-llms-upstage/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..dee59891494b6f63e77b8f1b0b014876c41dcf1a
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-upstage/README.md
@@ -0,0 +1,9 @@
+# LlamaIndex Llms Integration: Upstage
+
+[Upstage](https://www.upstage.ai/) is a leading artificial intelligence (AI) company specializing in delivering above-human-grade performance LLM components.
+
+## Solar LLM
+
+Solar Mini Chat is a fast yet powerful advanced large language model focusing on English and Korean. It has been specifically fine-tuned for multi-turn chat purposes, showing enhanced performance across a wide range of natural language processing tasks, like multi-turn conversation or tasks that require an understanding of long contexts, such as RAG (Retrieval-Augmented Generation), compared to other models of a similar size. This fine-tuning equips it with the ability to handle longer conversations more effectively, making it particularly adept for interactive applications.
+
+Check more details [here](https://developers.upstage.ai/docs/getting-started/quick-start).
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/llama_index/llms/upstage/BUILD b/llama-index-integrations/llms/llama-index-llms-upstage/llama_index/llms/upstage/BUILD
new file mode 100644
index 0000000000000000000000000000000000000000..db46e8d6c978c67e301dd6c47bee08c1b3fd141c
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-upstage/llama_index/llms/upstage/BUILD
@@ -0,0 +1 @@
+python_sources()
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/llama_index/llms/upstage/__init__.py b/llama-index-integrations/llms/llama-index-llms-upstage/llama_index/llms/upstage/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..5ff1dc4b442b3336fb35c88595dea43341fb5954
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-upstage/llama_index/llms/upstage/__init__.py
@@ -0,0 +1,4 @@
+from llama_index.llms.upstage.base import Upstage
+
+
+__all__ = ["Upstage"]
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/llama_index/llms/upstage/base.py b/llama-index-integrations/llms/llama-index-llms-upstage/llama_index/llms/upstage/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..0af0b41e46294212a4c04d398eaff133da690d44
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-upstage/llama_index/llms/upstage/base.py
@@ -0,0 +1,166 @@
+from typing import Optional, Dict, Any, Sequence, Callable
+
+import httpx
+from llama_index.core.base.llms.types import LLMMetadata, ChatMessage
+from llama_index.llms.openai import OpenAI
+
+from llama_index.llms.upstage.utils import (
+ resolve_upstage_credentials,
+ is_chat_model,
+ upstage_modelname_to_contextsize,
+ is_function_calling_model,
+)
+from llama_index.core.callbacks import CallbackManager
+from llama_index.core.constants import DEFAULT_TEMPERATURE
+from llama_index.core.types import BaseOutputParser, PydanticProgramMode
+from pydantic import Field, PrivateAttr
+from openai import OpenAI as SyncOpenAI
+from openai import AsyncOpenAI
+
+DEFAULT_UPSTAGE_MODEL = "solar-1-mini-chat"
+
+
+class Upstage(OpenAI):
+ """Upstage LLM.
+
+ Examples:
+ `pip install llama-index-llms-upstage`
+
+ ```python
+ from llama_index.llms.upstage import Upstage
+ import os
+
+ os.environ["UPSTAGE_API_KEY"] = "YOUR_API_KEY"
+
+ llm = Upstage()
+ stream = llm.stream("Hello, how are you?")
+
+ for response in stream:
+ print(response.delta, end="")
+
+ ```
+ """
+
+ model: str = Field(
+ default=DEFAULT_UPSTAGE_MODEL, description="The Upstage model to use."
+ )
+ temperature: float = Field(
+ default=DEFAULT_TEMPERATURE,
+ description="The temperature to use during generation.",
+ gte=0.0,
+ lte=1.0,
+ )
+ max_tokens: Optional[int] = Field(
+ description="The maximum number of tokens to generate."
+ )
+ logprobs: Optional[bool] = Field(
+ description="Whether to return logprobs per token."
+ )
+ top_logprobs: int = Field(
+ description="The number of top token logprobs to return.",
+ default=0,
+ gte=0,
+ lte=20,
+ )
+ additional_kwargs: Dict[str, Any] = Field(
+ description="Additional kwargs for the Upstage API.", default_factory=dict
+ )
+ max_retries: int = Field(
+ description="The maximum number of API retries.", default=3, gte=0
+ )
+ timeout: float = Field(
+ description="The timeout, in seconds, for API requests.", default=60.0, gte=0.0
+ )
+ reuse_client: bool = Field(
+ description=(
+ "Reuse the OpenAI client between requests. When doing anything with large "
+ "volumes of async API calls, setting this to false can improve stability."
+ ),
+ default=True,
+ )
+
+ api_key: str = Field(default=None, description="The Upstage API key.")
+ api_base: str = Field(
+ default="https://api.upstage.ai/v1/solar",
+ description="The Upstage API base URL.",
+ )
+
+ _client: Optional[SyncOpenAI] = PrivateAttr()
+ _aclient: Optional[AsyncOpenAI] = PrivateAttr()
+ _http_client: Optional[httpx.Client] = PrivateAttr()
+
+ def __init__(
+ self,
+ model: str = DEFAULT_UPSTAGE_MODEL,
+ temperature: float = DEFAULT_TEMPERATURE,
+ max_tokens: Optional[int] = None,
+ logprobs: Optional[bool] = None,
+ top_logprobs: int = 0,
+ additional_kwargs: Dict[str, Any] = None,
+ max_retries: int = 3,
+ timeout: float = 60.0,
+ reuse_client: bool = True,
+ api_key: Optional[str] = None,
+ api_base: Optional[str] = None,
+ callback_manager: Optional[CallbackManager] = None,
+ default_headers: Optional[Dict[str, str]] = None,
+ http_client: Optional[httpx.Client] = None, # from base class
+ system_prompt: Optional[str] = None,
+ messages_to_prompt: Optional[Callable[[Sequence[ChatMessage]], str]] = None,
+ completion_to_prompt: Optional[Callable[[str], str]] = None,
+ pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,
+ output_parser: Optional[BaseOutputParser] = None,
+ **kwargs: Any
+ ) -> None:
+ additional_kwargs = additional_kwargs or {}
+ api_key, api_base = resolve_upstage_credentials(
+ api_key=api_key, api_base=api_base
+ )
+
+ super().__init__(
+ model=model,
+ temperature=temperature,
+ max_tokens=max_tokens,
+ logprobs=logprobs,
+ top_logprobs=top_logprobs,
+ additional_kwargs=additional_kwargs,
+ max_retries=max_retries,
+ timeout=timeout,
+ reuse_client=reuse_client,
+ api_key=api_key,
+ api_base=api_base,
+ callback_manager=callback_manager,
+ default_headers=default_headers,
+ http_client=http_client,
+ system_prompt=system_prompt,
+ messages_to_prompt=messages_to_prompt,
+ completion_to_prompt=completion_to_prompt,
+ pydantic_program_mode=pydantic_program_mode,
+ output_parser=output_parser,
+ **kwargs
+ )
+
+ self._client = None
+ self._aclient = None
+ self._http_client = http_client
+
+ def _get_model_name(self) -> str:
+ return self.model
+
+ @classmethod
+ def class_name(cls) -> str:
+ return "upstage_llm"
+
+ @property
+ def metadata(self) -> LLMMetadata:
+ return LLMMetadata(
+ context_window=upstage_modelname_to_contextsize(
+ modelname=self._get_model_name()
+ ),
+ num_output=self.max_tokens or -1,
+ is_chat_model=is_chat_model(model=self._get_model_name()),
+ is_function_calling_model=is_function_calling_model(
+ model=self._get_model_name()
+ ),
+ model_name=self.model,
+ )
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/llama_index/llms/upstage/utils.py b/llama-index-integrations/llms/llama-index-llms-upstage/llama_index/llms/upstage/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..f0acca6d0bb1e2c6ba1d68c5c6a61143c24aa8b0
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-upstage/llama_index/llms/upstage/utils.py
@@ -0,0 +1,50 @@
+import logging
+from typing import Optional, Tuple
+
+from llama_index.core.base.llms.generic_utils import get_from_param_or_env
+
+DEFAULT_UPSTAGE_API_BASE = "https://api.upstage.ai/v1/solar"
+CHAT_MODELS = {
+ "solar-1-mini-chat": 32768,
+}
+
+ALL_AVAILABLE_MODELS = {**CHAT_MODELS}
+
+logger = logging.getLogger(__name__)
+
+
+def resolve_upstage_credentials(
+ api_key: Optional[str] = None,
+ api_base: Optional[str] = None,
+) -> Tuple[Optional[str], str]:
+ """Resolve Upstage credentials.
+
+ The order of precedence is:
+ 1. param
+ 2. env
+ 4. default
+ """
+ # resolve from param or env
+ api_key = get_from_param_or_env("api_key", api_key, "UPSTAGE_API_KEY", "")
+ api_base = get_from_param_or_env("api_base", api_base, "UPSTAGE_API_BASE", "")
+
+ final_api_key = api_key or ""
+ final_api_base = api_base or DEFAULT_UPSTAGE_API_BASE
+
+ return final_api_key, str(final_api_base)
+
+
+def is_chat_model(model: str) -> bool:
+ return model in CHAT_MODELS
+
+
+def is_function_calling_model(model: str) -> bool:
+ return is_chat_model(model)
+
+
+def upstage_modelname_to_contextsize(modelname: str) -> int:
+ if modelname not in ALL_AVAILABLE_MODELS:
+ raise ValueError(
+ f"Unknown model: {modelname}. Please provide a valid Upstage model name in: {', '.join(ALL_AVAILABLE_MODELS.keys())}"
+ )
+ return CHAT_MODELS[modelname]
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-upstage/pyproject.toml
new file mode 100644
index 0000000000000000000000000000000000000000..eb1c8bee6d2e91a4a46242bfb448b695700e4acf
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-upstage/pyproject.toml
@@ -0,0 +1,57 @@
+[build-system]
+build-backend = "poetry.core.masonry.api"
+requires = ["poetry-core"]
+
+[tool.codespell]
+check-filenames = true
+check-hidden = true
+# Feel free to un-skip examples, and experimental, you will just need to
+# work through many typos (--write-changes and --interactive will help)
+skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb"
+
+[tool.llamahub]
+contains_example = false
+import_path = "llama_index.llms.upstage"
+
+[tool.llamahub.class_authors]
+Upstage = "llama-index"
+
+[tool.mypy]
+disallow_untyped_defs = true
+# Remove venv skip when integrated with pre-commit
+exclude = ["_static", "build", "examples", "notebooks", "venv"]
+ignore_missing_imports = true
+python_version = "3.8"
+
+[tool.poetry]
+authors = ["Your Name <you@example.com>"]
+description = "llama-index llms upstage integration"
+license = "MIT"
+name = "llama-index-llms-upstage"
+packages = [{include = "llama_index/"}]
+readme = "README.md"
+version = "0.1.0"
+
+[tool.poetry.dependencies]
+python = ">=3.8.1,<4.0"
+llama-index-core = "^0.10.0"
+openai = "^1.21.2"
+
+[tool.poetry.group.dev.dependencies]
+black = {extras = ["jupyter"], version = "<=23.9.1,>=23.7.0"}
+codespell = {extras = ["toml"], version = ">=v2.2.6"}
+ipython = "8.10.0"
+jupyter = "^1.0.0"
+mypy = "0.991"
+pre-commit = "3.2.0"
+pylint = "2.15.10"
+pytest = "7.2.1"
+pytest-mock = "3.11.1"
+ruff = "0.0.292"
+tree-sitter-languages = "^1.8.0"
+types-Deprecated = ">=0.1.0"
+types-PyYAML = "^6.0.12.12"
+types-protobuf = "^4.24.0.4"
+types-redis = "4.5.5.0"
+types-requests = "2.28.11.8" # TODO: unpin when mypy>0.991
+types-setuptools = "67.1.0.0"
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/tests/BUILD b/llama-index-integrations/llms/llama-index-llms-upstage/tests/BUILD
new file mode 100644
index 0000000000000000000000000000000000000000..dabf212d7e7162849c24a733909ac4f645d75a31
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-upstage/tests/BUILD
@@ -0,0 +1 @@
+python_tests()
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/tests/__init__.py b/llama-index-integrations/llms/llama-index-llms-upstage/tests/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/tests/test_llms_upstage.py b/llama-index-integrations/llms/llama-index-llms-upstage/tests/test_llms_upstage.py
new file mode 100644
index 0000000000000000000000000000000000000000..31ed618400d6bee7b50b9af0f148dab2fdec605d
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-upstage/tests/test_llms_upstage.py
@@ -0,0 +1,7 @@
+from llama_index.core.base.llms.base import BaseLLM
+from llama_index.llms.upstage import Upstage
+
+
+def test_text_inference_embedding_class():
+ names_of_base_classes = [b.__name__ for b in Upstage.__mro__]
+ assert BaseLLM.__name__ in names_of_base_classes
diff --git a/llama-index-integrations/llms/llama-index-llms-upstage/tests/test_upstage.py b/llama-index-integrations/llms/llama-index-llms-upstage/tests/test_upstage.py
new file mode 100644
index 0000000000000000000000000000000000000000..f1cee132ec4f753d311f9642e65ff06b1cb1be68
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-upstage/tests/test_upstage.py
@@ -0,0 +1,175 @@
+import os
+from typing import Any, Generator, Optional
+from unittest.mock import MagicMock, patch
+
+from llama_index.core.base.llms.types import ChatMessage
+
+import openai
+from openai.types.chat.chat_completion import (
+ ChatCompletion,
+ ChatCompletionMessage,
+ Choice,
+)
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk, ChoiceDelta
+from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
+from openai.types.completion import CompletionUsage
+
+from llama_index.llms.upstage import Upstage
+
+
+class CachedUpstageApiKeys:
+ """
+ Saves the users' Upstage API key either in
+ the environment variable or set to the library itself.
+ This allows us to run tests by setting it without plowing over
+ the local environment.
+ """
+
+ def __init__(
+ self,
+ set_env_key_to: Optional[str] = "",
+ set_library_key_to: Optional[str] = None,
+ set_fake_key: bool = False,
+ set_env_type_to: Optional[str] = "",
+ set_library_type_to: str = "upstage", # default value in upstage package
+ ):
+ self.set_env_key_to = set_env_key_to
+ self.set_library_key_to = set_library_key_to
+ self.set_fake_key = set_fake_key
+ self.set_env_type_to = set_env_type_to
+ self.set_library_type_to = set_library_type_to
+
+ def __enter__(self) -> None:
+ self.api_env_variable_was = os.environ.get("UPSTAGE_API_KEY", "")
+ self.upstage_api_key_was = openai.api_key
+
+ os.environ["UPSTAGE_API_KEY"] = str(self.set_env_key_to)
+
+ if self.set_fake_key:
+ os.environ["UPSTAGE_API_KEY"] = "a" * 32
+
+ # No matter what, set the environment variable back to what it was
+ def __exit__(self, *exc: object) -> None:
+ os.environ["UPSTAGE_API_KEY"] = str(self.api_env_variable_was)
+ openai.api_key = self.upstage_api_key_was
+
+
+def mock_chat_completion(*args: Any, **kwargs: Any) -> ChatCompletion:
+ return ChatCompletion(
+ id="chatcmpl-abc123",
+ object="chat.completion",
+ created=1677858242,
+ model="solar-1-mini-chat",
+ usage=CompletionUsage(prompt_tokens=13, completion_tokens=7, total_tokens=20),
+ choices=[
+ Choice(
+ message=ChatCompletionMessage(
+ role="assistant", content="\n\nThis is a test!"
+ ),
+ finish_reason="stop",
+ index=0,
+ )
+ ],
+ )
+
+
+def mock_chat_completion_stream(
+ *args: Any, **kwargs: Any
+) -> Generator[ChatCompletionChunk, None, None]:
+ responses = [
+ ChatCompletionChunk(
+ id="chatcmpl-6ptKyqKOGXZT6iQnqiXAH8adNLUzD",
+ object="chat.completion.chunk",
+ created=1677825464,
+ model="solar-1-mini-chat",
+ choices=[
+ ChunkChoice(
+ delta=ChoiceDelta(role="assistant"), finish_reason=None, index=0
+ )
+ ],
+ ),
+ ChatCompletionChunk(
+ id="chatcmpl-6ptKyqKOGXZT6iQnqiXAH8adNLUzD",
+ object="chat.completion.chunk",
+ created=1677825464,
+ model="solar-1-mini-chat",
+ choices=[
+ ChunkChoice(
+ delta=ChoiceDelta(content="\n\n"), finish_reason=None, index=0
+ )
+ ],
+ ),
+ ChatCompletionChunk(
+ id="chatcmpl-6ptKyqKOGXZT6iQnqiXAH8adNLUzD",
+ object="chat.completion.chunk",
+ created=1677825464,
+ model="solar-1-mini-chat",
+ choices=[
+ ChunkChoice(delta=ChoiceDelta(content="2"), finish_reason=None, index=0)
+ ],
+ ),
+ ChatCompletionChunk(
+ id="chatcmpl-6ptKyqKOGXZT6iQnqiXAH8adNLUzD",
+ object="chat.completion.chunk",
+ created=1677825464,
+ model="solar-1-mini-chat",
+ choices=[ChunkChoice(delta=ChoiceDelta(), finish_reason="stop", index=0)],
+ ),
+ ]
+ yield from responses
+
+
+@patch("llama_index.llms.openai.base.SyncOpenAI")
+def test_chat_model_basic(MockSyncUpstage: MagicMock) -> None:
+ with CachedUpstageApiKeys(set_fake_key=True):
+ mock_instance = MockSyncUpstage.return_value
+ mock_instance.chat.completions.create.return_value = mock_chat_completion()
+
+ llm = Upstage(model="solar-1-mini-chat")
+ prompt = "test prompt"
+ message = ChatMessage(role="user", content="test message")
+
+ response = llm.complete(prompt)
+ assert response.text == "\n\nThis is a test!"
+
+ chat_response = llm.chat([message])
+ assert chat_response.message.content == "\n\nThis is a test!"
+
+
+@patch("llama_index.llms.openai.base.SyncOpenAI")
+def test_chat_model_streaming(MockSyncUpstage: MagicMock) -> None:
+ with CachedUpstageApiKeys(set_fake_key=True):
+ mock_instance = MockSyncUpstage.return_value
+ mock_instance.chat.completions.create.return_value = (
+ mock_chat_completion_stream()
+ )
+
+ llm = Upstage(model="solar-1-mini-chat")
+ prompt = "test prompt"
+ message = ChatMessage(role="user", content="test message")
+
+ response_gen = llm.stream_complete(prompt)
+ responses = list(response_gen)
+ assert responses[-1].text == "\n\n2"
+
+ mock_instance.chat.completions.create.return_value = (
+ mock_chat_completion_stream()
+ )
+ chat_response_gen = llm.stream_chat([message])
+ chat_responses = list(chat_response_gen)
+ assert chat_responses[-1].message.content == "\n\n2"
+ assert chat_responses[-1].message.role == "assistant"
+
+
+def test_validates_api_key_is_present() -> None:
+ with CachedUpstageApiKeys():
+ os.environ["UPSTAGE_API_KEY"] = "a" * 32
+
+ # We can create a new LLM when the env variable is set
+ assert Upstage()
+
+ os.environ["UPSTAGE_API_KEY"] = ""
+
+ # We can create a new LLM when the api_key is set on the
+ # class directly
+ assert Upstage(api_key="a" * 32)