save

lilac/garden_client.py

@@ -7,13 +7,16 @@
 
 import numpy as np
 import requests
+from pydantic import BaseModel
 
 from .env import env
 from .schema import Item
 from .utils import DebugTimer
 
 GARDEN_FRONT_GATE_URL = 'https://lilacai--front-gate-fastapi-app.modal.run'
 GARDEN_ENCODING_SCHEME_HEADER = 'X-Lilac-EncodingScheme'
+GARDEN_DOC_COUNT_HEADER = 'X-Lilac-DocCount'
+PLATFORM_LIST_ENDPOINTS = 'https://platform-v3omj4i5vq-uc.a.run.app/api/v1/entitlements/list'
 
 
 def _decode_b64_npy(b: bytes) -> np.ndarray:
@@ -33,7 +36,7 @@ def _call_garden(endpoint_name: str, docs: list[Any], **kwargs: Any) -> Iterator
       params={k: str(v) for k, v in kwargs.items()},
       headers={
         'Authorization': f'Bearer {lilac_api_key}',
-        'X-Lilac-DocCount': str(len(docs)),
+        GARDEN_DOC_COUNT_HEADER: str(len(docs)),
       },
       stream=True,
     ) as response:
@@ -44,4 +47,17 @@ def _call_garden(endpoint_name: str, docs: list[Any], **kwargs: Any) -> Iterator
         yield decoder(line)
 
 
-cluster = functools.partial(_call_garden, 'cluster')
+cluster = functools.partial(_call_garden, 'clustering')
+pii = functools.partial(_call_garden, 'pii')
+
+
+class GardenEndpoint(BaseModel):
+  entitlement_name: str
+  limitations: dict[str, Any]
+
+
+def list_endpoints() -> list[GardenEndpoint]:
+  """List the available endpoints."""
+  # TODO: put api key in Auth header instead of json body.
+  lilac_api_key = env('LILAC_API_KEY')
+  return requests.post(PLATFORM_LIST_ENDPOINTS, json={'api_key': lilac_api_key}).json()

lilac/router_garden.py

@@ -0,0 +1,14 @@
+"""Endpoints for Lilac Garden."""
+
+from fastapi import APIRouter
+
+from .garden_client import GardenEndpoint, list_endpoints
+from .router_utils import RouteErrorHandler
+
+router = APIRouter(route_class=RouteErrorHandler)
+
+
+@router.get('/')
+def get_available_endpoints() -> list[GardenEndpoint]:
+  """Get the list of available sources."""
+  return list_endpoints()

lilac/server.py

@@ -32,6 +32,7 @@
   router_data_loader,
   router_dataset,
   router_dataset_signals,
+  router_garden,
   router_google_login,
   router_rag,
   router_signal,
@@ -141,6 +142,7 @@ def module_not_found_error(request: Request, exc: ModuleNotFoundError) -> JSONRe
 v1_router.include_router(router_signal.router, prefix='/signals', tags=['signals'])
 v1_router.include_router(router_tasks.router, prefix='/tasks', tags=['tasks'])
 v1_router.include_router(router_rag.router, prefix='/rag', tags=['rag'])
+v1_router.include_router(router_garden.router, prefix='/garden', tags=['garden'])
 
 for source_name, source in registered_sources().items():
   if source.router:

lilac/signals/pii.py

@@ -1,14 +1,13 @@
 """Compute text statistics for a document."""
 from typing import ClassVar, Iterator, Optional
 
-import modal
 from typing_extensions import override
 
-from ..batch_utils import compress_docs
+from .. import garden_client
 from ..schema import Field, Item, RichData, SignalInputType, field
 from ..signal import TextSignal
 from ..tasks import TaskExecutionType
-from ..utils import DebugTimer, chunks
+from ..utils import DebugTimer
 
 SECRETS_KEY = 'secrets'
 # Selected categories. For all categories, see:
@@ -72,10 +71,5 @@ def compute(self, data: list[RichData]) -> list[Optional[Item]]:
 
   @override
   def compute_garden(self, docs: Iterator[str]) -> Iterator[Item]:
-    pii = modal.Function.lookup('pii', 'PII.detect')
     with DebugTimer('Computing PII on Lilac Garden'):
-      batches = chunks(docs, PII_REMOTE_BATCH_SIZE)
-      requests = ({'gzipped_docs': compress_docs(b)} for b in batches)
-      for response in pii.map(requests, order_outputs=True):
-        for item in response['result']:
-          yield item
+      yield from garden_client.pii(list(docs))

web/blueprint/src/lib/components/ComputeClusterModal.svelte

@@ -96,10 +96,10 @@
       options.namespace,
       options.datasetName,
       {
-        input: selectedFormatSelector == null ? options.input : null,
+        input: selectedFormatSelector === 'none' ? options.input : null,
         use_garden: options.use_garden,
         output_path: outputColumn,
-        input_selector: selectedFormatSelector,
+        input_selector: selectedFormatSelector === 'none' ? null : selectedFormatSelector,
         overwrite: options.overwrite
       }
     ]);

web/lib/fastapi_client/index.ts

@@ -47,6 +47,7 @@ export type { ExampleIn } from './models/ExampleIn';
 export type { ExampleOrigin } from './models/ExampleOrigin';
 export type { ExportOptions } from './models/ExportOptions';
 export type { Field } from './models/Field';
+export type { GardenEndpoint } from './models/GardenEndpoint';
 export type { GetStatsOptions } from './models/GetStatsOptions';
 export type { GroupsSortBy } from './models/GroupsSortBy';
 export type { HTTPValidationError } from './models/HTTPValidationError';
@@ -113,6 +114,7 @@ export { ConceptsService } from './services/ConceptsService';
 export { DataLoadersService } from './services/DataLoadersService';
 export { DatasetsService } from './services/DatasetsService';
 export { DefaultService } from './services/DefaultService';
+export { GardenService } from './services/GardenService';
 export { GoogleLoginService } from './services/GoogleLoginService';
 export { LangsmithService } from './services/LangsmithService';
 export { RagService } from './services/RagService';

web/lib/fastapi_client/models/GardenEndpoint.ts

@@ -0,0 +1,10 @@
+/* generated using openapi-typescript-codegen -- do no edit */
+/* istanbul ignore file */
+/* tslint:disable */
+/* eslint-disable */
+
+export type GardenEndpoint = {
+    entitlement_name: string;
+    limitations: Record<string, any>;
+};
+

web/lib/fastapi_client/services/GardenService.ts

@@ -0,0 +1,26 @@
+/* generated using openapi-typescript-codegen -- do no edit */
+/* istanbul ignore file */
+/* tslint:disable */
+/* eslint-disable */
+import type { GardenEndpoint } from '../models/GardenEndpoint';
+
+import type { CancelablePromise } from '../core/CancelablePromise';
+import { OpenAPI } from '../core/OpenAPI';
+import { request as __request } from '../core/request';
+
+export class GardenService {
+
+    /**
+     * Get Available Endpoints
+     * Get the list of available sources.
+     * @returns GardenEndpoint Successful Response
+     * @throws ApiError
+     */
+    public static getAvailableEndpoints(): CancelablePromise<Array<GardenEndpoint>> {
+        return __request(OpenAPI, {
+            method: 'GET',
+            url: '/api/v1/garden/',
+        });
+    }
+
+}