Merge pull request #9 from theam/run-evaluations-from-ui

Run evaluations from UI and add Auth

Author: alvaroloes

apps/aifindr-evaluations-runner/evaluator.py

@@ -1,12 +1,14 @@
-from opik import Opik
+import logging
+from opik.api_objects.opik_client import get_client_cached
+from opik.config import update_session_config
 from opik.evaluation import evaluate
 from opik.evaluation.metrics import (Hallucination, ContextRecall, ContextPrecision)
 from workflows import run_workflow
 from metrics.follows_criteria import FollowsCriteria
 from pydantic import BaseModel
 from enum import Enum
 
-client = Opik()
+logger = logging.getLogger(__name__)
 
 class ExperimentStatus(Enum):
     RUNNING = "running"
@@ -16,19 +18,30 @@ class ExperimentStatus(Enum):
 
 class EvaluationParams(BaseModel):
     task_id: str
+    workspace_name: str
     dataset_name: str
     experiment_name: str
     project_name: str
     base_prompt_name: str
     workflow: str
+    api_key: str
 
 def evaluation_task(dataset_item, workflow: str):
+    # TODO: validate properly dataset_item so that no field is empty
+    if not dataset_item['query']:
+        logger.error("Trying to run workflow with an empty query")
+        return {
+            "input": "invalid-query",
+            "output": "invalid-query",
+            "context": [],
+        }
+    
     response_content = run_workflow(workflow, dataset_item['query'])
 
     # parsed_response = json.loads(response_content.response)
-    # print(parsed_response)
-    # print(parsed_response.keys())
-    # print(parsed_response['text_response'])
+    # logger.info("------> Response: ", parsed_response)
+    # logger.info("------> Response keys: ", parsed_response.keys())
+    # logger.info("------> Response text_response: ", parsed_response['text_response'])
 
     result = {
         "input": dataset_item['query'],
@@ -42,8 +55,10 @@ def build_evaluation_task(params: EvaluationParams):
 
 
 def execute_evaluation(params: EvaluationParams):
+    client = build_opik_client(params.workspace_name, params.api_key)
     dataset = client.get_dataset(name=params.dataset_name)
     base_prompt = client.get_prompt(name=params.base_prompt_name)
+
     if not base_prompt:
         raise ValueError(f"No base prompt found with name '{params.base_prompt_name}'")
 
@@ -63,3 +78,16 @@ def execute_evaluation(params: EvaluationParams):
         prompt=base_prompt,
         task_threads=20
     )
+
+def build_opik_client(workspace_name: str, api_key: str):
+    # Normally you would create the Opik cliekt with Opik(workspace=workspace_name, api_key=api_key)
+    # However, the internal evaluate method gets the client from get_client_cached(), what creates 
+    # a client with the default workspace and without the api key. What is more, it is reused in every request
+
+    # So we use the functions update_session_config to set the parameters, as the Opik constructor can take them from there
+    # Then, we clear the cache to remove any previous client with possible different workspace or api key.
+    # Finally, we return the client created with the new parameters. Now "evaluate" will use this client we created.
+    update_session_config('workspace', workspace_name)
+    update_session_config('api_key', api_key)
+    get_client_cached.cache_clear()
+    return get_client_cached()

apps/aifindr-evaluations-runner/main.py

@@ -1,3 +1,5 @@
+from fastapi import FastAPI, HTTPException, Request
+from pydantic import BaseModel
 import logging
 import asyncio
 import uuid
@@ -6,9 +8,6 @@
 from urllib.parse import urljoin, urlparse
 from typing import Dict
 
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-
 from settings import settings
 from evaluator import EvaluationParams, ExperimentStatus, execute_evaluation
 
@@ -25,9 +24,10 @@
 
 
 class RunEvaluationsRequest(BaseModel):
+    workspace_name: str
     dataset_name: str
     experiment_name: str
-    project_name: str
+    project_name: str | None = None
     base_prompt_name: str
     workflow: str
 
@@ -113,23 +113,25 @@ async def health(timeout: int = 5):
 
 
 @app.post("/evaluations/run", response_model=RunEvaluationsResponse)
-async def run_evaluation(request: RunEvaluationsRequest):
+async def run_evaluation(input: RunEvaluationsRequest, req: Request):
     try:
         # Generate task ID
         task_id = str(uuid.uuid4())
         # Create EvaluationParams with all fields from request plus task_id
         evaluation_params = EvaluationParams(
             task_id=task_id,
-            dataset_name=request.dataset_name,
-            experiment_name=request.experiment_name,
-            project_name=request.project_name,
-            base_prompt_name=request.base_prompt_name,
-            workflow=request.workflow,
+            workspace_name=input.workspace_name,
+            dataset_name=input.dataset_name,
+            experiment_name=input.experiment_name,
+            project_name=input.project_name,
+            base_prompt_name=input.base_prompt_name,
+            workflow=input.workflow,
+            api_key=req.headers.get("Authorization")
         )
 
         try:
             TASK_QUEUE.put_nowait(evaluation_params)
-            logger.info(f"Evaluation task added to queue: {evaluation_params}")
+            logger.info("Evaluation task added to queue")
         except asyncio.QueueFull:
             logger.error(
                 f"Queue is full. Evaluation task not added to the queue: {evaluation_params}"

apps/aifindr-evaluations-runner/metrics/follows_criteria.py

@@ -35,8 +35,10 @@ def __init__(self, prompt_template: str, name: str = "Follows criteria", model_n
         self.name = name
         self.llm_client = models.LiteLLMChatModel(model_name=model_name)
         self.prompt_template = f"""
+# Instructions
 {prompt_template}
------
+
+# Answer format
 Answer with a json with the following format:
 
 {{{{
@@ -58,8 +60,6 @@ def score(self, output: str, criteria: str, **ignored_kwargs: Any):
             output=output,
             criteria=criteria
         )
-
-        print("Prompt total: ", prompt)
         # Generate and parse the response from the LLM
         response = self.llm_client.generate_string(input=prompt, response_format=FollowsCriteriaResult)
 

apps/aifindr-evaluations-runner/workflows.py

@@ -7,7 +7,7 @@
 from pydantic import BaseModel
 from typing import Optional, List, Any
 
-MAX_RETRIES = 3
+MAX_RETRIES = 5
 RETRIEVAL_EVENT_ID_PREFIX = "similarity_search_by_text"
 LLM_EVENT_ID_PREFIX = "llm"
 
@@ -35,7 +35,7 @@ def run_workflow(workflow: str, query: str) -> WorkflowResponse:
             return _make_workflow_request(workflow, query)
         except Exception as e:
             wait_time = 0.5 * (retry_count + 1)  # Increasing delay between retries
-            print(f"Request failed with error: {e}. Waiting {wait_time}s before retrying... ({retry_count + 1}/{MAX_RETRIES})")
+            logger.warn(f"Request failed with error: {e}. Waiting {wait_time}s before retrying... ({retry_count + 1}/{MAX_RETRIES})")
             retry_count += 1
             time.sleep(wait_time)
 
@@ -85,6 +85,7 @@ def _make_workflow_request(workflow: str, query: str) -> WorkflowResponse:
         try:
             data = json.loads(event.data)
             if event.id.startswith(RETRIEVAL_EVENT_ID_PREFIX):
+                logger.debug(f"Retrieval response: {data}")
                 retrieval_response = data['response']['hits']
             elif event.id.startswith(LLM_EVENT_ID_PREFIX) and 'content' in data['delta']['message']:
                 llm_response += data['delta']['message']['content']

apps/opik-backend/config.yml

@@ -129,11 +129,14 @@ authentication:
   # Default:
   # Description: Configures how to authenticate requests which originates from the sdk
   sdk:
-    url: ${AUTH_SDK_URL:-''}
+    url: ${AUTH_DOMAIN:-''}/userinfo
+    # AIFindr: Whether the URL is from auth0 (ended in /userinfo)
+    isAuth0: ${AUTH_IS_AUTH0:-false}
   # Default:
   # Description: Configures how to authenticate requests which originates from the ui
   ui:
-    url: ${AUTH_UI_URL:-''}
+    url: ${AUTH_DOMAIN:-''}/userinfo
+    isAuth0: false
 
 # https://www.dropwizard.io/en/stable/manual/configuration.html#servers
 server:
@@ -285,3 +288,6 @@ clickHouseLogAppender:
   # Default: PT0.500S or 500ms
   # Description: Time interval after which the log messages are sent to ClickHouse if the batch size is not reached
   flushIntervalDuration: ${CLICKHOUSE_LOG_APPENDER_FLUSH_INTERVAL_DURATION:-PT0.500S}
+
+experimentRunner:
+  url: ${EXPERIMENT_RUNNER_URL:-''}

apps/opik-backend/src/main/java/com/comet/opik/api/ExperimentRunRequest.java

@@ -0,0 +1,20 @@
+package com.comet.opik.api;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.databind.PropertyNamingStrategies;
+import com.fasterxml.jackson.databind.annotation.JsonNaming;
+
+import jakarta.validation.constraints.NotBlank;
+import lombok.Builder;
+
+@Builder(toBuilder = true)
+@JsonIgnoreProperties(ignoreUnknown = true)
+@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
+public record ExperimentRunRequest(
+    @NotBlank String workspaceName,
+    @NotBlank String datasetName,
+    @NotBlank String experimentName,
+    String projectName,
+    @NotBlank String basePromptName,
+    @NotBlank String workflow
+) {} 

apps/opik-backend/src/main/java/com/comet/opik/api/ExperimentRunResponse.java

@@ -0,0 +1,10 @@
+package com.comet.opik.api;
+
+import com.fasterxml.jackson.databind.PropertyNamingStrategies;
+import com.fasterxml.jackson.databind.annotation.JsonNaming;
+
+@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
+public record ExperimentRunResponse(
+    String status,
+    String taskId
+) {} 

apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/ExperimentsResource.java

@@ -1,12 +1,22 @@
 package com.comet.opik.api.resources.v1.priv;
 
+import java.util.Collections;
+import java.util.Optional;
+import java.util.Set;
+import java.util.UUID;
+import java.util.stream.Collectors;
+
+import org.glassfish.jersey.server.ChunkedOutput;
+
 import com.codahale.metrics.annotation.Timed;
 import com.comet.opik.api.Experiment;
 import com.comet.opik.api.ExperimentItem;
 import com.comet.opik.api.ExperimentItemSearchCriteria;
 import com.comet.opik.api.ExperimentItemStreamRequest;
 import com.comet.opik.api.ExperimentItemsBatch;
 import com.comet.opik.api.ExperimentItemsDelete;
+import com.comet.opik.api.ExperimentRunRequest;
+import com.comet.opik.api.ExperimentRunResponse;
 import com.comet.opik.api.ExperimentSearchCriteria;
 import com.comet.opik.api.ExperimentsDelete;
 import com.comet.opik.api.FeedbackDefinition;
@@ -15,14 +25,17 @@
 import com.comet.opik.api.resources.v1.priv.validate.IdParamsValidator;
 import com.comet.opik.domain.ExperimentItemService;
 import com.comet.opik.domain.ExperimentService;
+import com.comet.opik.domain.FeedbackScoreDAO.EntityType;
 import com.comet.opik.domain.FeedbackScoreService;
 import com.comet.opik.domain.IdGenerator;
 import com.comet.opik.domain.Streamer;
 import com.comet.opik.infrastructure.auth.RequestContext;
 import com.comet.opik.infrastructure.ratelimit.RateLimited;
 import com.comet.opik.utils.AsyncUtils;
+import static com.comet.opik.utils.AsyncUtils.setRequestContext;
 import com.fasterxml.jackson.annotation.JsonView;
 import com.fasterxml.jackson.databind.JsonNode;
+
 import io.dropwizard.jersey.errors.ErrorMessage;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.headers.Header;
@@ -40,6 +53,7 @@
 import jakarta.ws.rs.Consumes;
 import jakarta.ws.rs.DefaultValue;
 import jakarta.ws.rs.GET;
+import jakarta.ws.rs.HeaderParam;
 import jakarta.ws.rs.POST;
 import jakarta.ws.rs.Path;
 import jakarta.ws.rs.PathParam;
@@ -52,16 +66,6 @@
 import lombok.NonNull;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
-import org.glassfish.jersey.server.ChunkedOutput;
-
-import java.util.Collections;
-import java.util.Optional;
-import java.util.Set;
-import java.util.UUID;
-import java.util.stream.Collectors;
-
-import static com.comet.opik.domain.FeedbackScoreDAO.EntityType;
-import static com.comet.opik.utils.AsyncUtils.setRequestContext;
 
 @Path("/v1/private/experiments")
 @Produces(MediaType.APPLICATION_JSON)
@@ -304,4 +308,29 @@ public Response findFeedbackScoreNames(@QueryParam("experiment_ids") String expe
 
         return Response.ok(feedbackScoreNames).build();
     }
+
+    @POST
+    @Path("/run")
+    @Operation(operationId = "runExperiment", 
+        summary = "Run experiment", 
+        description = "Run experiment with specified dataset, project and workflow configuration", 
+        responses = {
+            @ApiResponse(responseCode = "200", description = "Experiment run started", 
+                content = @Content(schema = @Schema(implementation = ExperimentRunResponse.class)))})
+    @RateLimited
+    public Response runExperiment(
+            @RequestBody(content = @Content(schema = @Schema(implementation = ExperimentRunRequest.class))) 
+            @NotNull @Valid ExperimentRunRequest request,
+            @HeaderParam("Authorization") String authorization) {
+
+        log.info("Running experiment {}", request);
+        
+        ExperimentRunResponse response = experimentService.runExperiment(request, authorization)
+                .contextWrite(ctx -> setRequestContext(ctx, requestContext))
+                .block();
+        
+        log.info("Experiment started. Status: '{}', Task ID: '{}'", 
+                response.status(), response.taskId());
+        return Response.ok(response).build();
+    }
 }