enkryptai-sdk 0.1.6__py3-none-any.whl → 1.0.0__py3-none-any.whl

enkryptai_sdk/dto/ai_proxy.py

@@ -0,0 +1,325 @@
+from .base import BaseDTO
+from dataclasses import dataclass, field
+from typing import List, Dict, Any, Set, Optional
+
+
+@dataclass
+class ChatCompletionMessage(BaseDTO):
+    role: str
+    content: str
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionMessage":
+        return cls(
+            role=data.get("role", ""),
+            content=data.get("content", "")
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "role": self.role,
+            "content": self.content
+        }
+
+
+@dataclass
+class ChatCompletionRequest(BaseDTO):
+    model: str
+    messages: List[ChatCompletionMessage]
+    enkrypt_context: Optional[str] = None
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    max_tokens: Optional[int] = None
+    presence_penalty: Optional[float] = None
+    frequency_penalty: Optional[float] = None
+    n: Optional[int] = None
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionRequest":
+        messages_data = data.get("messages", [])
+        
+        # Create a copy of the data without the fields we're explicitly processing
+        extra_fields = {k: v for k, v in data.items() 
+                       if k not in ["messages", "model", "enkrypt_context", 
+                                   "temperature", "top_p", "max_tokens", 
+                                   "presence_penalty", "frequency_penalty", "n"]}
+        
+        return cls(
+            messages=[ChatCompletionMessage.from_dict(msg) for msg in messages_data],
+            model=data.get("model", ""),
+            enkrypt_context=data.get("enkrypt_context", None),
+            temperature=data.get("temperature", None),
+            top_p=data.get("top_p", None),
+            max_tokens=data.get("max_tokens", None),
+            presence_penalty=data.get("presence_penalty", None),
+            frequency_penalty=data.get("frequency_penalty", None),
+            n=data.get("n", 1),
+            _extra_fields=extra_fields
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "messages": [message.to_dict() for message in self.messages],
+            "model": self.model
+        }
+        
+        # Add optional fields only if they are not None or default values
+        if self.enkrypt_context is not None:
+            result["enkrypt_context"] = self.enkrypt_context
+        if self.temperature is not None:
+            result["temperature"] = self.temperature
+        if self.top_p is not None:
+            result["top_p"] = self.top_p
+        if self.max_tokens is not None:
+            result["max_tokens"] = self.max_tokens
+        if self.presence_penalty is not None:
+            result["presence_penalty"] = self.presence_penalty
+        if self.frequency_penalty is not None:
+            result["frequency_penalty"] = self.frequency_penalty
+        if self.n is not None:
+            result["n"] = self.n
+            
+        # Add any extra fields that weren't explicitly defined
+        result.update(self._extra_fields)
+        return result
+
+
+@dataclass
+class ChatCompletionResponseMessage(BaseDTO):
+    role: str
+    content: str
+    tool_calls: List[Dict[str, Any]] = field(default_factory=list)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionResponseMessage":
+        return cls(
+            role=data.get("role", ""),
+            content=data.get("content", ""),
+            tool_calls=data.get("tool_calls", [])
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "role": self.role,
+            "content": self.content
+        }
+        if self.tool_calls:
+            result["tool_calls"] = self.tool_calls
+        return result
+
+
+@dataclass
+class ChatCompletionChoice(BaseDTO):
+    index: int
+    message: ChatCompletionResponseMessage
+    finish_reason: Optional[str] = None
+    seed: Optional[int] = None
+    logprobs: Optional[Any] = None
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionChoice":
+        return cls(
+            index=data.get("index", 0),
+            message=ChatCompletionResponseMessage.from_dict(data.get("message", {})),
+            finish_reason=data.get("finish_reason", None),
+            seed=data.get("seed", None),
+            logprobs=data.get("logprobs", None)
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "index": self.index,
+            "message": self.message.to_dict()
+        }
+        if self.finish_reason is not None:
+            result["finish_reason"] = self.finish_reason
+        if self.seed is not None:
+            result["seed"] = self.seed
+        if self.logprobs is not None:
+            result["logprobs"] = self.logprobs
+        return result
+
+
+@dataclass
+class ChatCompletionUsage(BaseDTO):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionUsage":
+        return cls(
+            prompt_tokens=data.get("prompt_tokens", 0),
+            completion_tokens=data.get("completion_tokens", 0),
+            total_tokens=data.get("total_tokens", 0)
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "prompt_tokens": self.prompt_tokens,
+            "completion_tokens": self.completion_tokens,
+            "total_tokens": self.total_tokens
+        }
+
+
+@dataclass
+class ChatCompletionResponse(BaseDTO):
+    id: str
+    object: str
+    created: int
+    model: str
+    choices: List[ChatCompletionChoice]
+    usage: ChatCompletionUsage
+    prompt: List[Any] = field(default_factory=list)
+    enkrypt_policy_detections: Dict[str, Any] = field(default_factory=dict)
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionResponse":
+        # Extract fields we're explicitly handling
+        choices_data = data.get("choices", [])
+        usage_data = data.get("usage", {})
+        
+        # Create a copy of the data without the fields we're explicitly processing
+        extra_fields = {k: v for k, v in data.items() 
+                       if k not in ["id", "object", "created", "model", 
+                                    "choices", "usage", "prompt", "enkrypt_policy_detections"]}
+        
+        return cls(
+            id=data.get("id", ""),
+            object=data.get("object", ""),
+            created=data.get("created", 0),
+            model=data.get("model", ""),
+            choices=[ChatCompletionChoice.from_dict(choice) for choice in choices_data],
+            usage=ChatCompletionUsage.from_dict(usage_data),
+            prompt=data.get("prompt", []),
+            enkrypt_policy_detections=data.get("enkrypt_policy_detections", {}),
+            _extra_fields=extra_fields
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "id": self.id,
+            "object": self.object,
+            "created": self.created,
+            "model": self.model,
+            "choices": [choice.to_dict() for choice in self.choices],
+            "usage": self.usage.to_dict(),
+            "prompt": self.prompt,
+            "enkrypt_policy_detections": self.enkrypt_policy_detections
+        }
+        
+        # Add any extra fields
+        result.update(self._extra_fields)
+        return result
+
+
+@dataclass
+class ChatCompletionError(BaseDTO):
+    code: str
+    param: Optional[str] = None
+    type: Optional[str] = None
+    message: Optional[str] = None
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionError":
+        # If data is string, then parse the json
+        if isinstance(data, str):
+            import json
+            data = json.loads(data)
+        # Create a copy of the data without the fields we're explicitly processing
+        extra_fields = {k: v for k, v in data.items() 
+                       if k not in ["code", "param", "type", "message"]}
+        
+        return cls(
+            code=data.get("code", ""),
+            param=data.get("param", None),
+            type=data.get("type", None),
+            message=data.get("message", None),
+            _extra_fields=extra_fields
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "code": self.code
+        }
+        
+        # Add optional fields only if they are not None
+        if self.param is not None:
+            result["param"] = self.param
+        if self.type is not None:
+            result["type"] = self.type
+        if self.message is not None:
+            result["message"] = self.message
+            
+        # Add any extra fields
+        result.update(self._extra_fields)
+        return result
+
+
+@dataclass
+class ChatCompletionErrorResponse(BaseDTO):
+    error: ChatCompletionError
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionErrorResponse":
+        error_data = data.get("error", {})
+        
+        # Create a copy of the data without the fields we're explicitly processing
+        extra_fields = {k: v for k, v in data.items() 
+                       if k not in ["error"]}
+        
+        return cls(
+            error=ChatCompletionError.from_dict(error_data),
+            _extra_fields=extra_fields
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            "error": self.error.to_dict()
+        }
+        
+        # Add any extra fields
+        result.update(self._extra_fields)
+        return result
+
+
+@dataclass
+class ChatCompletionDirectErrorResponse(BaseDTO):
+    error: str
+    message: str
+    request_id: Optional[str] = None
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ChatCompletionDirectErrorResponse":
+        # Create a copy of the data without the fields we're explicitly processing
+        extra_fields = {k: v for k, v in data.items() 
+                       if k not in ["error", "message", "request_id"]}
+        
+        return cls(
+            error=data.get("error", None),
+            message=data.get("message", None),
+            request_id=data.get("request_id", None),
+            _extra_fields=extra_fields
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = {
+            # "message": self.message
+        }
+        
+        # Add optional fields only if they are not None
+        if self.error is not None:
+            result["error"] = self.error
+        if self.message is not None:
+            result["message"] = self.message
+        if self.request_id is not None:
+            result["request_id"] = self.request_id
+        
+        # Add any extra fields
+        result.update(self._extra_fields)
+        return result

enkryptai_sdk/dto/base.py

@@ -0,0 +1,70 @@
+import json
+from dataclasses import dataclass, field
+from typing import Type, Dict, Any, TypeVar
+
+T = TypeVar("T", bound="BaseDTO")
+
+
+@dataclass
+class BaseDTO:
+    """Base class for all DTO classes with common serialization methods."""
+
+    def __getattr__(self, name):
+        """Called when an attribute is not found in the normal places."""
+        try:
+            extra_fields = object.__getattribute__(self, "_extra_fields")
+            if name in extra_fields:
+                return extra_fields[name]
+        except AttributeError:
+            pass
+
+        raise AttributeError(
+            f"'{self.__class__.__name__}' object has no attribute '{name}'"
+        )
+
+    def __setattr__(self, name, value):
+        """Handle setting attributes, storing unknown ones in _extra_fields."""
+        if name in self.__class__.__annotations__ or name == "_extra_fields":
+            super().__setattr__(name, value)
+        else:
+            self._extra_fields[name] = value
+
+    @classmethod
+    def from_dict(cls: Type[T], data: Dict[str, Any]) -> T:
+        """Create an instance from a dictionary."""
+        known_fields = cls.__annotations__.keys()
+        base_data = {k: v for k, v in data.items() if k in known_fields}
+        extra_fields = {
+            k: v
+            for k, v in data.items()
+            if k not in known_fields and k != "_extra_fields"
+        }
+
+        instance = cls(**base_data)
+        instance._extra_fields = extra_fields
+        return instance
+
+    @classmethod
+    def from_json(cls: Type[T], json_str: str) -> T:
+        """Create an instance from a JSON string."""
+        data = json.loads(json_str)
+        return cls.from_dict(data)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the instance to a dictionary."""
+        d = {k: v for k, v in self.__dict__.items() if k != "_extra_fields"}
+        d.update(self._extra_fields)
+        return d
+
+    def to_json(self) -> str:
+        """Convert the instance to a JSON string."""
+        return json.dumps(self.to_dict())
+
+    def get_all_fields(self) -> Dict[str, Any]:
+        """Safely get all fields including extra fields."""
+        regular_fields = {k: v for k, v in vars(self).items() if not k.startswith("_")}
+        extra_fields = vars(self).get("_extra_fields", {})
+
+        result = regular_fields.copy()
+        result.update(extra_fields)
+        return result

enkryptai_sdk/dto/datasets.py

@@ -0,0 +1,152 @@
+from .base import BaseDTO
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Any
+
+
+@dataclass
+class DatasetConfig(BaseDTO):
+    dataset_name: str
+    system_description: str
+    policy_description: str = ""
+    tools: List[str] = field(default_factory=list)
+    info_pdf_url: str = ""
+    max_prompts: int = 100
+    # scenarios: int = 2
+    # categories: int = 2
+    # depth: int = 2
+
+
+@dataclass
+class DatasetCollection(BaseDTO):
+    datasets: List[str] = field(default_factory=list)
+
+
+@dataclass
+class DatasetSummary(BaseDTO):
+    test_types: List[str] = field(default_factory=list)
+
+
+@dataclass
+class DatasetDataPoint(BaseDTO):
+    test_type: str
+    scenario: str
+    category: str
+    prompt: str
+    source: str
+
+
+@dataclass
+class DatasetResponse(BaseDTO):
+    dataset: List[DatasetDataPoint] = field(default_factory=list)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "DatasetResponse":
+        dataset_items = [DatasetDataPoint.from_dict(item) for item in data.get("dataset", [])]
+        return cls(dataset=dataset_items)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {"dataset": [item.to_dict() for item in self.dataset]}
+
+
+@dataclass
+class DatasetContent(BaseDTO):
+    description: str
+    test_types: dict
+    scenarios: dict
+    categories: dict
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "DatasetContent":
+        return cls(
+            description=data.get("description", ""),
+            test_types=data.get("test_types", {}),
+            scenarios=data.get("scenarios", {}),
+            categories=data.get("categories", {}),
+        )
+
+
+@dataclass
+class DatasetTaskData(BaseDTO):
+    dataset_name: str
+    created_at: str
+    started_at: str
+    status: str
+    system_description: str
+    scenarios: int
+    categories: int
+    depth: int
+    policy_description: str
+    project_name: str
+    info_pdf_url: str
+    tools: list
+    task_id: str
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "DatasetTaskData":
+        return cls(
+            dataset_name=data.get("dataset_name", ""),
+            created_at=data.get("created_at", ""),
+            started_at=data.get("started_at", ""),
+            status=data.get("status", ""),
+            system_description=data.get("system_description", ""),
+            scenarios=data.get("scenarios", 0),
+            categories=data.get("categories", 0),
+            depth=data.get("depth", 0),
+            policy_description=data.get("policy_description", ""),
+            project_name=data.get("project_name", ""),
+            info_pdf_url=data.get("info_pdf_url", ""),
+            tools=data.get("tools", []),
+            task_id=data.get("task_id", ""),
+        )
+
+
+@dataclass
+class DatasetCard(BaseDTO):
+    datacard: DatasetContent
+    dataset_name: str
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "DatasetCard":
+        datacard_content = DatasetContent.from_dict(data.get("datacard", {}))
+        return cls(datacard=datacard_content, dataset_name=data.get("dataset_name", ""))
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {"datacard": self.datacard.to_dict(), "dataset_name": self.dataset_name}
+
+
+@dataclass
+class DatasetTaskStatus(BaseDTO):
+    status: str
+    dataset_name: str
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "DatasetTaskStatus":
+        return cls(status=data.get("status", ""), dataset_name=data.get("dataset_name", ""))
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {"status": self.status, "dataset_name": self.dataset_name}
+
+
+@dataclass
+class DatasetTask(BaseDTO):
+    data: DatasetTaskData
+    dataset_name: str
+    _extra_fields: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "DatasetTask":
+        task_data = DatasetTaskData.from_dict(data.get("data", {}))
+        return cls(data=task_data, dataset_name=data.get("dataset_name", ""))
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {"data": self.data.to_dict(), "dataset_name": self.dataset_name}
+
+
+@dataclass
+class DatasetAddTaskResponse(BaseDTO):
+    task_id: str
+    message: Optional[str] = None