From 205d063ec6c17a1b205c18e91ecbc895ab7ecf04 Mon Sep 17 00:00:00 2001
From: hallacy <hallacy@openai.com>
Date: Thu, 9 Sep 2021 13:46:36 -0400
Subject: [PATCH] Search v2 functionality (#33)

* Add support for search_indices (#64)

* Add support for search_indices

* Updated with Schnurr's comments

* Add version to search (#65)

* Make search query required (#67)
---
 openai/__init__.py               |  4 +--
 openai/api_resources/__init__.py |  7 +++--
 openai/api_resources/search.py   | 15 +++++++++
 openai/cli.py                    | 52 +++++++++++++++++++++++++++++---
 openai/version.py                |  2 +-
 5 files changed, 69 insertions(+), 11 deletions(-)
 create mode 100644 openai/api_resources/search.py

diff --git a/openai/__init__.py b/openai/__init__.py
index f9a6bad36e..054cec3d4d 100644
--- a/openai/__init__.py
+++ b/openai/__init__.py
@@ -34,6 +34,6 @@
     File,
     FineTune,
     Model,
+    Search,
 )
-
-from openai.error import OpenAIError, APIError, InvalidRequestError  # noqa: E402,F401
+from openai.error import APIError, InvalidRequestError, OpenAIError  # noqa: E402,F401
diff --git a/openai/api_resources/__init__.py b/openai/api_resources/__init__.py
index a6da13b2df..6d2cbe67cc 100644
--- a/openai/api_resources/__init__.py
+++ b/openai/api_resources/__init__.py
@@ -1,8 +1,9 @@
+from openai.api_resources.answer import Answer  # noqa: F401
+from openai.api_resources.classification import Classification  # noqa: F401
 from openai.api_resources.completion import Completion  # noqa: F401
 from openai.api_resources.engine import Engine  # noqa: F401
 from openai.api_resources.error_object import ErrorObject  # noqa: F401
 from openai.api_resources.file import File  # noqa: F401
-from openai.api_resources.answer import Answer  # noqa: F401
-from openai.api_resources.classification import Classification  # noqa: F401
-from openai.api_resources.model import Model  # noqa: F401
 from openai.api_resources.fine_tune import FineTune  # noqa: F401
+from openai.api_resources.model import Model  # noqa: F401
+from openai.api_resources.search import Search  # noqa: F401
diff --git a/openai/api_resources/search.py b/openai/api_resources/search.py
new file mode 100644
index 0000000000..4a6e9a4b46
--- /dev/null
+++ b/openai/api_resources/search.py
@@ -0,0 +1,15 @@
+from openai.api_resources.abstract.api_resource import APIResource
+
+
+class Search(APIResource):
+    api_prefix = "v1"
+    OBJECT_NAME = "search_indices"
+
+    @classmethod
+    def class_url(cls):
+        return "/%s/%s" % (cls.api_prefix, cls.OBJECT_NAME)
+
+    @classmethod
+    def create_alpha(cls, **params):
+        instance = cls()
+        return instance.request("post", f"{cls.class_url()}/search", params)
diff --git a/openai/cli.py b/openai/cli.py
index 6d67e2b0ef..e625d1bdf6 100644
--- a/openai/cli.py
+++ b/openai/cli.py
@@ -3,16 +3,16 @@
 import signal
 import sys
 import warnings
+
+import openai
 from openai.validators import (
-    write_out_file,
     apply_necessary_remediation,
     apply_optional_remediation,
-    read_any_format,
     get_validators,
+    read_any_format,
+    write_out_file,
 )
 
-import openai
-
 
 class bcolors:
     HEADER = "\033[95m"
@@ -102,7 +102,6 @@ def generate(cls, args):
 
     @classmethod
     def search(cls, args):
-        # Will soon be deprecated and replaced by a Search.create
         params = {
             "query": args.query,
             "max_rerank": args.max_rerank,
@@ -113,6 +112,9 @@ def search(cls, args):
         if args.file:
             params["file"] = args.file
 
+        if args.version:
+            params["version"] = args.version
+
         resp = openai.Engine(id=args.id).search(**params)
         scores = [
             (search_result["score"], search_result["document"])
@@ -221,6 +223,17 @@ def list(cls, args):
         print(file)
 
 
+class Search:
+    @classmethod
+    def create_alpha(cls, args):
+        resp = openai.Search.create_alpha(
+            query=[args.query],
+            max_documents=args.max_documents,
+            file_id=args.file,
+        )
+        print(resp)
+
+
 class FineTune:
     @classmethod
     def list(cls, args):
@@ -589,6 +602,11 @@ def help(args):
         type=bool,
         default=False,
     )
+    sub.add_argument(
+        "--version",
+        help="The version of the search routing to use",
+    )
+
     sub.add_argument("-q", "--query", required=True, help="Search query")
     sub.set_defaults(func=Engine.search)
 
@@ -688,6 +706,30 @@ def help(args):
     sub = subparsers.add_parser("files.list")
     sub.set_defaults(func=File.list)
 
+    # Search
+    sub = subparsers.add_parser("search.create_alpha")
+
+    sub.add_argument(
+        "-f",
+        "--file",
+        required=True,
+        help="ID for previously uploaded file that contains the documents you want to search",
+    )
+    sub.add_argument(
+        "-m",
+        "--max_documents",
+        help="The maximum number of documents to return",
+        type=int,
+        default=200,
+    )
+    sub.add_argument(
+        "-q",
+        "--query",
+        required=True,
+        help="Search query",
+    )
+    sub.set_defaults(func=Search.create_alpha)
+
     # Finetune
     sub = subparsers.add_parser("fine_tunes.list")
     sub.set_defaults(func=FineTune.list)
diff --git a/openai/version.py b/openai/version.py
index d3e0bd4aa1..cae39d7ebb 100644
--- a/openai/version.py
+++ b/openai/version.py
@@ -1 +1 @@
-VERSION = "0.10.3"
+VERSION = "0.10.4"