GEN-1292 new protos (#9)

* max_rows in gravity

* Squashed 'protos/' changes from 8c94d7c..69137fc

69137fc Merge pull request #13 from macrocosm-os/gen-1222_add_missing_nebula_to_dataset_query
c024b86 GEN-1222 Add missing nebula to dataset query
ee2edaa Merge pull request #12 from macrocosm-os/optional-fields-2
ada0239 max_rows is not optional - revert this change
d452c9a make some fields in the proto files optional

git-subtree-dir: protos
git-subtree-split: 69137fcf99e56523bb9a6f74e267caf080056e4f

* gen protos

* improve working example

* bump version

* add type

Author: cassova

examples/gravity_workflow_example.py

@@ -20,12 +20,18 @@
 
 class GravityWorkflow:
     def __init__(
-        self, task_name: str, email: str, reddit_subreddit: str, x_hashtag: str
+        self,
+        task_name: str,
+        email: str,
+        reddit_subreddit: str,
+        x_hashtag: str,
+        max_rows: int,
     ):
         self.task_name = task_name
         self.email = email
         self.reddit_subreddit = reddit_subreddit
         self.x_hashtag = x_hashtag
+        self.max_rows = max_rows
         self.api_key = os.environ.get(
             "GRAVITY_API_KEY", os.environ.get("MACROCOSMOS_API_KEY")
         )
@@ -261,7 +267,9 @@ async def build_datasets(self, crawler_ids: Set[str]):
 
                 # Build dataset
                 response = await self.client.gravity.BuildDataset(
-                    crawler_id=crawler_id, notification_requests=[notification]
+                    crawler_id=crawler_id,
+                    max_rows=self.max_rows,
+                    notification_requests=[notification],
                 )
 
                 if response and response.dataset_id:
@@ -477,7 +485,22 @@ def get_user_input():
     if not task_name:
         task_name = "MyTestTask"
 
-    return email, reddit, x_hashtag, task_name
+    # Get max rows with default
+    while True:
+        max_rows_input = input("Max rows per dataset [1000]: ").strip()
+        if not max_rows_input:
+            max_rows = 1000
+            break
+        try:
+            max_rows = int(max_rows_input)
+            if max_rows <= 0:
+                print("Please enter a positive number.")
+                continue
+            break
+        except ValueError:
+            print("Please enter a valid number.")
+
+    return email, reddit, x_hashtag, task_name, max_rows
 
 
 async def main():
@@ -486,12 +509,12 @@ async def main():
     print("════════════════════════════════════════════════════")
 
     # Get user input with defaults
-    email, reddit, x_hashtag, task_name = get_user_input()
+    email, reddit, x_hashtag, task_name, max_rows = get_user_input()
 
     # Set up signal handlers for graceful shutdown
     loop = asyncio.get_running_loop()
 
-    workflow = GravityWorkflow(task_name, email, reddit, x_hashtag)
+    workflow = GravityWorkflow(task_name, email, reddit, x_hashtag, max_rows)
 
     # Register signal handlers
     for sig in (signal.SIGINT, signal.SIGTERM):

protos/README.md

@@ -1,7 +1,9 @@
 # Macrocosmos Protobufs
+
 Protobufs used by Macrocosmos products
 
-# Usage
+## Usage
+
 Here are instructions on how to add and update the protos from your own project repo.
 
 > ⚠️ These commands must be run from the target repo's root directory

protos/apex/v1/apex.proto

@@ -23,25 +23,25 @@ message ChatCompletionRequest {
   // messages: the messages to generate completions for.
   repeated ChatMessage messages = 2;
   // seed: the seed to use for the completion.
-  int64 seed = 3;
+  optional int64 seed = 3;
   // task: the task to generate completions for (e.g. "InferenceTask").
-  string task = 4;
+  optional string task = 4;
   // model: the LLM name to use for the completion. (optional, suggest leaving this empty as not all LLMs are supported)
-  string model = 5;
+  optional string model = 5;
   // test_time_inference: whether to use test time inference.
-  bool test_time_inference = 6;
+  optional bool test_time_inference = 6;
   // mixture: whether to use a mixture of miners to create a slower but better answer.
-  bool mixture = 7;
+  optional bool mixture = 7;
   // sampling_parameters: the sampling parameters to use for the completion.
-  SamplingParameters sampling_parameters = 8;
+  optional SamplingParameters sampling_parameters = 8;
   // inference_mode: the inference mode to use for the completion.
-  string inference_mode = 9;
+  optional string inference_mode = 9;
   // json_format: whether to use JSON format for the completion.
-  bool json_format = 10;
+  optional bool json_format = 10;
   // stream: whether to stream the completion.
-  bool stream = 11;
+  optional bool stream = 11;
   // timeout: the timeout for the completion in seconds.
-  int64 timeout = 12;
+  optional int64 timeout = 12;
 }
 
 // The sampling parameters for the completion.
@@ -52,7 +52,7 @@ message SamplingParameters {
   // top_p: the top_p to use for the completion.
   double top_p = 2;
   // top_k: the top_k to use for the completion.
-  double top_k = 3;
+  optional double top_k = 3;
   // max_new_tokens: the max_new_tokens to use for the completion.
   int64 max_new_tokens = 4;
   // do_sample: whether to do sample for the completion.
@@ -330,13 +330,13 @@ message WebRetrievalRequest {
   // search_query: the search query.
   string search_query = 2;
   // n_miners: the number of miners to use for the query.
-  int64 n_miners = 3;
+  optional int64 n_miners = 3;
   // n_results: the number of results to return.
-  int64 n_results = 4;
+  optional int64 n_results = 4;
   // max_response_time: the max response time to allow for the miners to respond in seconds.
-  int64 max_response_time = 5;
+  optional int64 max_response_time = 5;
   // timeout: the timeout for the web retrieval in seconds.
-  int64 timeout = 6;
+  optional int64 timeout = 6;
 }
 
 // A web search result from Apex

protos/billing/v1/billing.proto

@@ -12,7 +12,7 @@ service BillingService {
 // GetUsageRequest is the request message for getting the usage of the user's credits
 message GetUsageRequest {
   // product_type: the type of the product (i.e. "gravity")
-  string product_type = 1;
+  optional string product_type = 1;
 }
 
 // ProductPlan is details of the subscription plan for a product

protos/gravity/v1/gravity.proto

@@ -115,9 +115,9 @@ message GravityTaskState {
 // GetGravityTasksRequest is the request message for listing gravity tasks for a user
 message GetGravityTasksRequest {
   // gravity_task_id: the ID of the gravity task (optional, if not provided, all gravity tasks for the user will be returned)
-  string gravity_task_id = 1;
+  optional string gravity_task_id = 1;
   // include_crawlers: whether to include the crawler states in the response
-  bool include_crawlers = 2;
+  optional bool include_crawlers = 2;
 }
 
 // GetGravityTasksResponse is the response message for listing gravity tasks for a user
@@ -141,7 +141,7 @@ message NotificationRequest {
   // address: the address to send the notification to (only email addresses are supported currently)
   string address = 2;
   // redirect_url: the URL to include in the notication message that redirects the user to any built datasets
-  string redirect_url = 3;
+  optional string redirect_url = 3;
 }
 
 // GetCrawlerRequest is the request message for getting a crawler
@@ -166,7 +166,7 @@ message CreateGravityTaskRequest {
   //   that is automatically generated upon completion of the crawler is ready to download (optional)
   repeated NotificationRequest notification_requests = 3;
   // gravity_task_id: the ID of the gravity task (optional, default will generate a random ID)
-  string gravity_task_id = 4;
+  optional string gravity_task_id = 4;
 }
 
 // CreateGravityTaskResponse is the response message for creating a new gravity task
@@ -194,6 +194,15 @@ message BuildDatasetResponse {
   Dataset dataset = 2;
 }
 
+message Nebula {
+  // error: nebula build error message
+  string error = 1;
+  // file_size_bytes: the size of the file in bytes
+  int64 file_size_bytes = 2;
+  // url: the URL of the file
+  string url = 3;
+}
+
 // Dataset contains the progress and results of a dataset build
 message Dataset {
   // crawler_workflow_id: the ID of the parent crawler for this dataset
@@ -212,6 +221,8 @@ message Dataset {
   repeated DatasetStep steps = 7;
   // total_steps: the total number of steps in the dataset build
   int64 total_steps = 8;
+  // nebula: the details about the nebula that was built
+  Nebula nebula = 9;
 }
 
 // DatasetFile contains the details about a dataset file

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "macrocosmos"
-version = "1.0.0"
+version = "1.0.1"
 description = "The official Python SDK for Macrocosmos"
 readme = "README.md"
 license = "Apache-2.0"

src/macrocosmos/generated/apex/v1/apex_p2p.py

@@ -30,7 +30,7 @@ class SamplingParameters(BaseModel):
 # top_p: the top_p to use for the completion.
     top_p: float = Field(default=0.0)
 # top_k: the top_k to use for the completion.
-    top_k: float = Field(default=0.0)
+    top_k: typing.Optional[float] = Field(default=0.0)
 # max_new_tokens: the max_new_tokens to use for the completion.
     max_new_tokens: int = Field(default=0)
 # do_sample: whether to do sample for the completion.
@@ -47,25 +47,25 @@ class ChatCompletionRequest(BaseModel):
 # messages: the messages to generate completions for.
     messages: typing.List[ChatMessage] = Field(default_factory=list)
 # seed: the seed to use for the completion.
-    seed: int = Field(default=0)
+    seed: typing.Optional[int] = Field(default=0)
 # task: the task to generate completions for (e.g. "InferenceTask").
-    task: str = Field(default="")
+    task: typing.Optional[str] = Field(default="")
 # model: the LLM name to use for the completion. (optional, suggest leaving this empty as not all LLMs are supported)
-    model: str = Field(default="")
+    model: typing.Optional[str] = Field(default="")
 # test_time_inference: whether to use test time inference.
-    test_time_inference: bool = Field(default=False)
+    test_time_inference: typing.Optional[bool] = Field(default=False)
 # mixture: whether to use a mixture of miners to create a slower but better answer.
-    mixture: bool = Field(default=False)
+    mixture: typing.Optional[bool] = Field(default=False)
 # sampling_parameters: the sampling parameters to use for the completion.
-    sampling_parameters: SamplingParameters = Field(default_factory=SamplingParameters)
+    sampling_parameters: typing.Optional[SamplingParameters] = Field(default_factory=SamplingParameters)
 # inference_mode: the inference mode to use for the completion.
-    inference_mode: str = Field(default="")
+    inference_mode: typing.Optional[str] = Field(default="")
 # json_format: whether to use JSON format for the completion.
-    json_format: bool = Field(default=False)
+    json_format: typing.Optional[bool] = Field(default=False)
 # stream: whether to stream the completion.
-    stream: bool = Field(default=False)
+    stream: typing.Optional[bool] = Field(default=False)
 # timeout: the timeout for the completion in seconds.
-    timeout: int = Field(default=0)
+    timeout: typing.Optional[int] = Field(default=0)
 
 class TopLogprob(BaseModel):
     """
@@ -372,13 +372,13 @@ class WebRetrievalRequest(BaseModel):
 # search_query: the search query.
     search_query: str = Field(default="")
 # n_miners: the number of miners to use for the query.
-    n_miners: int = Field(default=0)
+    n_miners: typing.Optional[int] = Field(default=0)
 # n_results: the number of results to return.
-    n_results: int = Field(default=0)
+    n_results: typing.Optional[int] = Field(default=0)
 # max_response_time: the max response time to allow for the miners to respond in seconds.
-    max_response_time: int = Field(default=0)
+    max_response_time: typing.Optional[int] = Field(default=0)
 # timeout: the timeout for the web retrieval in seconds.
-    timeout: int = Field(default=0)
+    timeout: typing.Optional[int] = Field(default=0)
 
 class WebSearchResult(BaseModel):
     """