codemie-sdk-python 0.1.228__tar.gz → 0.1.230__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (45) hide show
  1. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/PKG-INFO +48 -9
  2. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/README.md +47 -8
  3. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/pyproject.toml +1 -1
  4. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/__init__.py +27 -1
  5. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/client/client.py +7 -0
  6. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/assistant.py +8 -0
  7. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/datasource.py +1 -0
  8. codemie_sdk_python-0.1.230/src/codemie_sdk/models/vendor_guardrail.py +152 -0
  9. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/workflow_execution_payload.py +4 -0
  10. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/assistant.py +40 -0
  11. codemie_sdk_python-0.1.230/src/codemie_sdk/services/vendor_guardrail.py +375 -0
  12. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/workflow.py +10 -1
  13. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/workflow_execution.py +26 -4
  14. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/utils/http.py +13 -2
  15. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/auth/__init__.py +0 -0
  16. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/auth/credentials.py +0 -0
  17. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/client/__init__.py +0 -0
  18. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/exceptions.py +0 -0
  19. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/__init__.py +0 -0
  20. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/common.py +0 -0
  21. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/conversation.py +0 -0
  22. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/file_operation.py +0 -0
  23. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/integration.py +0 -0
  24. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/llm.py +0 -0
  25. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/task.py +0 -0
  26. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/user.py +0 -0
  27. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/vendor_assistant.py +0 -0
  28. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/vendor_knowledgebase.py +0 -0
  29. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/vendor_workflow.py +0 -0
  30. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/workflow.py +0 -0
  31. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/workflow_state.py +0 -0
  32. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/models/workflow_thoughts.py +0 -0
  33. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/conversation.py +0 -0
  34. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/datasource.py +0 -0
  35. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/files.py +0 -0
  36. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/integration.py +0 -0
  37. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/llm.py +0 -0
  38. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/task.py +0 -0
  39. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/user.py +0 -0
  40. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/vendor_assistant.py +0 -0
  41. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/vendor_knowledgebase.py +0 -0
  42. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/vendor_workflow.py +0 -0
  43. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/webhook.py +0 -0
  44. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/services/workflow_execution_state.py +0 -0
  45. {codemie_sdk_python-0.1.228 → codemie_sdk_python-0.1.230}/src/codemie_sdk/utils/__init__.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: codemie-sdk-python
3
- Version: 0.1.228
3
+ Version: 0.1.230
4
4
  Summary: CodeMie SDK for Python
5
5
  Author: Vadym Vlasenko
6
6
  Author-email: vadym_vlasenko@epam.com
@@ -174,16 +174,41 @@ result = client.assistant.delete("assistant-id")
174
174
 
175
175
  #### Advanced Features
176
176
 
177
- 6. **Chat with Assistant**
177
+ 6. **Chat with Assistant (with MCP header propagation)**
178
178
  ```python
179
179
  from codemie_sdk.models.assistant import AssistantChatRequest
180
180
 
181
181
  chat_request = AssistantChatRequest(
182
182
  text="Your message here",
183
183
  stream=False, # Set to True for streaming response
184
- # Additional parameters
184
+ propagate_headers=True, # Enable propagation of X-* headers to MCP servers
185
+ )
186
+ # Pass X-* headers to forward to MCP servers
187
+ response = client.assistant.chat(
188
+ "assistant-id",
189
+ chat_request,
190
+ headers={
191
+ "X-Tenant-ID": "tenant-abc-123",
192
+ "X-User-ID": "user-456",
193
+ "X-Request-ID": "req-123",
194
+ },
195
+ )
196
+ ```
197
+
198
+ 6.a. **Chat with Assistant by slug (with MCP header propagation)**
199
+ ```python
200
+ chat_request = AssistantChatRequest(
201
+ text="Your message here",
202
+ propagate_headers=True,
203
+ )
204
+ response = client.assistants.chat_by_slug(
205
+ "assistant-slug",
206
+ chat_request,
207
+ headers={
208
+ "X-Environment": "production",
209
+ "X-Feature-Flag-Beta": "true",
210
+ },
185
211
  )
186
- response = client.assistant.chat("assistant-id", chat_request)
187
212
  ```
188
213
 
189
214
  7. **Utilize structured outputs with Assistant**
@@ -548,10 +573,18 @@ result = client.workflow.delete("workflow-id")
548
573
 
549
574
  The SDK provides comprehensive workflow execution management through the WorkflowExecutionService:
550
575
 
551
- 1. **Run Workflow**
576
+ 1. **Run Workflow (with MCP header propagation)**
552
577
  ```python
553
- # Simple workflow execution
554
- execution = client.workflow.run("workflow-id", user_input="optional input")
578
+ # Enable propagation in payload and pass X-* headers to forward to MCP servers
579
+ execution = client.workflow.run(
580
+ "workflow-id",
581
+ user_input="optional input",
582
+ propagate_headers=True,
583
+ headers={
584
+ "X-Request-ID": "req-abc-123",
585
+ "X-Source-App": "analytics-ui",
586
+ },
587
+ )
555
588
 
556
589
  # Get execution service for advanced operations
557
590
  execution_service = client.workflow.executions("workflow-id")
@@ -571,8 +604,14 @@ execution = execution_service.get("execution-id")
571
604
  # Abort running execution
572
605
  result = execution_service.abort("execution-id")
573
606
 
574
- # Resume interrupted execution
575
- result = execution_service.resume("execution-id")
607
+ # Resume interrupted execution with header propagation (query param + headers)
608
+ result = execution_service.resume(
609
+ "execution-id",
610
+ propagate_headers=True,
611
+ headers={
612
+ "X-Correlation-ID": "corr-456",
613
+ },
614
+ )
576
615
 
577
616
  # Delete all executions
578
617
  result = execution_service.delete_all()
@@ -161,16 +161,41 @@ result = client.assistant.delete("assistant-id")
161
161
 
162
162
  #### Advanced Features
163
163
 
164
- 6. **Chat with Assistant**
164
+ 6. **Chat with Assistant (with MCP header propagation)**
165
165
  ```python
166
166
  from codemie_sdk.models.assistant import AssistantChatRequest
167
167
 
168
168
  chat_request = AssistantChatRequest(
169
169
  text="Your message here",
170
170
  stream=False, # Set to True for streaming response
171
- # Additional parameters
171
+ propagate_headers=True, # Enable propagation of X-* headers to MCP servers
172
+ )
173
+ # Pass X-* headers to forward to MCP servers
174
+ response = client.assistant.chat(
175
+ "assistant-id",
176
+ chat_request,
177
+ headers={
178
+ "X-Tenant-ID": "tenant-abc-123",
179
+ "X-User-ID": "user-456",
180
+ "X-Request-ID": "req-123",
181
+ },
182
+ )
183
+ ```
184
+
185
+ 6.a. **Chat with Assistant by slug (with MCP header propagation)**
186
+ ```python
187
+ chat_request = AssistantChatRequest(
188
+ text="Your message here",
189
+ propagate_headers=True,
190
+ )
191
+ response = client.assistants.chat_by_slug(
192
+ "assistant-slug",
193
+ chat_request,
194
+ headers={
195
+ "X-Environment": "production",
196
+ "X-Feature-Flag-Beta": "true",
197
+ },
172
198
  )
173
- response = client.assistant.chat("assistant-id", chat_request)
174
199
  ```
175
200
 
176
201
  7. **Utilize structured outputs with Assistant**
@@ -535,10 +560,18 @@ result = client.workflow.delete("workflow-id")
535
560
 
536
561
  The SDK provides comprehensive workflow execution management through the WorkflowExecutionService:
537
562
 
538
- 1. **Run Workflow**
563
+ 1. **Run Workflow (with MCP header propagation)**
539
564
  ```python
540
- # Simple workflow execution
541
- execution = client.workflow.run("workflow-id", user_input="optional input")
565
+ # Enable propagation in payload and pass X-* headers to forward to MCP servers
566
+ execution = client.workflow.run(
567
+ "workflow-id",
568
+ user_input="optional input",
569
+ propagate_headers=True,
570
+ headers={
571
+ "X-Request-ID": "req-abc-123",
572
+ "X-Source-App": "analytics-ui",
573
+ },
574
+ )
542
575
 
543
576
  # Get execution service for advanced operations
544
577
  execution_service = client.workflow.executions("workflow-id")
@@ -558,8 +591,14 @@ execution = execution_service.get("execution-id")
558
591
  # Abort running execution
559
592
  result = execution_service.abort("execution-id")
560
593
 
561
- # Resume interrupted execution
562
- result = execution_service.resume("execution-id")
594
+ # Resume interrupted execution with header propagation (query param + headers)
595
+ result = execution_service.resume(
596
+ "execution-id",
597
+ propagate_headers=True,
598
+ headers={
599
+ "X-Correlation-ID": "corr-456",
600
+ },
601
+ )
563
602
 
564
603
  # Delete all executions
565
604
  result = execution_service.delete_all()
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "codemie-sdk-python"
3
- version = "0.1.228"
3
+ version = "0.1.230"
4
4
  description = "CodeMie SDK for Python"
5
5
  authors = [
6
6
  "Vadym Vlasenko <vadym_vlasenko@epam.com>",
@@ -60,11 +60,25 @@ from .models.vendor_knowledgebase import (
60
60
  VendorKnowledgeBaseInstallResponse,
61
61
  VendorKnowledgeBaseUninstallResponse,
62
62
  )
63
+ from .models.vendor_guardrail import (
64
+ VendorGuardrailSetting,
65
+ VendorGuardrailSettingsResponse,
66
+ VendorGuardrail,
67
+ VendorGuardrailStatus,
68
+ VendorGuardrailsResponse,
69
+ VendorGuardrailVersion,
70
+ VendorGuardrailVersionsResponse,
71
+ VendorGuardrailInstallRequest,
72
+ VendorGuardrailInstallSummary,
73
+ VendorGuardrailInstallResponse,
74
+ VendorGuardrailUninstallResponse,
75
+ )
63
76
  from .services.vendor_assistant import VendorAssistantService
64
77
  from .services.vendor_workflow import VendorWorkflowService
65
78
  from .services.vendor_knowledgebase import VendorKnowledgeBaseService
79
+ from .services.vendor_guardrail import VendorGuardrailService
66
80
 
67
- __version__ = "0.2.11"
81
+ __version__ = "0.2.12"
68
82
  __all__ = [
69
83
  "CodeMieClient",
70
84
  "VendorType",
@@ -106,4 +120,16 @@ __all__ = [
106
120
  "VendorKnowledgeBaseInstallResponse",
107
121
  "VendorKnowledgeBaseUninstallResponse",
108
122
  "VendorKnowledgeBaseService",
123
+ "VendorGuardrailSetting",
124
+ "VendorGuardrailSettingsResponse",
125
+ "VendorGuardrail",
126
+ "VendorGuardrailStatus",
127
+ "VendorGuardrailsResponse",
128
+ "VendorGuardrailVersion",
129
+ "VendorGuardrailVersionsResponse",
130
+ "VendorGuardrailInstallRequest",
131
+ "VendorGuardrailInstallSummary",
132
+ "VendorGuardrailInstallResponse",
133
+ "VendorGuardrailUninstallResponse",
134
+ "VendorGuardrailService",
109
135
  ]
@@ -16,6 +16,7 @@ from ..services.webhook import WebhookService
16
16
  from ..services.vendor_assistant import VendorAssistantService
17
17
  from ..services.vendor_workflow import VendorWorkflowService
18
18
  from ..services.vendor_knowledgebase import VendorKnowledgeBaseService
19
+ from ..services.vendor_guardrail import VendorGuardrailService
19
20
 
20
21
 
21
22
  class CodeMieClient:
@@ -101,6 +102,9 @@ class CodeMieClient:
101
102
  self.vendor_knowledgebases = VendorKnowledgeBaseService(
102
103
  self._api_domain, self._token, verify_ssl=self._verify_ssl
103
104
  )
105
+ self.vendor_guardrails = VendorGuardrailService(
106
+ self._api_domain, self._token, verify_ssl=self._verify_ssl
107
+ )
104
108
 
105
109
  @property
106
110
  def token(self) -> str:
@@ -156,4 +160,7 @@ class CodeMieClient:
156
160
  self.vendor_knowledgebases = VendorKnowledgeBaseService(
157
161
  self._api_domain, self._token, verify_ssl=self._verify_ssl
158
162
  )
163
+ self.vendor_guardrails = VendorGuardrailService(
164
+ self._api_domain, self._token, verify_ssl=self._verify_ssl
165
+ )
159
166
  return self._token
@@ -284,6 +284,14 @@ class AssistantChatRequest(BaseModel):
284
284
  default=None, description="DataSource in conversation history"
285
285
  )
286
286
  stream: bool = Field(default=False, description="Enable streaming response")
287
+ propagate_headers: bool = Field(
288
+ default=False,
289
+ description="Enable propagation of X-* HTTP headers to MCP servers during tool execution",
290
+ )
291
+ custom_metadata: Optional[dict[str, Any]] = Field(
292
+ default=None,
293
+ description="Custom metadata for the AI Assistant",
294
+ )
287
295
  top_k: int = Field(default=10, description="Top K results to consider")
288
296
  system_prompt: str = Field(default="", description="Override system prompt")
289
297
  background_task: bool = Field(default=False, description="Run as background task")
@@ -25,6 +25,7 @@ class DataSourceType(str, Enum):
25
25
  CHUNK_SUMMARY = "chunk-summary"
26
26
  JSON = "knowledge_base_json"
27
27
  BEDROCK = "knowledge_base_bedrock"
28
+ PLATFORM = "platform_marketplace_assistant"
28
29
 
29
30
 
30
31
  class DataSourceStatus(str, Enum):
@@ -0,0 +1,152 @@
1
+ """Models for vendor guardrail settings."""
2
+
3
+ from datetime import datetime
4
+ from enum import Enum
5
+ from typing import Optional, List
6
+
7
+ from pydantic import BaseModel, ConfigDict, Field
8
+
9
+ from .vendor_assistant import PaginationInfo, TokenPagination
10
+
11
+
12
+ class VendorGuardrailSetting(BaseModel):
13
+ """Model representing a vendor guardrail setting."""
14
+
15
+ model_config = ConfigDict(extra="ignore")
16
+
17
+ setting_id: str = Field(..., description="Unique identifier for the setting")
18
+ setting_name: str = Field(..., description="Name of the setting")
19
+ project: str = Field(..., description="Project associated with the setting")
20
+ entities: List[str] = Field(
21
+ default_factory=list, description="List of entities associated with the setting"
22
+ )
23
+ invalid: Optional[bool] = Field(None, description="Whether the setting is invalid")
24
+ error: Optional[str] = Field(
25
+ None, description="Error message if the setting is invalid"
26
+ )
27
+
28
+
29
+ class VendorGuardrailSettingsResponse(BaseModel):
30
+ """Response model for vendor guardrail settings list."""
31
+
32
+ model_config = ConfigDict(extra="ignore")
33
+
34
+ data: List[VendorGuardrailSetting] = Field(
35
+ ..., description="List of vendor guardrail settings"
36
+ )
37
+ pagination: PaginationInfo = Field(..., description="Pagination information")
38
+
39
+
40
+ class VendorGuardrailStatus(str, Enum):
41
+ """Status of vendor guardrail."""
42
+
43
+ PREPARED = "PREPARED"
44
+ NOT_PREPARED = "NOT_PREPARED"
45
+
46
+
47
+ class VendorGuardrail(BaseModel):
48
+ """Model representing a vendor guardrail."""
49
+
50
+ model_config = ConfigDict(extra="ignore")
51
+
52
+ id: str = Field(..., description="Unique identifier for the guardrail")
53
+ name: str = Field(..., description="Name of the guardrail")
54
+ status: VendorGuardrailStatus = Field(..., description="Status of the guardrail")
55
+ description: Optional[str] = Field(None, description="Description of the guardrail")
56
+ version: str = Field(..., description="Version of the guardrail")
57
+ createdAt: datetime = Field(
58
+ ..., description="Creation timestamp", alias="createdAt"
59
+ )
60
+ updatedAt: datetime = Field(
61
+ ..., description="Last update timestamp", alias="updatedAt"
62
+ )
63
+
64
+
65
+ class VendorGuardrailsResponse(BaseModel):
66
+ """Response model for vendor guardrails list."""
67
+
68
+ model_config = ConfigDict(extra="ignore")
69
+
70
+ data: List[VendorGuardrail] = Field(..., description="List of vendor guardrails")
71
+ pagination: TokenPagination = Field(
72
+ ..., description="Token-based pagination information"
73
+ )
74
+
75
+
76
+ class VendorGuardrailVersion(BaseModel):
77
+ """Model representing a vendor guardrail version."""
78
+
79
+ model_config = ConfigDict(extra="ignore")
80
+
81
+ id: str = Field(..., description="Unique identifier for the guardrail")
82
+ version: str = Field(..., description="Version of the guardrail")
83
+ name: str = Field(..., description="Name of the guardrail")
84
+ status: VendorGuardrailStatus = Field(..., description="Status of the version")
85
+ description: Optional[str] = Field(None, description="Description of the version")
86
+ blockedInputMessaging: Optional[str] = Field(
87
+ None,
88
+ description="Message to display when input is blocked by guardrail",
89
+ alias="blockedInputMessaging",
90
+ )
91
+ blockedOutputsMessaging: Optional[str] = Field(
92
+ None,
93
+ description="Message to display when output is blocked by guardrail",
94
+ alias="blockedOutputsMessaging",
95
+ )
96
+ createdAt: datetime = Field(
97
+ ..., description="Creation timestamp", alias="createdAt"
98
+ )
99
+ updatedAt: datetime = Field(
100
+ ..., description="Last update timestamp", alias="updatedAt"
101
+ )
102
+
103
+
104
+ class VendorGuardrailVersionsResponse(BaseModel):
105
+ """Response model for vendor guardrail versions list."""
106
+
107
+ model_config = ConfigDict(extra="ignore")
108
+
109
+ data: List[VendorGuardrailVersion] = Field(
110
+ ..., description="List of vendor guardrail versions"
111
+ )
112
+ pagination: TokenPagination = Field(
113
+ ..., description="Token-based pagination information"
114
+ )
115
+
116
+
117
+ class VendorGuardrailInstallRequest(BaseModel):
118
+ """Model for a single guardrail installation request."""
119
+
120
+ model_config = ConfigDict(extra="ignore")
121
+
122
+ id: str = Field(..., description="Guardrail ID to install")
123
+ version: str = Field(..., description="Version to use for the guardrail")
124
+ setting_id: str = Field(..., description="Vendor setting ID")
125
+
126
+
127
+ class VendorGuardrailInstallSummary(BaseModel):
128
+ """Model for guardrail installation summary."""
129
+
130
+ model_config = ConfigDict(extra="ignore")
131
+
132
+ guardrailId: str = Field(..., description="Installed guardrail ID")
133
+ version: str = Field(..., description="Version used for installation")
134
+ aiRunId: str = Field(..., description="AI run ID for the installation")
135
+
136
+
137
+ class VendorGuardrailInstallResponse(BaseModel):
138
+ """Response model for guardrail installation."""
139
+
140
+ model_config = ConfigDict(extra="ignore")
141
+
142
+ summary: List[VendorGuardrailInstallSummary] = Field(
143
+ ..., description="List of installation summaries"
144
+ )
145
+
146
+
147
+ class VendorGuardrailUninstallResponse(BaseModel):
148
+ """Response model for guardrail uninstallation."""
149
+
150
+ model_config = ConfigDict(extra="ignore")
151
+
152
+ success: bool = Field(..., description="Whether the uninstallation was successful")
@@ -15,3 +15,7 @@ class WorkflowExecutionCreateRequest(BaseModel):
15
15
  file_name: Optional[str] = Field(
16
16
  None, description="File name associated with the workflow execution"
17
17
  )
18
+ propagate_headers: bool = Field(
19
+ default=False,
20
+ description="Enable propagation of X-* HTTP headers to MCP servers during tool execution",
21
+ )
@@ -309,12 +309,14 @@ class AssistantService:
309
309
  self,
310
310
  assistant_id: str,
311
311
  request: AssistantChatRequest,
312
+ headers: Optional[Dict[str, str]] = None,
312
313
  ) -> Union[requests.Response, BaseModelResponse]:
313
314
  """Send a chat request to an assistant.
314
315
 
315
316
  Args:
316
317
  assistant_id: ID of the assistant to chat with
317
318
  request: Chat request details
319
+ headers: Optional additional HTTP headers (e.g., X-* for MCP propagation)
318
320
 
319
321
  Returns:
320
322
  Chat response or streaming response
@@ -333,6 +335,7 @@ class AssistantService:
333
335
  BaseModelResponse,
334
336
  json_data=request.model_dump(exclude_none=True, by_alias=True),
335
337
  stream=request.stream,
338
+ extra_headers=headers,
336
339
  )
337
340
  if not request.stream and pydantic_schema:
338
341
  # we do conversion to the BaseModel here because self._parse_response don't see actual request model,
@@ -380,6 +383,43 @@ class AssistantService:
380
383
 
381
384
  return response
382
385
 
386
+ def chat_by_slug(
387
+ self,
388
+ assistant_slug: str,
389
+ request: AssistantChatRequest,
390
+ headers: Optional[Dict[str, str]] = None,
391
+ ) -> Union[requests.Response, BaseModelResponse]:
392
+ """Send a chat request to an assistant by slug.
393
+
394
+ Args:
395
+ assistant_slug: Slug of the assistant to chat with
396
+ request: Chat request details
397
+ headers: Optional additional HTTP headers (e.g., X-* for MCP propagation)
398
+
399
+ Returns:
400
+ Chat response or streaming response
401
+ """
402
+ pydantic_schema = None
403
+ if (
404
+ request.output_schema is not None
405
+ and inspect.isclass(request.output_schema)
406
+ and issubclass(request.output_schema, BaseModel)
407
+ ):
408
+ pydantic_schema = deepcopy(request.output_schema)
409
+ request.output_schema = request.output_schema.model_json_schema()
410
+
411
+ response = self._api.post(
412
+ f"/v1/assistants/slug/{assistant_slug}/model",
413
+ BaseModelResponse,
414
+ json_data=request.model_dump(exclude_none=True, by_alias=True),
415
+ stream=request.stream,
416
+ extra_headers=headers,
417
+ )
418
+ if not request.stream and pydantic_schema:
419
+ response.generated = pydantic_schema.model_validate(response.generated)
420
+
421
+ return response
422
+
383
423
  def upload_file_to_chat(self, file_path: Path):
384
424
  """Upload a file to assistant chat and return the response containing file_url."""
385
425
 
@@ -0,0 +1,375 @@
1
+ """Vendor guardrail service implementation for managing cloud vendor guardrail settings."""
2
+
3
+ from typing import Union, Optional, List
4
+
5
+ from ..models.vendor_assistant import VendorType
6
+ from ..models.vendor_guardrail import (
7
+ VendorGuardrailSettingsResponse,
8
+ VendorGuardrailsResponse,
9
+ VendorGuardrail,
10
+ VendorGuardrailVersion,
11
+ VendorGuardrailVersionsResponse,
12
+ VendorGuardrailInstallRequest,
13
+ VendorGuardrailInstallResponse,
14
+ VendorGuardrailUninstallResponse,
15
+ )
16
+ from ..utils import ApiRequestHandler
17
+
18
+
19
+ class VendorGuardrailService:
20
+ """Service for managing cloud vendor guardrail settings (AWS, Azure, GCP)."""
21
+
22
+ def __init__(self, api_domain: str, token: str, verify_ssl: bool = True):
23
+ """Initialize the vendor guardrail service.
24
+
25
+ Args:
26
+ api_domain: Base URL for the CodeMie API
27
+ token: Authentication token
28
+ verify_ssl: Whether to verify SSL certificates
29
+ """
30
+ self._api = ApiRequestHandler(api_domain, token, verify_ssl)
31
+
32
+ def get_guardrail_settings(
33
+ self,
34
+ vendor: Union[VendorType, str],
35
+ page: int = 0,
36
+ per_page: int = 10,
37
+ ) -> VendorGuardrailSettingsResponse:
38
+ """Get guardrail settings for a specific cloud vendor.
39
+
40
+ Args:
41
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
42
+ page: Page number for pagination (0-based)
43
+ per_page: Number of items per page
44
+
45
+ Returns:
46
+ VendorGuardrailSettingsResponse containing list of settings and pagination info
47
+
48
+ Example:
49
+ >>> # Using enum
50
+ >>> settings = client.vendor_guardrails.get_guardrail_settings(VendorType.AWS, page=0, per_page=10)
51
+ >>> # Using string
52
+ >>> settings = client.vendor_guardrails.get_guardrail_settings("aws", page=0, per_page=10)
53
+ >>> # Access settings data
54
+ >>> for setting in settings.data:
55
+ ... print(f"Setting: {setting.setting_name}, Project: {setting.project}")
56
+ ... if setting.invalid:
57
+ ... print(f"Error: {setting.error}")
58
+ """
59
+ # Convert enum to string value if needed
60
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
61
+
62
+ params = {
63
+ "page": page,
64
+ "per_page": per_page,
65
+ }
66
+
67
+ return self._api.get(
68
+ f"/v1/vendors/{vendor_str}/guardrails/settings",
69
+ VendorGuardrailSettingsResponse,
70
+ params=params,
71
+ wrap_response=False,
72
+ )
73
+
74
+ def get_guardrails(
75
+ self,
76
+ vendor: Union[VendorType, str],
77
+ setting_id: str,
78
+ per_page: int = 10,
79
+ next_token: Optional[str] = None,
80
+ ) -> VendorGuardrailsResponse:
81
+ """Get guardrails for a specific vendor setting.
82
+
83
+ Args:
84
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
85
+ setting_id: ID of the vendor setting to retrieve guardrails for
86
+ per_page: Number of items per page
87
+ next_token: Token for pagination (optional, for retrieving next page)
88
+
89
+ Returns:
90
+ VendorGuardrailsResponse containing list of guardrails and pagination token
91
+
92
+ Example:
93
+ >>> # Get first page
94
+ >>> guardrails = client.vendor_guardrails.get_guardrails(
95
+ ... vendor=VendorType.AWS,
96
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c",
97
+ ... per_page=8
98
+ ... )
99
+ >>> # Access guardrail data
100
+ >>> for guardrail in guardrails.data:
101
+ ... print(f"Name: {guardrail.name}, Status: {guardrail.status}")
102
+ ... print(f"Version: {guardrail.version}, Description: {guardrail.description}")
103
+ >>> # Get next page if available
104
+ >>> if guardrails.pagination.next_token:
105
+ ... next_page = client.vendor_guardrails.get_guardrails(
106
+ ... vendor=VendorType.AWS,
107
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c",
108
+ ... per_page=8,
109
+ ... next_token=guardrails.pagination.next_token
110
+ ... )
111
+ """
112
+ # Convert enum to string value if needed
113
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
114
+
115
+ params = {
116
+ "setting_id": setting_id,
117
+ "per_page": per_page,
118
+ }
119
+
120
+ if next_token:
121
+ params["next_token"] = next_token
122
+
123
+ return self._api.get(
124
+ f"/v1/vendors/{vendor_str}/guardrails",
125
+ VendorGuardrailsResponse,
126
+ params=params,
127
+ wrap_response=False,
128
+ )
129
+
130
+ def get_guardrail(
131
+ self,
132
+ vendor: Union[VendorType, str],
133
+ guardrail_id: str,
134
+ setting_id: str,
135
+ ) -> VendorGuardrail:
136
+ """Get a specific guardrail by ID for a vendor setting.
137
+
138
+ Args:
139
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
140
+ guardrail_id: ID of the guardrail to retrieve
141
+ setting_id: ID of the vendor setting
142
+
143
+ Returns:
144
+ VendorGuardrail containing guardrail details
145
+
146
+ Example:
147
+ >>> guardrail = client.vendor_guardrails.get_guardrail(
148
+ ... vendor=VendorType.AWS,
149
+ ... guardrail_id="lss9vxro9oxg",
150
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c"
151
+ ... )
152
+ >>> print(f"Name: {guardrail.name}, Status: {guardrail.status}")
153
+ >>> print(f"Version: {guardrail.version}, Description: {guardrail.description}")
154
+ """
155
+ # Convert enum to string value if needed
156
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
157
+
158
+ params = {
159
+ "setting_id": setting_id,
160
+ }
161
+
162
+ return self._api.get(
163
+ f"/v1/vendors/{vendor_str}/guardrails/{guardrail_id}",
164
+ VendorGuardrail,
165
+ params=params,
166
+ wrap_response=False,
167
+ )
168
+
169
+ def get_guardrail_version(
170
+ self,
171
+ vendor: Union[VendorType, str],
172
+ guardrail_id: str,
173
+ version: str,
174
+ setting_id: str,
175
+ ) -> VendorGuardrailVersion:
176
+ """Get a specific version of a vendor guardrail with detailed information.
177
+
178
+ Args:
179
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
180
+ guardrail_id: ID of the guardrail to retrieve
181
+ version: Version number to retrieve (e.g., "1", "2", "DRAFT")
182
+ setting_id: ID of the vendor setting
183
+
184
+ Returns:
185
+ VendorGuardrailVersion containing detailed version information including
186
+ blocked messaging settings and timestamps
187
+
188
+ Example:
189
+ >>> version_details = client.vendor_guardrails.get_guardrail_version(
190
+ ... vendor=VendorType.AWS,
191
+ ... guardrail_id="lss9vxro9oxg",
192
+ ... version="1",
193
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c"
194
+ ... )
195
+ >>> print(f"Name: {version_details.name}")
196
+ >>> print(f"Version: {version_details.version}")
197
+ >>> print(f"Status: {version_details.status}")
198
+ >>> print(f"Blocked Input Message: {version_details.blockedInputMessaging}")
199
+ >>> print(f"Blocked Output Message: {version_details.blockedOutputsMessaging}")
200
+ """
201
+ # Convert enum to string value if needed
202
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
203
+
204
+ params = {
205
+ "setting_id": setting_id,
206
+ }
207
+
208
+ return self._api.get(
209
+ f"/v1/vendors/{vendor_str}/guardrails/{guardrail_id}/{version}",
210
+ VendorGuardrailVersion,
211
+ params=params,
212
+ wrap_response=False,
213
+ )
214
+
215
+ def get_guardrail_versions(
216
+ self,
217
+ vendor: Union[VendorType, str],
218
+ guardrail_id: str,
219
+ setting_id: str,
220
+ per_page: int = 10,
221
+ next_token: Optional[str] = None,
222
+ ) -> VendorGuardrailVersionsResponse:
223
+ """Get versions for a specific vendor guardrail.
224
+
225
+ Args:
226
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
227
+ guardrail_id: ID of the guardrail to retrieve versions for
228
+ setting_id: ID of the vendor setting
229
+ per_page: Number of items per page
230
+ next_token: Token for pagination (optional, for retrieving next page)
231
+
232
+ Returns:
233
+ VendorGuardrailVersionsResponse containing list of versions and pagination token
234
+
235
+ Example:
236
+ >>> # Get first page of versions
237
+ >>> versions = client.vendor_guardrails.get_guardrail_versions(
238
+ ... vendor=VendorType.AWS,
239
+ ... guardrail_id="lss9vxro9oxg",
240
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c",
241
+ ... per_page=5
242
+ ... )
243
+ >>> for version in versions.data:
244
+ ... print(f"{version.name} (v{version.version}): {version.status}")
245
+ >>> # Get next page if available
246
+ >>> if versions.pagination.next_token:
247
+ ... next_page = client.vendor_guardrails.get_guardrail_versions(
248
+ ... vendor=VendorType.AWS,
249
+ ... guardrail_id="lss9vxro9oxg",
250
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c",
251
+ ... per_page=5,
252
+ ... next_token=versions.pagination.next_token
253
+ ... )
254
+ """
255
+ # Convert enum to string value if needed
256
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
257
+
258
+ params = {
259
+ "setting_id": setting_id,
260
+ "per_page": per_page,
261
+ }
262
+
263
+ if next_token:
264
+ params["next_token"] = next_token
265
+
266
+ return self._api.get(
267
+ f"/v1/vendors/{vendor_str}/guardrails/{guardrail_id}/versions",
268
+ VendorGuardrailVersionsResponse,
269
+ params=params,
270
+ wrap_response=False,
271
+ )
272
+
273
+ def install_guardrails(
274
+ self,
275
+ vendor: Union[VendorType, str],
276
+ guardrails: List[VendorGuardrailInstallRequest],
277
+ ) -> VendorGuardrailInstallResponse:
278
+ """Install/activate vendor guardrails.
279
+
280
+ Args:
281
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
282
+ guardrails: List of guardrail installation requests with guardrail ID, version, and setting ID
283
+
284
+ Returns:
285
+ VendorGuardrailInstallResponse containing installation summary with AI run IDs
286
+
287
+ Example:
288
+ >>> from codemie_sdk import VendorGuardrailInstallRequest
289
+ >>> # Install single guardrail
290
+ >>> install_request = VendorGuardrailInstallRequest(
291
+ ... id="lss9vxro9oxg",
292
+ ... version="1",
293
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c"
294
+ ... )
295
+ >>> response = client.vendor_guardrails.install_guardrails(
296
+ ... vendor=VendorType.AWS,
297
+ ... guardrails=[install_request]
298
+ ... )
299
+ >>> for item in response.summary:
300
+ ... print(f"Installed {item.guardrailId} version {item.version} with run ID: {item.aiRunId}")
301
+ >>>
302
+ >>> # Install multiple guardrails
303
+ >>> requests = [
304
+ ... VendorGuardrailInstallRequest(
305
+ ... id="GUARDRAIL_ID_1",
306
+ ... version="1",
307
+ ... setting_id="SETTING_ID"
308
+ ... ),
309
+ ... VendorGuardrailInstallRequest(
310
+ ... id="GUARDRAIL_ID_2",
311
+ ... version="2",
312
+ ... setting_id="SETTING_ID"
313
+ ... )
314
+ ... ]
315
+ >>> response = client.vendor_guardrails.install_guardrails(
316
+ ... vendor=VendorType.AWS,
317
+ ... guardrails=requests
318
+ ... )
319
+ """
320
+ # Convert enum to string value if needed
321
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
322
+
323
+ # Convert list of Pydantic models to list of dicts
324
+ payload = [guardrail.model_dump(by_alias=True) for guardrail in guardrails]
325
+
326
+ return self._api.post(
327
+ f"/v1/vendors/{vendor_str}/guardrails",
328
+ VendorGuardrailInstallResponse,
329
+ json_data=payload,
330
+ wrap_response=False,
331
+ )
332
+
333
+ def uninstall_guardrail(
334
+ self,
335
+ vendor: Union[VendorType, str],
336
+ ai_run_id: str,
337
+ ) -> VendorGuardrailUninstallResponse:
338
+ """Uninstall/deactivate a vendor guardrail.
339
+
340
+ Args:
341
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
342
+ ai_run_id: AI run ID returned from the install operation
343
+
344
+ Returns:
345
+ VendorGuardrailUninstallResponse with success status
346
+
347
+ Example:
348
+ >>> # First, install a guardrail
349
+ >>> install_request = VendorGuardrailInstallRequest(
350
+ ... id="lss9vxro9oxg",
351
+ ... version="1",
352
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c"
353
+ ... )
354
+ >>> install_response = client.vendor_guardrails.install_guardrails(
355
+ ... vendor=VendorType.AWS,
356
+ ... guardrails=[install_request]
357
+ ... )
358
+ >>> ai_run_id = install_response.summary[0].aiRunId
359
+ >>>
360
+ >>> # Later, uninstall the guardrail using the AI run ID
361
+ >>> response = client.vendor_guardrails.uninstall_guardrail(
362
+ ... vendor=VendorType.AWS,
363
+ ... ai_run_id=ai_run_id
364
+ ... )
365
+ >>> if response.success:
366
+ ... print("Guardrail successfully uninstalled!")
367
+ """
368
+ # Convert enum to string value if needed
369
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
370
+
371
+ return self._api.delete(
372
+ f"/v1/vendors/{vendor_str}/guardrails/{ai_run_id}",
373
+ VendorGuardrailUninstallResponse,
374
+ wrap_response=False,
375
+ )
@@ -135,6 +135,8 @@ class WorkflowService:
135
135
  workflow_id: str,
136
136
  user_input: Optional[str] = None,
137
137
  file_name: Optional[str] = None,
138
+ propagate_headers: bool = False,
139
+ headers: Optional[dict[str, str]] = None,
138
140
  ) -> dict:
139
141
  """Run a workflow with optional input parameters.
140
142
 
@@ -142,11 +144,18 @@ class WorkflowService:
142
144
  workflow_id: ID of the workflow to run
143
145
  user_input: Optional user input for the workflow execution
144
146
  file_name: Optional file name for the workflow execution
147
+ propagate_headers: Enable propagation of X-* HTTP headers to MCP servers
148
+ headers: Optional additional HTTP headers (e.g., X-* for MCP propagation)
145
149
 
146
150
  Returns:
147
151
  dict: Created workflow execution details
148
152
  """
149
- return self.executions(workflow_id).create(user_input, file_name)
153
+ return self.executions(workflow_id).create(
154
+ user_input=user_input,
155
+ file_name=file_name,
156
+ propagate_headers=propagate_headers,
157
+ headers=headers,
158
+ )
150
159
 
151
160
  def executions(self, workflow_id: str) -> WorkflowExecutionService:
152
161
  """Get workflow execution service for the specified workflow.
@@ -46,24 +46,33 @@ class WorkflowExecutionService:
46
46
  )
47
47
 
48
48
  def create(
49
- self, user_input: Optional[str] = None, file_name: Optional[str] = None
49
+ self,
50
+ user_input: Optional[str] = None,
51
+ file_name: Optional[str] = None,
52
+ propagate_headers: bool = False,
53
+ headers: Optional[dict[str, str]] = None,
50
54
  ) -> dict:
51
55
  """Create a new workflow execution.
52
56
 
53
57
  Args:
54
58
  user_input: Optional input data for the workflow execution.
55
59
  file_name: Optional file name associated with the workflow execution.
60
+ propagate_headers: Enable propagation of X-* HTTP headers to MCP servers.
61
+ headers: Optional additional HTTP headers (e.g., X-* for MCP propagation)
56
62
 
57
63
  Returns:
58
64
  dict: Created workflow execution details
59
65
  """
60
66
  payload = WorkflowExecutionCreateRequest(
61
- user_input=user_input, file_name=file_name
67
+ user_input=user_input,
68
+ file_name=file_name,
69
+ propagate_headers=propagate_headers,
62
70
  )
63
71
  return self._api.post(
64
72
  f"/v1/workflows/{self._workflow_id}/executions",
65
73
  dict,
66
74
  json_data=payload.model_dump(),
75
+ extra_headers=headers,
67
76
  )
68
77
 
69
78
  def get(self, execution_id: str) -> WorkflowExecution:
@@ -108,17 +117,30 @@ class WorkflowExecutionService:
108
117
  f"/v1/workflows/{self._workflow_id}/executions/{execution_id}/abort", dict
109
118
  )
110
119
 
111
- def resume(self, execution_id: str) -> dict:
120
+ def resume(
121
+ self,
122
+ execution_id: str,
123
+ propagate_headers: bool = False,
124
+ headers: Optional[dict[str, str]] = None,
125
+ ) -> dict:
112
126
  """Resume an interrupted workflow execution.
113
127
 
114
128
  Args:
115
129
  execution_id: ID of the execution to resume
130
+ propagate_headers: Enable propagation of X-* HTTP headers to MCP servers.
131
+ headers: Optional additional HTTP headers (e.g., X-* for MCP propagation)
116
132
 
117
133
  Returns:
118
134
  dict: Updated workflow execution details
119
135
  """
136
+ params = {"propagate_headers": propagate_headers}
137
+ # Empty body per API; passing empty dict to satisfy typing
120
138
  return self._api.put(
121
- f"/v1/workflows/{self._workflow_id}/executions/{execution_id}/resume", dict
139
+ f"/v1/workflows/{self._workflow_id}/executions/{execution_id}/resume",
140
+ dict,
141
+ json_data={},
142
+ params=params,
143
+ extra_headers=headers,
122
144
  )
123
145
 
124
146
  def get_thoughts(
@@ -161,6 +161,7 @@ class ApiRequestHandler:
161
161
  stream: bool = False,
162
162
  wrap_response: bool = True,
163
163
  raise_on_error: bool = True,
164
+ extra_headers: Optional[Dict[str, str]] = None,
164
165
  ) -> Union[T, requests.Response]:
165
166
  """Makes a POST request and parses the response.
166
167
 
@@ -171,6 +172,7 @@ class ApiRequestHandler:
171
172
  stream: Whether to return streaming response
172
173
  wrap_response: Whether response is wrapped in 'data' field
173
174
  raise_on_error: Whether to raise exception on HTTP error status codes
175
+ extra_headers: Optional additional HTTP headers to include (e.g., X-* for MCP propagation)
174
176
 
175
177
  Returns:
176
178
  Parsed response object/list or streaming response
@@ -178,9 +180,13 @@ class ApiRequestHandler:
178
180
  if json_data:
179
181
  logger.debug(f"Request body: {json_data}")
180
182
 
183
+ headers = self._get_headers()
184
+ if extra_headers:
185
+ headers.update(extra_headers)
186
+
181
187
  response = requests.post(
182
188
  url=f"{self._base_url}{endpoint}",
183
- headers=self._get_headers(),
189
+ headers=headers,
184
190
  json=json_data,
185
191
  verify=self._verify_ssl,
186
192
  stream=stream,
@@ -242,6 +248,7 @@ class ApiRequestHandler:
242
248
  json_data: Dict[str, Any],
243
249
  params: Optional[Dict[str, Any]] = None,
244
250
  wrap_response: bool = True,
251
+ extra_headers: Optional[Dict[str, str]] = None,
245
252
  ) -> T:
246
253
  """Makes a PUT request and parses the response.
247
254
 
@@ -251,14 +258,18 @@ class ApiRequestHandler:
251
258
  json_data: JSON request body
252
259
  params: Query parameters
253
260
  wrap_response: Whether response is wrapped in 'data' field
261
+ extra_headers: Optional additional HTTP headers to include (e.g., X-* for MCP propagation)
254
262
 
255
263
  Returns:
256
264
  Parsed response object or list of objects
257
265
  """
258
266
  logger.debug(f"Request body: {json_data}")
267
+ headers = self._get_headers()
268
+ if extra_headers:
269
+ headers.update(extra_headers)
259
270
  response = requests.put(
260
271
  url=f"{self._base_url}{endpoint}",
261
- headers=self._get_headers(),
272
+ headers=headers,
262
273
  json=json_data,
263
274
  params=params,
264
275
  verify=self._verify_ssl,