codemie-sdk-python 0.1.204__py3-none-any.whl → 0.1.230__py3-none-any.whl

codemie_sdk/models/vendor_knowledgebase.py

@@ -0,0 +1,151 @@
+"""Models for vendor knowledge base settings."""
+
+from datetime import datetime
+from enum import Enum
+from typing import Optional, List
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .vendor_assistant import PaginationInfo, TokenPagination
+
+
+class VendorKnowledgeBaseSetting(BaseModel):
+    """Model representing a vendor knowledge base setting."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    setting_id: str = Field(..., description="Unique identifier for the setting")
+    setting_name: str = Field(..., description="Name of the setting")
+    project: str = Field(..., description="Project associated with the setting")
+    entities: List[str] = Field(
+        default_factory=list, description="List of entities associated with the setting"
+    )
+    invalid: Optional[bool] = Field(None, description="Whether the setting is invalid")
+    error: Optional[str] = Field(
+        None, description="Error message if the setting is invalid"
+    )
+
+
+class VendorKnowledgeBaseSettingsResponse(BaseModel):
+    """Response model for vendor knowledge base settings list."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    data: List[VendorKnowledgeBaseSetting] = Field(
+        ..., description="List of vendor knowledge base settings"
+    )
+    pagination: PaginationInfo = Field(..., description="Pagination information")
+
+
+class VendorKnowledgeBaseStatus(str, Enum):
+    """Status of vendor knowledge base."""
+
+    PREPARED = "PREPARED"
+    NOT_PREPARED = "NOT_PREPARED"
+
+
+class VendorKnowledgeBase(BaseModel):
+    """Model representing a vendor knowledge base."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    id: str = Field(..., description="Unique identifier for the knowledge base")
+    name: str = Field(..., description="Name of the knowledge base")
+    status: VendorKnowledgeBaseStatus = Field(
+        ..., description="Status of the knowledge base"
+    )
+    description: Optional[str] = Field(
+        None, description="Description of the knowledge base"
+    )
+    updatedAt: datetime = Field(
+        ..., description="Last update timestamp", alias="updatedAt"
+    )
+    aiRunId: Optional[str] = Field(
+        None,
+        description="AI run ID if the knowledge base is installed",
+        alias="aiRunId",
+    )
+
+
+class VendorKnowledgeBasesResponse(BaseModel):
+    """Response model for vendor knowledge bases list."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    data: List[VendorKnowledgeBase] = Field(
+        ..., description="List of vendor knowledge bases"
+    )
+    pagination: TokenPagination = Field(
+        ..., description="Token-based pagination information"
+    )
+
+
+class VendorKnowledgeBaseDetail(BaseModel):
+    """Model representing detailed information about a vendor knowledge base."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    id: str = Field(..., description="Unique identifier for the knowledge base")
+    name: str = Field(..., description="Name of the knowledge base")
+    description: Optional[str] = Field(
+        None, description="Description of the knowledge base"
+    )
+    type: str = Field(..., description="Type of knowledge base (e.g., VECTOR)")
+    status: VendorKnowledgeBaseStatus = Field(
+        ..., description="Status of the knowledge base"
+    )
+    embeddingModel: str = Field(
+        ...,
+        description="Embedding model used by the knowledge base",
+        alias="embeddingModel",
+    )
+    kendraIndexArn: Optional[str] = Field(
+        None, description="Kendra index ARN if applicable", alias="kendraIndexArn"
+    )
+    createdAt: datetime = Field(
+        ..., description="Creation timestamp", alias="createdAt"
+    )
+    updatedAt: datetime = Field(
+        ..., description="Last update timestamp", alias="updatedAt"
+    )
+    aiRunId: Optional[str] = Field(
+        None,
+        description="AI run ID if the knowledge base is installed",
+        alias="aiRunId",
+    )
+
+
+class VendorKnowledgeBaseInstallRequest(BaseModel):
+    """Model for a single knowledge base installation request."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    id: str = Field(..., description="Knowledge base ID to install")
+    setting_id: str = Field(..., description="Vendor setting ID")
+
+
+class VendorKnowledgeBaseInstallSummary(BaseModel):
+    """Model for knowledge base installation summary."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    knowledgeBaseId: str = Field(..., description="Installed knowledge base ID")
+    aiRunId: str = Field(..., description="AI run ID for the installation")
+
+
+class VendorKnowledgeBaseInstallResponse(BaseModel):
+    """Response model for knowledge base installation."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    summary: List[VendorKnowledgeBaseInstallSummary] = Field(
+        ..., description="List of installation summaries"
+    )
+
+
+class VendorKnowledgeBaseUninstallResponse(BaseModel):
+    """Response model for knowledge base uninstallation."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    success: bool = Field(..., description="Whether the uninstallation was successful")

codemie_sdk/models/vendor_workflow.py

@@ -0,0 +1,145 @@
+"""Models for vendor workflow settings."""
+
+from datetime import datetime
+from enum import Enum
+from typing import Optional, List
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .vendor_assistant import PaginationInfo, TokenPagination
+
+
+class VendorWorkflowSetting(BaseModel):
+    """Model representing a vendor workflow setting."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    setting_id: str = Field(..., description="Unique identifier for the setting")
+    setting_name: str = Field(..., description="Name of the setting")
+    project: str = Field(..., description="Project associated with the setting")
+    entities: List[str] = Field(
+        default_factory=list, description="List of entities associated with the setting"
+    )
+    invalid: Optional[bool] = Field(None, description="Whether the setting is invalid")
+    error: Optional[str] = Field(
+        None, description="Error message if the setting is invalid"
+    )
+
+
+class VendorWorkflowSettingsResponse(BaseModel):
+    """Response model for vendor workflow settings list."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    data: List[VendorWorkflowSetting] = Field(
+        ..., description="List of vendor workflow settings"
+    )
+    pagination: PaginationInfo = Field(..., description="Pagination information")
+
+
+class VendorWorkflowStatus(str, Enum):
+    """Status of vendor workflow."""
+
+    PREPARED = "PREPARED"
+    NOT_PREPARED = "NOT_PREPARED"
+
+
+class VendorWorkflow(BaseModel):
+    """Model representing a vendor workflow."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    id: str = Field(..., description="Unique identifier for the workflow")
+    name: str = Field(..., description="Name of the workflow")
+    status: VendorWorkflowStatus = Field(..., description="Status of the workflow")
+    description: Optional[str] = Field(None, description="Description of the workflow")
+    version: str = Field(..., description="Version of the workflow")
+    createdAt: datetime = Field(
+        ..., description="Creation timestamp", alias="createdAt"
+    )
+    updatedAt: datetime = Field(
+        ..., description="Last update timestamp", alias="updatedAt"
+    )
+
+
+class VendorWorkflowsResponse(BaseModel):
+    """Response model for vendor workflows list."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    data: List[VendorWorkflow] = Field(..., description="List of vendor workflows")
+    pagination: TokenPagination = Field(
+        ..., description="Token-based pagination information"
+    )
+
+
+class VendorWorkflowAlias(BaseModel):
+    """Model representing a vendor workflow alias."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    id: str = Field(..., description="Unique identifier for the alias")
+    name: str = Field(..., description="Name of the alias")
+    status: VendorWorkflowStatus = Field(..., description="Status of the alias")
+    description: Optional[str] = Field(None, description="Description of the alias")
+    version: str = Field(..., description="Version of the alias")
+    createdAt: datetime = Field(
+        ..., description="Creation timestamp", alias="createdAt"
+    )
+    updatedAt: datetime = Field(
+        ..., description="Last update timestamp", alias="updatedAt"
+    )
+    aiRunId: Optional[str] = Field(
+        None, description="AI run ID if the alias is installed", alias="aiRunId"
+    )
+
+
+class VendorWorkflowAliasesResponse(BaseModel):
+    """Response model for vendor workflow aliases list."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    data: List[VendorWorkflowAlias] = Field(
+        ..., description="List of vendor workflow aliases"
+    )
+    pagination: TokenPagination = Field(
+        ..., description="Token-based pagination information"
+    )
+
+
+class VendorWorkflowInstallRequest(BaseModel):
+    """Model for a single workflow installation request."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    id: str = Field(..., description="Workflow ID to install")
+    flowAliasId: str = Field(..., description="Flow alias ID to use for the workflow")
+    setting_id: str = Field(..., description="Vendor setting ID")
+
+
+class VendorWorkflowInstallSummary(BaseModel):
+    """Model for workflow installation summary."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    flowId: str = Field(..., description="Installed workflow ID")
+    flowAliasId: str = Field(..., description="Flow alias ID used for installation")
+    aiRunId: str = Field(..., description="AI run ID for the installation")
+
+
+class VendorWorkflowInstallResponse(BaseModel):
+    """Response model for workflow installation."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    summary: List[VendorWorkflowInstallSummary] = Field(
+        ..., description="List of installation summaries"
+    )
+
+
+class VendorWorkflowUninstallResponse(BaseModel):
+    """Response model for workflow uninstallation."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    success: bool = Field(..., description="Whether the uninstallation was successful")

codemie_sdk/models/workflow_execution_payload.py

@@ -1,6 +1,6 @@
 """Workflow execution payload models."""
 
-from typing import Optional
+from typing import Optional, Union
 from pydantic import BaseModel, ConfigDict, Field
 
 
@@ -9,9 +9,13 @@ class WorkflowExecutionCreateRequest(BaseModel):
 
     model_config = ConfigDict(populate_by_name=True)
 
-    user_input: Optional[str] = Field(
+    user_input: Optional[Union[str, dict, list, int, float, bool]] = Field(
         None, description="User input for the workflow execution"
     )
     file_name: Optional[str] = Field(
         None, description="File name associated with the workflow execution"
     )
+    propagate_headers: bool = Field(
+        default=False,
+        description="Enable propagation of X-* HTTP headers to MCP servers during tool execution",
+    )

codemie_sdk/services/assistant.py

@@ -1,5 +1,6 @@
 """Assistant service implementation."""
 
+import inspect
 import json
 from pathlib import Path
 from typing import List, Union, Optional, Dict, Any, Literal
@@ -169,23 +170,162 @@ class AssistantService:
         """
         return self._api.get(f"/v1/assistants/prebuilt/{slug}", Assistant)
 
+    def list_versions(
+        self, assistant_id: str, page: int = 0, per_page: Optional[int] = None
+    ):
+        """List assistant versions.
+
+        Args:
+            assistant_id: Assistant identifier
+            page: Page number for pagination
+            per_page: Items per page (optional). If not provided, backend defaults are used.
+
+        Returns:
+            List of AssistantVersion objects
+        """
+
+        params: Dict[str, Any] = {"page": page}
+        if per_page is not None:
+            params["per_page"] = per_page
+        from ..models.assistant import AssistantVersion
+
+        raw = self._api.get(
+            f"/v1/assistants/{assistant_id}/versions",
+            dict,
+            params=params,
+            wrap_response=False,
+        )
+        items = []
+        if isinstance(raw, list):
+            items = raw
+        elif isinstance(raw, dict):
+            items = raw.get("data") or raw.get("versions") or []
+        else:
+            items = []
+        return [AssistantVersion.model_validate(it) for it in items]
+
+    def get_version(self, assistant_id: str, version_number: int):
+        """Get a specific assistant version by number.
+
+        Args:
+            assistant_id: Assistant identifier
+            version_number: Version number to retrieve
+
+        Returns:
+            AssistantVersion object
+        """
+        from ..models.assistant import AssistantVersion
+
+        raw = self._api.get(
+            f"/v1/assistants/{assistant_id}/versions/{version_number}", AssistantVersion
+        )
+        if isinstance(raw, dict):
+            return AssistantVersion.model_validate(raw)
+        return raw
+
+    def compare_versions(self, assistant_id: str, v1: int, v2: int) -> Dict[str, Any]:
+        """Compare two assistant versions and return diff summary.
+
+        Args:
+            assistant_id: Assistant identifier
+            v1: First version number
+            v2: Second version number
+
+        Returns:
+            Generic dictionary with comparison result (diff, summary, etc.)
+        """
+        return self._api.get(
+            f"/v1/assistants/{assistant_id}/versions/{v1}/compare/{v2}",
+            dict,
+        )
+
+    def rollback_to_version(
+        self, assistant_id: str, version_number: int, change_notes: Optional[str] = None
+    ) -> dict:
+        """Rollback assistant to a specific version. Creates a new version mirroring target.
+
+        Args:
+            assistant_id: Assistant identifier
+            version_number: Target version to rollback to
+            change_notes: Optional description of why rollback is performed
+
+        Returns:
+            Backend response (dict)
+        """
+        payload: Dict[str, Any] = {}
+        if change_notes:
+            payload["change_notes"] = change_notes
+        try:
+            return self._api.post(
+                f"/v1/assistants/{assistant_id}/versions/{version_number}/rollback",
+                dict,
+                json_data=payload,
+            )
+        except requests.HTTPError as err:
+            try:
+                assistant = self.get(assistant_id)
+                version = self.get_version(assistant_id, version_number)
+
+                update_req = AssistantUpdateRequest(
+                    name=assistant.name,
+                    description=assistant.description or "",
+                    system_prompt=version.system_prompt,
+                    project=assistant.project,
+                    llm_model_type=version.llm_model_type or assistant.llm_model_type,
+                    temperature=version.temperature
+                    if hasattr(version, "temperature")
+                    else assistant.temperature,
+                    top_p=version.top_p
+                    if hasattr(version, "top_p")
+                    else assistant.top_p,
+                    context=version.context
+                    if hasattr(version, "context")
+                    else assistant.context,
+                    toolkits=version.toolkits
+                    if hasattr(version, "toolkits")
+                    else assistant.toolkits,
+                    user_prompts=assistant.user_prompts,
+                    shared=assistant.shared,
+                    is_react=assistant.is_react,
+                    is_global=assistant.is_global,
+                    slug=assistant.slug,
+                    mcp_servers=version.mcp_servers
+                    if hasattr(version, "mcp_servers")
+                    else assistant.mcp_servers,
+                    assistant_ids=version.assistant_ids
+                    if hasattr(version, "assistant_ids")
+                    else assistant.assistant_ids,
+                )
+                resp = self.update(assistant_id, update_req)
+                resp["_rollback_fallback"] = True
+                resp["_target_version"] = version_number
+                if change_notes:
+                    resp["change_notes"] = change_notes
+                return resp
+            except Exception:
+                raise err
+
     def chat(
         self,
         assistant_id: str,
         request: AssistantChatRequest,
+        headers: Optional[Dict[str, str]] = None,
     ) -> Union[requests.Response, BaseModelResponse]:
         """Send a chat request to an assistant.
 
         Args:
             assistant_id: ID of the assistant to chat with
             request: Chat request details
+            headers: Optional additional HTTP headers (e.g., X-* for MCP propagation)
 
         Returns:
             Chat response or streaming response
         """
         pydantic_schema = None
-        if request.output_schema is not None and issubclass(
-            request.output_schema, BaseModel
+        if (
+            request.output_schema is not None
+            and inspect.isclass(request.output_schema)
+            and issubclass(request.output_schema, BaseModel)
         ):
             pydantic_schema = deepcopy(request.output_schema)
             request.output_schema = request.output_schema.model_json_schema()
@@ -195,6 +335,7 @@ class AssistantService:
             BaseModelResponse,
             json_data=request.model_dump(exclude_none=True, by_alias=True),
             stream=request.stream,
+            extra_headers=headers,
         )
         if not request.stream and pydantic_schema:
             # we do conversion to the BaseModel here because self._parse_response don't see actual request model,
@@ -203,6 +344,82 @@ class AssistantService:
 
         return response
 
+    def chat_with_version(
+        self,
+        assistant_id: str,
+        version_number: int,
+        request: AssistantChatRequest,
+    ) -> Union[requests.Response, BaseModelResponse]:
+        """Send a chat request to a specific assistant version.
+
+        Uses the stable chat endpoint with an explicit `version` parameter to
+        ensure compatibility with environments that don't expose
+        /versions/{version}/model.
+
+        Args:
+            assistant_id: ID of the assistant to chat with
+            version_number: version to pin chat to
+            request: Chat request details
+
+        Returns:
+            Chat response or streaming response
+        """
+        pydantic_schema = None
+        if issubclass(request.output_schema, BaseModel):
+            pydantic_schema = deepcopy(request.output_schema)
+            request.output_schema = request.output_schema.model_json_schema()
+
+        payload = request.model_dump(exclude_none=True, by_alias=True)
+        payload["version"] = version_number
+
+        response = self._api.post(
+            f"/v1/assistants/{assistant_id}/model",
+            BaseModelResponse,
+            json_data=payload,
+            stream=request.stream,
+        )
+        if not request.stream and pydantic_schema:
+            response.generated = pydantic_schema.model_validate(response.generated)
+
+        return response
+
+    def chat_by_slug(
+        self,
+        assistant_slug: str,
+        request: AssistantChatRequest,
+        headers: Optional[Dict[str, str]] = None,
+    ) -> Union[requests.Response, BaseModelResponse]:
+        """Send a chat request to an assistant by slug.
+
+        Args:
+            assistant_slug: Slug of the assistant to chat with
+            request: Chat request details
+            headers: Optional additional HTTP headers (e.g., X-* for MCP propagation)
+
+        Returns:
+            Chat response or streaming response
+        """
+        pydantic_schema = None
+        if (
+            request.output_schema is not None
+            and inspect.isclass(request.output_schema)
+            and issubclass(request.output_schema, BaseModel)
+        ):
+            pydantic_schema = deepcopy(request.output_schema)
+            request.output_schema = request.output_schema.model_json_schema()
+
+        response = self._api.post(
+            f"/v1/assistants/slug/{assistant_slug}/model",
+            BaseModelResponse,
+            json_data=request.model_dump(exclude_none=True, by_alias=True),
+            stream=request.stream,
+            extra_headers=headers,
+        )
+        if not request.stream and pydantic_schema:
+            response.generated = pydantic_schema.model_validate(response.generated)
+
+        return response
+
     def upload_file_to_chat(self, file_path: Path):
         """Upload a file to assistant chat and return the response containing file_url."""