codemie-sdk-python 0.1.1__py3-none-any.whl

codemie_sdk/models/datasource.py

@@ -0,0 +1,293 @@
+import re
+from datetime import datetime
+from enum import Enum
+from typing import Optional, List
+
+from pydantic import BaseModel, Field, model_validator, ConfigDict, field_validator
+
+from .common import TokensUsage, User
+
+
+class CodeDataSourceType(str, Enum):
+    CODE = "code"
+    SUMMARY = "summary"
+    CHUNK_SUMMARY = "chunk-summary"
+
+
+class DataSourceType(str, Enum):
+    CODE = "code"
+    CONFLUENCE = "knowledge_base_confluence"
+    JIRA = "knowledge_base_jira"
+    FILE = "knowledge_base_file"
+    GOOGLE = "llm_routing_google"
+
+
+class DataSourceStatus(str, Enum):
+    COMPLETED = "completed"
+    FAILED = "failed"
+    FETCHING = "fetching"
+    IN_PROGRESS = "in_progress"
+
+
+class DataSourceProcessingInfo(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    total_documents_count: Optional[int] = Field(None, alias="total_documents")
+    skipped_documents_count: Optional[int] = Field(None, alias="skipped_documents")
+    total_size_kb: Optional[float] = None
+    average_file_size_bytes: Optional[float] = None
+    unique_extensions: Optional[List[str]] = None
+    filtered_documents: Optional[int] = None
+    processed_documents_count: Optional[int] = Field(None, alias="documents_count_key")
+
+
+# Base request models
+class Confluence(BaseModel):
+    """Model for Confluence-specific response fields"""
+
+    cql: Optional[str] = None
+    include_restricted_content: Optional[bool] = None
+    include_archived_content: Optional[bool] = None
+    include_attachments: Optional[bool] = None
+    include_comments: Optional[bool] = None
+    keep_markdown_format: Optional[bool] = None
+    keep_newlines: Optional[bool] = None
+    max_pages: Optional[int] = None
+    pages_per_request: Optional[int] = None
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class Jira(BaseModel):
+    """Model for Jira-specific response fields"""
+
+    jql: Optional[str] = None
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class Google(BaseModel):
+    """Model for Jira-specific response fields"""
+
+    google_doc: str = Field(None, alias="googleDoc")
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class Code(BaseModel):
+    """Model for code repository datasource creation"""
+
+    link: Optional[str] = Field(..., min_length=1, max_length=1000)
+    branch: Optional[str] = Field(..., min_length=1, max_length=1000)
+    index_type: Optional[CodeDataSourceType] = Field(None, alias="indexType")
+    files_filter: Optional[str] = Field(default="", alias="filesFilter")
+    embeddings_model: Optional[str] = Field(None, alias="embeddingsModel")
+    summarization_model: Optional[str] = Field(None, alias="summarizationModel")
+    prompt: Optional[str] = None
+    docs_generation: bool = Field(False, alias="docsGeneration")
+
+
+class BaseDataSourceRequest(BaseModel):
+    """Base model for all datasource creation requests"""
+
+    name: str = Field(
+        ...,
+        description="Name must contain only lowercase letters and underscores.",
+        max_length=25,
+    )
+    project_name: str
+    description: str = Field(..., max_length=100)
+    shared_with_project: bool = Field(False, alias="project_space_visible")
+    setting_id: str = Field(None)
+    type: DataSourceType
+
+    model_config = ConfigDict(populate_by_name=True, extra="allow")
+
+    @classmethod
+    def required_fields(cls) -> List[str]:
+        return []
+
+    @field_validator("name")
+    def validate_name_field(cls, value):
+        if not re.fullmatch(r"^[a-z][a-z_-]*$", value):
+            raise ValueError(
+                "Name must contain only lowercase letters and underscores, and cannot begin with '_' or '-'."
+            )
+        return value
+
+    @model_validator(mode="before")
+    def pre_init_validator(cls, values):
+        # Ensure that at least one of the fields is not None
+        if cls.required_fields() and not any(
+            values.get(field) for field in cls.required_fields()
+        ):
+            raise ValueError(
+                f"At least one of the following fields must be set: {', '.join(cls.required_fields())}"
+            )
+        return values
+
+
+class CodeDataSourceRequest(BaseDataSourceRequest, Code):
+    # Override to comply with keys in json
+    project_space_visible: bool = Field(False, alias="projectSpaceVisible")
+    setting_id: str = Field(None, alias="settingId")
+
+    def __init__(self, **data):
+        super().__init__(type=DataSourceType.CODE, **data)
+
+    @classmethod
+    def required_fields(cls) -> List[str]:
+        return ["link", "embeddings_model", "branch", "index_type"]
+
+
+class JiraDataSourceRequest(BaseDataSourceRequest, Jira):
+    def __init__(self, **data):
+        super().__init__(type=DataSourceType.JIRA, **data)
+
+    @classmethod
+    def required_fields(cls) -> List[str]:
+        return ["jql"]
+
+
+class ConfluenceDataSourceRequest(BaseDataSourceRequest, Confluence):
+    def __init__(self, **data):
+        super().__init__(type=DataSourceType.CONFLUENCE, **data)
+
+    @classmethod
+    def required_fields(cls) -> List[str]:
+        return ["cql"]
+
+
+class GoogleDataSourceRequest(BaseDataSourceRequest, Google):
+    def __init__(self, **data):
+        super().__init__(type=DataSourceType.GOOGLE, **data)
+
+    @classmethod
+    def required_fields(cls) -> List[str]:
+        return ["google_doc"]
+
+
+class BaseUpdateDataSourceRequest(BaseDataSourceRequest):
+    """Mixin update-specific reindex fields"""
+
+    full_reindex: Optional[bool] = Field(None)
+    skip_reindex: Optional[bool] = Field(None)
+    resume_indexing: Optional[bool] = Field(None)
+    incremental_reindex: Optional[bool] = Field(None)
+
+    model_config = ConfigDict(populate_by_name=True, extra="allow")
+
+    @model_validator(mode="after")
+    def validate_reindex_options(self) -> "BaseUpdateDataSourceRequest":
+        ds_type = self.type
+
+        if ds_type == DataSourceType.CONFLUENCE:
+            if self.incremental_reindex:
+                raise ValueError(
+                    "Confluence data sources only support full_reindex and resume_indexing"
+                )
+
+        elif ds_type == DataSourceType.JIRA:
+            if self.resume_indexing:
+                raise ValueError(
+                    "Jira data sources only support full_reindex and incremental_reindex"
+                )
+
+        elif ds_type == DataSourceType.CODE:
+            if self.incremental_reindex:
+                raise ValueError("Code data sources do not support incremental_reindex")
+
+        elif ds_type == DataSourceType.GOOGLE:
+            if self.resume_indexing or self.incremental_reindex:
+                raise ValueError("Google data sources only support full_reindex")
+
+        return self
+
+
+class UpdateCodeDataSourceRequest(BaseUpdateDataSourceRequest, Code):
+    """Model for code repository datasource updates"""
+
+    def __init__(self, **data):
+        super().__init__(type=DataSourceType.CODE, **data)
+
+
+class UpdateConfluenceDataSourceRequest(BaseUpdateDataSourceRequest, Confluence):
+    def __init__(self, **data):
+        super().__init__(type=DataSourceType.CONFLUENCE, **data)
+
+
+class UpdateJiraDataSourceRequest(BaseUpdateDataSourceRequest, Jira):
+    def __init__(self, **data):
+        super().__init__(type=DataSourceType.JIRA, **data)
+
+
+class UpdateGoogleDataSourceRequest(BaseUpdateDataSourceRequest):
+    """Model for Google docs datasource updates"""
+
+    google_doc: Optional[str] = Field(None, alias="googleDoc")
+
+    def __init__(self, **data):
+        super().__init__(type=DataSourceType.GOOGLE, **data)
+
+
+class DataSource(BaseModel):
+    model_config = ConfigDict(
+        extra="ignore",
+        populate_by_name=True,
+    )
+
+    id: str
+    project_name: str
+    name: str = Field(None, alias="repo_name")
+    description: Optional[str] = Field(None)
+    type: DataSourceType = Field(None, alias="index_type")
+    embeddings_model: Optional[str] = Field(None)
+    status: DataSourceStatus
+    setting_id: Optional[str] = Field(None)
+    created_date: datetime = Field(None, alias="date")
+    created_by: User
+    shared_with_project: bool = Field(None, alias="project_space_visible")
+    update_date: datetime
+    error_message: Optional[str] = Field(None, alias="text")
+    user_abilities: List[str]
+    processing_info: Optional[DataSourceProcessingInfo] = Field(None)
+    processed_documents: Optional[List[str]] = Field(None, alias="processed_files")
+    tokens_usage: Optional[TokensUsage] = Field(None)
+    # Code specific fields
+    code: Optional[Code] = None
+    # Jira specific fields
+    jira: Optional[Jira] = None
+    # Confluence specific fields
+    confluence: Optional[Confluence] = None
+    # Google doc specific fields
+    google_doc_link: Optional[str] = None
+
+    @model_validator(mode="before")
+    def before_init(cls, values):
+        if values.get("error"):
+            values["status"] = DataSourceStatus.FAILED
+        elif values.get("completed"):
+            values["status"] = DataSourceStatus.COMPLETED
+        elif values.get("is_fetching"):
+            values["status"] = DataSourceStatus.FETCHING
+        else:
+            values["status"] = DataSourceStatus.IN_PROGRESS
+
+        if values.get("index_type") in [
+            DataSourceType.CONFLUENCE,
+            DataSourceType.JIRA,
+            DataSourceType.GOOGLE,
+        ]:
+            complete_state = values.get("complete_state", 0)
+            if complete_state is not None:
+                values["processing_info"] = {"documents_count_key": complete_state}
+        elif values.get("index_type") == DataSourceType.CODE:
+            values["code"] = {
+                "link": values.get("link"),
+                "branch": values.get("branch"),
+                "files_filter": values.get("files_filter"),
+                "summarization_prompt": values.get("prompt"),
+                "summarization_model": values.get("summarization_model"),
+                "summarization_docs_generation": values.get("docs_generation"),
+            }
+        return values

codemie_sdk/models/integration.py

@@ -0,0 +1,68 @@
+"""Models for assistant-related data structures."""
+
+from datetime import datetime
+from enum import Enum
+from typing import List, Optional, Any
+
+from pydantic import BaseModel, Field, ConfigDict
+
+
+class CredentialTypes(str, Enum):
+    """Enum for credential types."""
+
+    JIRA = "Jira"
+    CONFLUENCE = "Confluence"
+    GIT = "Git"
+    KUBERNETES = "Kubernetes"
+    AWS = "AWS"
+    GCP = "GCP"
+    KEYCLOAK = "Keycloak"
+    AZURE = "Azure"
+    ELASTIC = "Elastic"
+    OPENAPI = "OpenAPI"
+    PLUGIN = "Plugin"
+    FILESYSTEM = "FileSystem"
+    SCHEDULER = "Scheduler"
+    WEBHOOK = "Webhook"
+    EMAIL = "Email"
+    AZURE_DEVOPS = "AzureDevOps"
+    SONAR = "Sonar"
+    SQL = "SQL"
+    TELEGRAM = "Telegram"
+    ZEPHYR_CLOUD = "ZephyrCloud"
+    ZEPHYR_SQUAD = "ZephyrSquad"
+    SERVICE_NOW = "ServiceNow"
+    DIAL = "DIAL"
+
+
+class IntegrationType(str, Enum):
+    """Enum for setting types."""
+
+    USER = "user"
+    PROJECT = "project"
+
+
+class CredentialValues(BaseModel):
+    """Model for credential values."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    key: str
+    value: Any
+
+
+class Integration(BaseModel):
+    """Model for settings configuration."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    id: Optional[str] = None
+    date: Optional[datetime] = None
+    update_date: Optional[datetime] = None
+    user_id: Optional[str] = None
+    project_name: str
+    alias: Optional[str] = None
+    default: bool = False
+    credential_type: CredentialTypes
+    credential_values: List[CredentialValues]
+    setting_type: IntegrationType = Field(default=IntegrationType.USER)

codemie_sdk/models/llm.py

@@ -0,0 +1,48 @@
+"""Models for LLM service."""
+
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+
+class LLMProvider(str, Enum):
+    """LLM provider options."""
+
+    AZURE_OPENAI = "azure_openai"
+    AWS_BEDROCK = "aws_bedrock"
+    GOOGLE_VERTEXAI = "google_vertexai"
+
+
+class CostConfig(BaseModel):
+    """Cost configuration for LLM model."""
+
+    input: float
+    output: float
+
+
+class LLMFeatures(BaseModel):
+    """Features supported by LLM model."""
+
+    streaming: Optional[bool] = True
+    tools: Optional[bool] = True
+    temperature: Optional[bool] = True
+    parallel_tool_calls: Optional[bool] = True
+    system_prompt: Optional[bool] = True
+    max_tokens: Optional[bool] = True
+
+
+class LLMModel(BaseModel):
+    """LLM model configuration."""
+
+    base_name: str
+    deployment_name: str
+    label: Optional[str] = None
+    multimodal: Optional[bool] = None
+    react_agent: Optional[bool] = None
+    enabled: bool
+    provider: Optional[LLMProvider] = None
+    default: Optional[bool] = None
+    cost: Optional[CostConfig] = None
+    max_output_tokens: Optional[int] = None
+    features: Optional[LLMFeatures] = Field(default_factory=lambda: LLMFeatures())

codemie_sdk/models/task.py

@@ -0,0 +1,44 @@
+from datetime import datetime
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class BackgroundTaskStatus(str, Enum):
+    """Enum for background task statuses."""
+
+    STARTED = "STARTED"
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+
+
+class TaskUser(BaseModel):
+    """Model representing task user information."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    user_id: str = Field(description="Unique identifier of the user")
+    username: str = Field(default="", description="Username of the task owner")
+    name: str = Field(default="", description="Display name of the task owner")
+
+
+class BackgroundTaskEntity(BaseModel):
+    """Model representing a background task."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    id: str = Field(description="Unique identifier of the task")
+    task: str = Field(description="Task description or name")
+    user: TaskUser = Field(description="Information about the task owner")
+    final_output: Optional[str] = Field(
+        default="", description="The final result or output of the task"
+    )
+    current_step: Optional[str] = Field(
+        default="", description="Current step or stage of the task"
+    )
+    status: BackgroundTaskStatus = Field(
+        description="Task status (STARTED, COMPLETED, or FAILED)"
+    )
+    date: datetime = Field(description="Task creation timestamp")
+    update_date: datetime = Field(description="Last update timestamp")

codemie_sdk/models/user.py

@@ -0,0 +1,50 @@
+"""Models for user-related data structures."""
+
+from typing import List, Optional
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class User(BaseModel):
+    """Model representing a user profile."""
+
+    model_config = ConfigDict(populate_by_name=True, extra="ignore")
+
+    user_id: str = Field(
+        description="Unique identifier of the user", validation_alias="userId"
+    )
+    name: str = Field(description="Full name of the user")
+    username: str = Field(description="Username for authentication")
+    is_admin: bool = Field(
+        description="Whether the user has admin privileges", validation_alias="isAdmin"
+    )
+    applications: List[str] = Field(
+        default_factory=list, description="List of applications the user has access to"
+    )
+    applications_admin: List[str] = Field(
+        default_factory=list,
+        description="List of applications where user has admin rights",
+        validation_alias="applicationsAdmin",
+    )
+    picture: str = Field(default="", description="URL to user's profile picture")
+    knowledge_bases: List[str] = Field(
+        default_factory=list,
+        description="List of knowledge bases the user has access to",
+        validation_alias="knowledgeBases",
+    )
+
+
+class UserData(BaseModel):
+    """Model representing user data."""
+
+    model_config = ConfigDict(populate_by_name=True, extra="ignore")
+
+    id: Optional[str] = Field(
+        default=None, description="Unique identifier of the user data record"
+    )
+    date: Optional[str] = Field(default=None, description="Creation timestamp")
+    update_date: Optional[str] = Field(
+        default=None, description="Last update timestamp"
+    )
+    user_id: Optional[str] = Field(
+        default=None, description="Associated user identifier"
+    )

codemie_sdk/models/workflow.py

@@ -0,0 +1,86 @@
+"""Workflow models."""
+
+from datetime import datetime
+from enum import StrEnum, Enum
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from codemie_sdk.models.common import User, TokensUsage
+
+
+class WorkflowMode(StrEnum):
+    """Available workflow modes."""
+
+    SEQUENTIAL = "Sequential"
+    AUTONOMOUS = "Autonomous"
+
+
+class WorkflowCreateRequest(BaseModel):
+    """Request model for workflow creation."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    project: str = Field(..., min_length=1)
+    name: str = Field(..., min_length=1)
+    description: Optional[str] = None
+    yaml_config: str = Field(..., min_length=1)
+    mode: WorkflowMode = WorkflowMode.SEQUENTIAL
+    shared: bool = False
+    icon_url: Optional[str] = None
+
+
+class WorkflowUpdateRequest(BaseModel):
+    """Request model for workflow updates."""
+
+    model_config = ConfigDict(populate_by_name=True)
+    project: str = Field(..., min_length=1)
+    name: str = Field(..., min_length=1)
+    description: str = Field(..., min_length=1)
+    yaml_config: str = Field(..., min_length=1)
+    mode: Optional[WorkflowMode] = None
+    shared: Optional[bool] = None
+    icon_url: Optional[str] = None
+
+
+class Workflow(BaseModel):
+    """Workflow template model."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    id: Optional[str] = None
+    project: str
+    name: str
+    description: Optional[str] = None
+    yaml_config: Optional[str] = None
+    mode: WorkflowMode = WorkflowMode.SEQUENTIAL
+    shared: bool = False
+    icon_url: Optional[str] = None
+    created_date: datetime = Field(None, alias="date")
+    update_date: datetime = Field(None)
+    created_by: User
+
+
+class ExecutionStatus(str, Enum):
+    IN_PROGRESS = "In Progress"
+    NOT_STARTED = "Not Started"
+    INTERRUPTED = "Interrupted"
+    FAILED = "Failed"
+    SUCCEEDED = "Succeeded"
+    ABORTED = "Aborted"
+
+
+class WorkflowExecution(BaseModel):
+    """Model representing a workflow execution."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    id: str
+    execution_id: str
+    workflow_id: str
+    status: ExecutionStatus = Field(alias="overall_status")
+    created_date: datetime = Field(alias="date")
+    prompt: str
+    updated_date: Optional[datetime] = None
+    created_by: User
+    tokens_usage: Optional[TokensUsage] = None