[OPIK-355] [Prompt library] SDK Prompt implementation - link with exp…

…eriment (#590) * update openapi spec * add id field to prompt * [OPIK-321] Add prompt id to experiments * rename field to version_id * add prompt field to experiment * link prompt to experiment * fix linter warnings * Rename experiment prompt link dto * update openapi spec * sync API-call params with recent changes in openapi * move prompt from experiment config to method params --------- Co-authored-by: Thiago Hora <[email protected]>
comet-ml · Nov 8, 2024 · f32c146 · f32c146
1 parent a8914c6
commit f32c146
Show file tree

Hide file tree

Showing 22 changed files with 716 additions and 16 deletions.
diff --git a/sdks/python/code_generation/fern/openapi/openapi.yaml b/sdks/python/code_generation/fern/openapi/openapi.yaml
@@ -52,6 +52,48 @@ tags:
 - name: Traces
   description: Trace related resources
 paths:
+  /v1/internal/usage/bi-datasets:
+    get:
+      tags:
+      - System usage
+      summary: Get datasets information for BI events
+      description: Get datasets information for BI events per user per workspace
+      operationId: getDatasetBiInfo
+      responses:
+        "200":
+          description: Datasets BiInformationResponse resource
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/BiInformationResponse'
+  /v1/internal/usage/bi-experiments:
+    get:
+      tags:
+      - System usage
+      summary: Get experiments information for BI events
+      description: Get experiments information for BI events per user per workspace
+      operationId: getExperimentBiInfo
+      responses:
+        "200":
+          description: Experiments BiInformationResponse resource
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/BiInformationResponse'
+  /v1/internal/usage/bi-traces:
+    get:
+      tags:
+      - System usage
+      summary: Get traces information for BI events
+      description: Get traces information for BI events per user per workspace
+      operationId: getTracesBiInfo
+      responses:
+        "200":
+          description: Traces BiInformationResponse resource
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/BiInformationResponse'
   /v1/internal/usage/workspace-trace-counts:
     get:
       tags:
@@ -1187,6 +1229,10 @@ paths:
         in: query
         schema:
           type: string
+      - name: truncate
+        in: query
+        schema:
+          type: boolean
       responses:
         "200":
           description: Spans resource
@@ -1391,6 +1437,10 @@ paths:
         in: query
         schema:
           type: string
+      - name: truncate
+        in: query
+        schema:
+          type: boolean
       responses:
         "200":
           description: Trace resource
@@ -1562,6 +1612,23 @@ paths:
             application/json: {}
 components:
   schemas:
+    BiInformation:
+      type: object
+      properties:
+        workspace_id:
+          type: string
+        user:
+          type: string
+        count:
+          type: integer
+          format: int64
+    BiInformationResponse:
+      type: object
+      properties:
+        bi_information:
+          type: array
+          items:
+            $ref: '#/components/schemas/BiInformation'
     TraceCountResponse:
       type: object
       properties:
@@ -2294,6 +2361,8 @@ components:
         last_updated_by:
           type: string
           readOnly: true
+        prompt_version:
+          $ref: '#/components/schemas/PromptVersionLink'
     FeedbackScoreAverage:
       required:
       - name
@@ -2305,6 +2374,21 @@ components:
         value:
           type: number
       readOnly: true
+    PromptVersionLink:
+      required:
+      - id
+      type: object
+      properties:
+        id:
+          type: string
+          format: uuid
+        commit:
+          type: string
+          readOnly: true
+        prompt_id:
+          type: string
+          format: uuid
+          readOnly: true
     Experiment_Write:
       required:
       - dataset_name
@@ -2319,6 +2403,16 @@ components:
           type: string
         metadata:
           $ref: '#/components/schemas/JsonNode_Write'
+        prompt_version:
+          $ref: '#/components/schemas/PromptVersionLink_Write'
+    PromptVersionLink_Write:
+      required:
+      - id
+      type: object
+      properties:
+        id:
+          type: string
+          format: uuid
     ExperimentItemsBatch:
       required:
       - experiment_items
@@ -2414,6 +2508,8 @@ components:
         last_updated_by:
           type: string
           readOnly: true
+        prompt_version:
+          $ref: '#/components/schemas/PromptVersionLink_Public'
     FeedbackScoreAverage_Public:
       required:
       - name
@@ -2425,6 +2521,21 @@ components:
         value:
           type: number
       readOnly: true
+    PromptVersionLink_Public:
+      required:
+      - id
+      type: object
+      properties:
+        id:
+          type: string
+          format: uuid
+        commit:
+          type: string
+          readOnly: true
+        prompt_id:
+          type: string
+          format: uuid
+          readOnly: true
     ErrorMessage_Public:
       type: object
       properties:

diff --git a/sdks/python/src/opik/api_objects/experiment/experiment.py b/sdks/python/src/opik/api_objects/experiment/experiment.py
@@ -5,7 +5,7 @@
 from opik.rest_api.types import experiment_item as rest_experiment_item
 from . import experiment_item
 from .. import helpers, constants
-
+from ... import Prompt
 
 LOGGER = logging.getLogger(__name__)
 
@@ -17,11 +17,13 @@ def __init__(
         name: Optional[str],
         dataset_name: str,
         rest_client: rest_api_client.OpikApi,
+        prompt: Optional[Prompt] = None,
     ) -> None:
         self.id = id
         self.name = name
         self.dataset_name = dataset_name
         self._rest_client = rest_client
+        self.prompt = prompt
 
     def insert(self, experiment_items: List[experiment_item.ExperimentItem]) -> None:
         rest_experiment_items = [

diff --git a/sdks/python/src/opik/api_objects/opik_client.py b/sdks/python/src/opik/api_objects/opik_client.py
@@ -431,43 +431,53 @@ def create_experiment(
         dataset_name: str,
         name: Optional[str] = None,
         experiment_config: Optional[Dict[str, Any]] = None,
+        prompt: Optional[Prompt] = None,
     ) -> experiment.Experiment:
         """
         Creates a new experiment using the given dataset name and optional parameters.
 
         Args:
-            dataset_name (str): The name of the dataset to associate with the experiment.
-            name (Optional[str]): The optional name for the experiment. If None, a generated name will be used.
-            experiment_config (Optional[Dict[str, Any]]): Optional experiment configuration parameters. Must be a dictionary if provided.
+            dataset_name: The name of the dataset to associate with the experiment.
+            name: The optional name for the experiment. If None, a generated name will be used.
+            experiment_config: Optional experiment configuration parameters. Must be a dictionary if provided.
+            prompt: Prompt object to associate with the experiment.
 
         Returns:
             experiment.Experiment: The newly created experiment object.
         """
         id = helpers.generate_id()
+        metadata = None
+        prompt_version: Optional[Dict[str, str]] = None
 
         if isinstance(experiment_config, Mapping):
+            if prompt is not None:
+                prompt_version = {"id": prompt.__internal_api__version_id__}
+
+                if "prompt" not in experiment_config:
+                    experiment_config["prompt"] = prompt.prompt
+
             metadata = jsonable_encoder.jsonable_encoder(experiment_config)
+
         elif experiment_config is not None:
             LOGGER.error(
                 "Experiment config must be dictionary, but %s was provided. Config will not be logged.",
                 experiment_config,
             )
-            metadata = None
-        else:
-            metadata = None
 
         self._rest_client.experiments.create_experiment(
             name=name,
             dataset_name=dataset_name,
             id=id,
             metadata=metadata,
+            prompt_version=prompt_version,
         )
 
         experiment_ = experiment.Experiment(
             id=id,
             name=name,
             dataset_name=dataset_name,
             rest_client=self._rest_client,
+            prompt=prompt,
         )
 
         return experiment_

diff --git a/sdks/python/src/opik/api_objects/prompt/prompt.py b/sdks/python/src/opik/api_objects/prompt/prompt.py
@@ -33,6 +33,9 @@ def __init__(
         self._name = new_instance.name
         self._prompt = new_instance.prompt
         self._commit = new_instance.commit
+        self.__internal_api__version_id__: str = (
+            new_instance.__internal_api__version_id__
+        )
         self.__internal_api__prompt_id__: str = new_instance.__internal_api__prompt_id__
 
     @property
@@ -75,7 +78,8 @@ def from_fern_prompt_version(
         # will not call __init__ to avoid API calls, create new instance with __new__
         prompt = cls.__new__(cls)
 
-        prompt.__internal_api__prompt_id__ = prompt_version.id
+        prompt.__internal_api__version_id__ = prompt_version.id
+        prompt.__internal_api__prompt_id__ = prompt_version.prompt_id
         prompt._name = name
         prompt._prompt = prompt_version.template
         prompt._commit = prompt_version.commit

diff --git a/sdks/python/src/opik/evaluation/evaluator.py b/sdks/python/src/opik/evaluation/evaluator.py
@@ -3,6 +3,7 @@
 
 from .types import LLMTask
 from .metrics import base_metric
+from .. import Prompt
 from ..api_objects.dataset import dataset
 from ..api_objects.experiment import experiment_item
 from ..api_objects import opik_client
@@ -20,6 +21,7 @@ def evaluate(
     verbose: int = 1,
     nb_samples: Optional[int] = None,
     task_threads: int = 16,
+    prompt: Optional[Prompt] = None,
 ) -> evaluation_result.EvaluationResult:
     """
     Performs task evaluation on a given dataset.
@@ -52,6 +54,8 @@ def evaluate(
             threads are created, all tasks executed in the current thread sequentially.
             are executed sequentially in the current thread.
             Use more than 1 worker if your task object is compatible with sharing across threads.
+
+        prompt: Prompt object to link with experiment.
     """
     client = opik_client.get_client_cached()
     start_time = time.time()
@@ -78,6 +82,7 @@ def evaluate(
         name=experiment_name,
         dataset_name=dataset.name,
         experiment_config=experiment_config,
+        prompt=prompt,
     )
 
     report.display_experiment_link(dataset.name, experiment.id)

diff --git a/sdks/python/src/opik/rest_api/__init__.py b/sdks/python/src/opik/rest_api/__init__.py
@@ -1,6 +1,8 @@
 # This file was auto-generated by Fern from our API Definition.
 
 from .types import (
+    BiInformation,
+    BiInformationResponse,
     CategoricalFeedbackDefinition,
     CategoricalFeedbackDefinitionCreate,
     CategoricalFeedbackDefinitionPublic,
@@ -87,6 +89,9 @@
     PromptPublic,
     PromptVersion,
     PromptVersionDetail,
+    PromptVersionLink,
+    PromptVersionLinkPublic,
+    PromptVersionLinkWrite,
     PromptVersionPagePublic,
     PromptVersionPublic,
     Span,
@@ -128,6 +133,8 @@
 
 __all__ = [
     "BadRequestError",
+    "BiInformation",
+    "BiInformationResponse",
     "CategoricalFeedbackDefinition",
     "CategoricalFeedbackDefinitionCreate",
     "CategoricalFeedbackDefinitionPublic",
@@ -220,6 +227,9 @@
     "PromptPublic",
     "PromptVersion",
     "PromptVersionDetail",
+    "PromptVersionLink",
+    "PromptVersionLinkPublic",
+    "PromptVersionLinkWrite",
     "PromptVersionPagePublic",
     "PromptVersionPublic",
     "Span",