106 lines
2.6 KiB
Python
106 lines
2.6 KiB
Python
from datetime import datetime
|
|
from typing import List, Dict, Any, Optional, Union
|
|
from pydantic import BaseModel, Field
|
|
from traceloop.sdk.evaluator.config import EvaluatorDetails
|
|
|
|
EvaluatorSpec = Union[str, EvaluatorDetails]
|
|
|
|
|
|
class TaskResponse(BaseModel):
|
|
"""Model for a single task process (row)"""
|
|
|
|
task_result: Optional[Dict[str, Any]] = None
|
|
evaluations: Dict[str, Any] = Field(default_factory=dict)
|
|
error: Optional[str] = None
|
|
|
|
|
|
class InitExperimentRequest(BaseModel):
|
|
"""Model for initializing an experiment"""
|
|
|
|
slug: str
|
|
dataset_slug: Optional[str] = None
|
|
dataset_version: Optional[str] = None
|
|
evaluator_slugs: Optional[List[str]] = None
|
|
experiment_metadata: Optional[Dict[str, Any]] = None
|
|
experiment_run_metadata: Optional[Dict[str, Any]] = None
|
|
|
|
|
|
class ExperimentResponse(BaseModel):
|
|
"""Model for experiment response"""
|
|
|
|
id: str
|
|
slug: str
|
|
metadata: Optional[Dict[str, Any]] = None
|
|
created_at: datetime
|
|
updated_at: datetime
|
|
|
|
|
|
class ExperimentRunResponse(BaseModel):
|
|
"""Model for experiment run response"""
|
|
|
|
id: str
|
|
metadata: Optional[Dict[str, Any]] = None
|
|
dataset_id: Optional[str] = None
|
|
dataset_version: Optional[str] = None
|
|
evaluator_ids: Optional[List[str]] = None
|
|
created_at: datetime
|
|
updated_at: datetime
|
|
|
|
|
|
class ExperimentInitResponse(BaseModel):
|
|
"""Model for experiment and run response"""
|
|
|
|
experiment: ExperimentResponse
|
|
run: ExperimentRunResponse
|
|
|
|
|
|
class CreateTaskRequest(BaseModel):
|
|
"""Model for create task request"""
|
|
|
|
input: Dict[str, Any]
|
|
output: Dict[str, Any]
|
|
|
|
|
|
class CreateTaskResponse(BaseModel):
|
|
"""Model for create task response"""
|
|
|
|
id: str
|
|
|
|
|
|
class TaskResult(BaseModel):
|
|
"""Model for a single task result"""
|
|
|
|
input: Dict[str, Any]
|
|
output: Optional[Dict[str, Any]] = None
|
|
error: Optional[str] = None
|
|
|
|
|
|
class GithubContext(BaseModel):
|
|
"""Model for GitHub context"""
|
|
|
|
repository: str
|
|
pr_url: str
|
|
commit_hash: str
|
|
actor: str
|
|
|
|
|
|
class RunInGithubRequest(BaseModel):
|
|
"""Model for bulk GitHub experiment execution request"""
|
|
|
|
experiment_slug: str
|
|
dataset_slug: Optional[str] = None
|
|
dataset_version: Optional[str] = None
|
|
evaluator_slugs: Optional[List[str]] = None
|
|
task_results: List[TaskResult]
|
|
github_context: GithubContext
|
|
experiment_metadata: Optional[Dict[str, Any]] = None
|
|
experiment_run_metadata: Optional[Dict[str, Any]] = None
|
|
|
|
|
|
class RunInGithubResponse(BaseModel):
|
|
"""Model for bulk GitHub experiment execution response"""
|
|
|
|
experiment_id: str
|
|
experiment_slug: str
|
|
run_id: str
|