Feat/mymagic structured output (#12021)

d14737c1 · singhmanas1 · GitHub · 57edd036 · d14737c1 · d14737c1
Unverified Commit d14737c1 authored 1 year ago by singhmanas1 Committed by GitHub 1 year ago
--- a/docs/examples/llm/mymagic.ipynb
+++ b/docs/examples/llm/mymagic.ipynb
@@ -68,9 +68,18 @@
    "    system_prompt=\"your-system-prompt\",\n",
    "    region=\"your-bucket-region\",\n",
    "    return_output=False,  # Whether you want MyMagic API to return the output json\n",
+    "    input_json_file=None,  # name of the input file (stored on the bucket)\n",
+    "    structured_output=None,  # json schema of the output\n",
    ")"
   ]
  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Note: if return_output is set True above, max_tokens should be set to atleast 100 "
+   ]
+  },
  {
   "cell_type": "code",
   "execution_count": null,
@@ -118,18 +127,9 @@
   "outputs": [],
   "source": [
    "async def main():\n",
-    "    allm = MyMagicAI(\n",
-    "        api_key=\"your-api-key\",\n",
-    "        storage_provider=\"s3\",  # s3, gcs\n",
-    "        bucket_name=\"your-bucket-name\",\n",
-    "        session=\"your-session-name\",  # files should be located in this folder on which batch inference will be run\n",
-    "        role_arn=\"your-role-arn\",\n",
-    "        system_prompt=\"your-system-prompt\",\n",
-    "        region=\"your-bucket-region\",\n",
-    "    )\n",
-    "    response = await allm.acomplete(\n",
+    "    response = await llm.acomplete(\n",
    "        question=\"your-question\",\n",
-    "        model=\"chhoose-model\",  # currently we support mistral7b, llama7b, mixtral8x7b,codellama70b, llama70b, more to come...\n",
+    "        model=\"chhoose-model\",  # currently we support mistral7b, llama7b, mixtral8x7,codellama70b, llama70b, more to come...\n",
    "        max_tokens=5,  # number of tokens to generate, default is 10\n",
    "    )\n",
    "\n",

 %% Cell type:markdown id: tags:

 # MyMagic AI LLM

 %% Cell type:markdown id: tags:

 ## Introduction
 This notebook demonstrates how to use MyMagicAI for batch inference on massive data stored in cloud buckets. The only enpoints implemented are `complete` and `acomplete` which can work on many use cases including Completion, Summariation and Extraction.
 To use this notebook, you need an API key (Personal Access Token) from MyMagicAI and data stored in cloud buckets.
 Sign up by clicking Get Started at [MyMagicAI's website](https://mymagic.ai/) to get your API key.

 %% Cell type:markdown id: tags:

 ## Setup
 To set up your bucket and grant MyMagic API a secure access to your cloud storage, please visit [MyMagic docs](https://docs.mymagic.ai/) for reference.
 If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.

 %% Cell type:code id: tags:

 ``` python
 %pip install llama-index-llms-mymagic
 ```

 %% Cell type:code id: tags:

 ``` python
 !pip install llama-index
 ```

 %% Cell type:code id: tags:

 ``` python
 from llama_index.llms.mymagic import MyMagicAI
 ```

 %% Cell type:code id: tags:

 ``` python
 llm = MyMagicAI(
    api_key="your-api-key",
    storage_provider="s3",  # s3, gcs
    bucket_name="your-bucket-name",
    session="your-session-name",  # files should be located in this folder on which batch inference will be run
    role_arn="your-role-arn",
    system_prompt="your-system-prompt",
    region="your-bucket-region",
    return_output=False,  # Whether you want MyMagic API to return the output json
+    input_json_file=None,  # name of the input file (stored on the bucket)
+    structured_output=None,  # json schema of the output
 )
 ```

+%% Cell type:markdown id: tags:
+
+Note: if return_output is set True above, max_tokens should be set to atleast 100
+
 %% Cell type:code id: tags:

 ``` python
 resp = llm.complete(
    question="your-question",
    model="chhoose-model",  # currently we support mistral7b, llama7b, mixtral8x7b, codellama70b, llama70b, more to come...
    max_tokens=5,  # number of tokens to generate, default is 10
 )
 ```

 %% Cell type:code id: tags:

 ``` python
 # The response indicated that the final output is stored in your bucket or raises an exception if the job failed
 print(resp)
 ```

 %% Cell type:markdown id: tags:

 ## Asynchronous Requests by using `acomplete` endpoint
 For asynchronous operations, use the following approach.

 %% Cell type:code id: tags:

 ``` python
 import asyncio
 ```

 %% Cell type:code id: tags:

 ``` python
 async def main():
-    allm = MyMagicAI(
-        api_key="your-api-key",
-        storage_provider="s3",  # s3, gcs
-        bucket_name="your-bucket-name",
-        session="your-session-name",  # files should be located in this folder on which batch inference will be run
-        role_arn="your-role-arn",
-        system_prompt="your-system-prompt",
-        region="your-bucket-region",
-    )
-    response = await allm.acomplete(
+    response = await llm.acomplete(
        question="your-question",
-        model="chhoose-model",  # currently we support mistral7b, llama7b, mixtral8x7b,codellama70b, llama70b, more to come...
+        model="chhoose-model",  # currently we support mistral7b, llama7b, mixtral8x7,codellama70b, llama70b, more to come...
        max_tokens=5,  # number of tokens to generate, default is 10
    )

    print("Async completion response:", response)
 ```

 %% Cell type:code id: tags:

 ``` python
 await main()
 ```

--- a/llama-index-integrations/llms/llama-index-llms-mymagic/llama_index/llms/mymagic/base.py
+++ b/llama-index-integrations/llms/llama-index-llms-mymagic/llama_index/llms/mymagic/base.py
@@ -51,6 +51,11 @@ class MyMagicAI(LLM):
    return_output: Optional[bool] = Field(
        False, description="Whether MyMagic API should return the output json"
    )
+    input_json_file: Optional[str] = None
+
+    structured_output: Optional[Dict[str, Any]] = Field(
+        None, description="User-defined structure for the response output"
+    )

    def __init__(
        self,
@@ -62,9 +67,12 @@ class MyMagicAI(LLM):
        role_arn: Optional[str] = None,
        region: Optional[str] = None,
        return_output: Optional[bool] = False,
+        input_json_file: Optional[str] = None,
+        structured_output: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
+        self.return_output = return_output

        self.question_data = {
            "storage_provider": storage_provider,
@@ -76,6 +84,8 @@ class MyMagicAI(LLM):
            "system_prompt": system_prompt,
            "region": region,
            "return_output": return_output,
+            "input_json_file": input_json_file,
+            "structured_output": structured_output,
        }

    @classmethod
@@ -87,7 +97,9 @@ class MyMagicAI(LLM):
        return self.base_url_template.format(model=model)

    async def _submit_question(self, question_data: Dict[str, Any]) -> Dict[str, Any]:
-        async with httpx.AsyncClient() as client:
+        timeout_config = httpx.Timeout(600.0, connect=60.0)
+
+        async with httpx.AsyncClient(timeout=timeout_config) as client:
            url = f"{self._construct_url(self.model)}/submit_question"
            resp = await client.post(url, json=question_data)
            resp.raise_for_status()
@@ -129,6 +141,10 @@ class MyMagicAI(LLM):
        )

        task_response = await self._submit_question(self.question_data)
+
+        if self.return_output:
+            return task_response
+
        task_id = task_response.get("task_id")
        while True:
            result = await self._get_result(task_id)
@@ -150,6 +166,9 @@ class MyMagicAI(LLM):
        )

        task_response = self._submit_question_sync(self.question_data)
+        if self.return_output:
+            return task_response
+
        task_id = task_response.get("task_id")
        while True:
            result = self._get_result_sync(task_id)

--- a/llama-index-integrations/llms/llama-index-llms-mymagic/pyproject.toml
+++ b/llama-index-integrations/llms/llama-index-llms-mymagic/pyproject.toml
@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-llms-mymagic"
 readme = "README.md"
-version = "0.1.5"
+version = "0.1.6"

 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"