better docstrings

segment_search_queries.py.py

@@ -1,3 +1,21 @@
+"""
+This script is used to segment a request into multiple search queries and perform them asynchronously.
+The `Search` class represents a single search query and has the `execute` method to perform the search.
+The `MultiSearch` class represents multiple searches and has an `execute` method that runs all the 
+searches concurrently using asyncio.
+The `segment` function uses OpenAI's GPT-3 model to convert a given string into multiple search queries,
+which are then run by calling the `execute` method of the returned `MultiSearch` object.
+
+Examples:
+>>> queries = segment(
+...     "Please send me the video from last week about the investment case study and also documents about your GPDR policy?"
+... )
+>>> queries.execute()
+# Expected output:
+# >>> Searching for `Video` with query `investment case study` using `SearchType.VIDEO`
+# >>> Searching for `Documents` with query `GPDR policy` using `SearchType.EMAIL`
+"""
+
 from openai_function_call import OpenAISchema
 from pydantic import Field
 from typing import List
@@ -7,17 +25,20 @@ import enum
 
 
 class SearchType(str, enum.Enum):
+    """Enumeration representing the types of searches that can be performed."""
+
     VIDEO = "video"
     EMAIL = "email"
 
 
 class Search(OpenAISchema):
     """
-    Search query for a single request
+    Class representing a single search query.
 
-    Tips:
-    - Be specific with your query, use key words and multiple representations of the same thing, e.g. "video" and "video clip" or "SSO" and "single sign on"
-    - Use the title to describe the request, e.g. "Video from last week about the investment case study"
+    Args:
+        title (str): The title of the request.
+        query (str): The query string to search for.
+        type (SearchType): The type of search to perform.
     """
 
     title: str = Field(..., description="Title of the request")
@@ -35,10 +56,10 @@ class Search(OpenAISchema):
 
 class MultiSearch(OpenAISchema):
     """
-    Segment a request into multiple search queries
+    Class representing multiple search queries.
 
-    Tips:
-    - Do not overlap queries, e.g. "video" and "video clip" are too similar
+    Args:
+        searches (List[Search]): The list of searches to perform.
     """
 
     searches: List[Search] = Field(..., description="List of searches")
@@ -54,6 +75,16 @@ class MultiSearch(OpenAISchema):
 
 @retry(stop=stop_after_attempt(3))
 def segment(data: str) -> MultiSearch:
+    """
+    Convert a string into multiple search queries using OpenAI's GPT-3 model.
+
+    Args:
+        data (str): The string to convert into search queries.
+
+    Returns:
+        MultiSearch: An object representing the multiple search queries.
+    """
+
     completion = openai.ChatCompletion.create(
         model="gpt-3.5-turbo-0613",
         temperature=0,