codemie-sdk-python 0.1.92__py3-none-any.whl → 0.1.258__py3-none-any.whl

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.

Potentially problematic release.


This version of codemie-sdk-python might be problematic. Click here for more details.

Files changed (33) hide show
  1. codemie_sdk/__init__.py +114 -2
  2. codemie_sdk/auth/credentials.py +5 -4
  3. codemie_sdk/client/client.py +66 -5
  4. codemie_sdk/models/assistant.py +111 -8
  5. codemie_sdk/models/conversation.py +169 -0
  6. codemie_sdk/models/datasource.py +80 -1
  7. codemie_sdk/models/file_operation.py +25 -0
  8. codemie_sdk/models/integration.py +3 -1
  9. codemie_sdk/models/vendor_assistant.py +187 -0
  10. codemie_sdk/models/vendor_guardrail.py +152 -0
  11. codemie_sdk/models/vendor_knowledgebase.py +151 -0
  12. codemie_sdk/models/vendor_workflow.py +145 -0
  13. codemie_sdk/models/workflow.py +1 -1
  14. codemie_sdk/models/workflow_execution_payload.py +21 -0
  15. codemie_sdk/models/workflow_state.py +6 -3
  16. codemie_sdk/models/workflow_thoughts.py +26 -0
  17. codemie_sdk/services/assistant.py +220 -1
  18. codemie_sdk/services/conversation.py +90 -0
  19. codemie_sdk/services/datasource.py +67 -0
  20. codemie_sdk/services/files.py +82 -0
  21. codemie_sdk/services/vendor_assistant.py +364 -0
  22. codemie_sdk/services/vendor_guardrail.py +375 -0
  23. codemie_sdk/services/vendor_knowledgebase.py +270 -0
  24. codemie_sdk/services/vendor_workflow.py +330 -0
  25. codemie_sdk/services/webhook.py +41 -0
  26. codemie_sdk/services/workflow.py +26 -2
  27. codemie_sdk/services/workflow_execution.py +54 -6
  28. codemie_sdk/utils/http.py +43 -16
  29. codemie_sdk_python-0.1.258.dist-info/METADATA +1404 -0
  30. codemie_sdk_python-0.1.258.dist-info/RECORD +45 -0
  31. codemie_sdk_python-0.1.92.dist-info/METADATA +0 -892
  32. codemie_sdk_python-0.1.92.dist-info/RECORD +0 -30
  33. {codemie_sdk_python-0.1.92.dist-info → codemie_sdk_python-0.1.258.dist-info}/WHEEL +0 -0
@@ -0,0 +1,330 @@
1
+ """Vendor workflow service implementation for managing cloud vendor workflow settings."""
2
+
3
+ from typing import Union, Optional, List
4
+
5
+ from ..models.vendor_assistant import VendorType
6
+ from ..models.vendor_workflow import (
7
+ VendorWorkflowSettingsResponse,
8
+ VendorWorkflowsResponse,
9
+ VendorWorkflow,
10
+ VendorWorkflowAliasesResponse,
11
+ VendorWorkflowInstallRequest,
12
+ VendorWorkflowInstallResponse,
13
+ VendorWorkflowUninstallResponse,
14
+ )
15
+ from ..utils import ApiRequestHandler
16
+
17
+
18
+ class VendorWorkflowService:
19
+ """Service for managing cloud vendor workflow settings (AWS, Azure, GCP)."""
20
+
21
+ def __init__(self, api_domain: str, token: str, verify_ssl: bool = True):
22
+ """Initialize the vendor workflow service.
23
+
24
+ Args:
25
+ api_domain: Base URL for the CodeMie API
26
+ token: Authentication token
27
+ verify_ssl: Whether to verify SSL certificates
28
+ """
29
+ self._api = ApiRequestHandler(api_domain, token, verify_ssl)
30
+
31
+ def get_workflow_settings(
32
+ self,
33
+ vendor: Union[VendorType, str],
34
+ page: int = 0,
35
+ per_page: int = 10,
36
+ ) -> VendorWorkflowSettingsResponse:
37
+ """Get workflow settings for a specific cloud vendor.
38
+
39
+ Args:
40
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
41
+ page: Page number for pagination (0-based)
42
+ per_page: Number of items per page
43
+
44
+ Returns:
45
+ VendorWorkflowSettingsResponse containing list of settings and pagination info
46
+
47
+ Example:
48
+ >>> # Using enum
49
+ >>> settings = client.vendor_workflows.get_workflow_settings(VendorType.AWS, page=0, per_page=10)
50
+ >>> # Using string
51
+ >>> settings = client.vendor_workflows.get_workflow_settings("aws", page=0, per_page=10)
52
+ >>> # Access settings data
53
+ >>> for setting in settings.data:
54
+ ... print(f"Setting: {setting.setting_name}, Project: {setting.project}")
55
+ ... if setting.invalid:
56
+ ... print(f"Error: {setting.error}")
57
+ """
58
+ # Convert enum to string value if needed
59
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
60
+
61
+ params = {
62
+ "page": page,
63
+ "per_page": per_page,
64
+ }
65
+
66
+ return self._api.get(
67
+ f"/v1/vendors/{vendor_str}/workflows/settings",
68
+ VendorWorkflowSettingsResponse,
69
+ params=params,
70
+ wrap_response=False,
71
+ )
72
+
73
+ def get_workflows(
74
+ self,
75
+ vendor: Union[VendorType, str],
76
+ setting_id: str,
77
+ per_page: int = 10,
78
+ next_token: Optional[str] = None,
79
+ ) -> VendorWorkflowsResponse:
80
+ """Get workflows for a specific vendor setting.
81
+
82
+ Args:
83
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
84
+ setting_id: ID of the vendor setting to retrieve workflows for
85
+ per_page: Number of items per page
86
+ next_token: Token for pagination (optional, for retrieving next page)
87
+
88
+ Returns:
89
+ VendorWorkflowsResponse containing list of workflows and pagination token
90
+
91
+ Example:
92
+ >>> # Get first page
93
+ >>> workflows = client.vendor_workflows.get_workflows(
94
+ ... vendor=VendorType.AWS,
95
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c",
96
+ ... per_page=8
97
+ ... )
98
+ >>> # Access workflow data
99
+ >>> for workflow in workflows.data:
100
+ ... print(f"Name: {workflow.name}, Status: {workflow.status}")
101
+ ... print(f"Version: {workflow.version}, Description: {workflow.description}")
102
+ >>> # Get next page if available
103
+ >>> if workflows.pagination.next_token:
104
+ ... next_page = client.vendor_workflows.get_workflows(
105
+ ... vendor=VendorType.AWS,
106
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c",
107
+ ... per_page=8,
108
+ ... next_token=workflows.pagination.next_token
109
+ ... )
110
+ """
111
+ # Convert enum to string value if needed
112
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
113
+
114
+ params = {
115
+ "setting_id": setting_id,
116
+ "per_page": per_page,
117
+ }
118
+
119
+ if next_token:
120
+ params["next_token"] = next_token
121
+
122
+ return self._api.get(
123
+ f"/v1/vendors/{vendor_str}/workflows",
124
+ VendorWorkflowsResponse,
125
+ params=params,
126
+ wrap_response=False,
127
+ )
128
+
129
+ def get_workflow(
130
+ self,
131
+ vendor: Union[VendorType, str],
132
+ workflow_id: str,
133
+ setting_id: str,
134
+ ) -> VendorWorkflow:
135
+ """Get a specific workflow by ID for a vendor setting.
136
+
137
+ Args:
138
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
139
+ workflow_id: ID of the workflow to retrieve
140
+ setting_id: ID of the vendor setting
141
+
142
+ Returns:
143
+ VendorWorkflow containing workflow details
144
+
145
+ Example:
146
+ >>> workflow = client.vendor_workflows.get_workflow(
147
+ ... vendor=VendorType.AWS,
148
+ ... workflow_id="9HXLQ7J9YP",
149
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c"
150
+ ... )
151
+ >>> print(f"Name: {workflow.name}, Status: {workflow.status}")
152
+ >>> print(f"Version: {workflow.version}, Description: {workflow.description}")
153
+ """
154
+ # Convert enum to string value if needed
155
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
156
+
157
+ params = {
158
+ "setting_id": setting_id,
159
+ }
160
+
161
+ return self._api.get(
162
+ f"/v1/vendors/{vendor_str}/workflows/{workflow_id}",
163
+ VendorWorkflow,
164
+ params=params,
165
+ wrap_response=False,
166
+ )
167
+
168
+ def get_workflow_aliases(
169
+ self,
170
+ vendor: Union[VendorType, str],
171
+ workflow_id: str,
172
+ setting_id: str,
173
+ per_page: int = 10,
174
+ next_token: Optional[str] = None,
175
+ ) -> VendorWorkflowAliasesResponse:
176
+ """Get aliases for a specific vendor workflow.
177
+
178
+ Args:
179
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
180
+ workflow_id: ID of the workflow to retrieve aliases for
181
+ setting_id: ID of the vendor setting
182
+ per_page: Number of items per page
183
+ next_token: Token for pagination (optional, for retrieving next page)
184
+
185
+ Returns:
186
+ VendorWorkflowAliasesResponse containing list of aliases and pagination token
187
+
188
+ Example:
189
+ >>> # Get first page of aliases
190
+ >>> aliases = client.vendor_workflows.get_workflow_aliases(
191
+ ... vendor=VendorType.AWS,
192
+ ... workflow_id="9HXLQ7J9YP",
193
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c",
194
+ ... per_page=5
195
+ ... )
196
+ >>> for alias in aliases.data:
197
+ ... print(f"{alias.name} (v{alias.version}): {alias.status}")
198
+ ... if alias.aiRunId:
199
+ ... print(f" AI Run ID: {alias.aiRunId}")
200
+ >>> # Get next page if available
201
+ >>> if aliases.pagination.next_token:
202
+ ... next_page = client.vendor_workflows.get_workflow_aliases(
203
+ ... vendor=VendorType.AWS,
204
+ ... workflow_id="9HXLQ7J9YP",
205
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c",
206
+ ... per_page=5,
207
+ ... next_token=aliases.pagination.next_token
208
+ ... )
209
+ """
210
+ # Convert enum to string value if needed
211
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
212
+
213
+ params = {
214
+ "setting_id": setting_id,
215
+ "per_page": per_page,
216
+ }
217
+
218
+ if next_token:
219
+ params["next_token"] = next_token
220
+
221
+ return self._api.get(
222
+ f"/v1/vendors/{vendor_str}/workflows/{workflow_id}/aliases",
223
+ VendorWorkflowAliasesResponse,
224
+ params=params,
225
+ wrap_response=False,
226
+ )
227
+
228
+ def install_workflows(
229
+ self,
230
+ vendor: Union[VendorType, str],
231
+ workflows: List[VendorWorkflowInstallRequest],
232
+ ) -> VendorWorkflowInstallResponse:
233
+ """Install/activate vendor workflows.
234
+
235
+ Args:
236
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
237
+ workflows: List of workflow installation requests with workflow ID, flow alias ID, and setting ID
238
+
239
+ Returns:
240
+ VendorWorkflowInstallResponse containing installation summary with AI run IDs
241
+
242
+ Example:
243
+ >>> from codemie_sdk import VendorWorkflowInstallRequest
244
+ >>> # Install single workflow
245
+ >>> install_request = VendorWorkflowInstallRequest(
246
+ ... id="9HXLQ7J9YP",
247
+ ... flowAliasId="9RUV0BI2L7",
248
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c"
249
+ ... )
250
+ >>> response = client.vendor_workflows.install_workflows(
251
+ ... vendor=VendorType.AWS,
252
+ ... workflows=[install_request]
253
+ ... )
254
+ >>> for item in response.summary:
255
+ ... print(f"Installed workflow {item.flowId} with run ID: {item.aiRunId}")
256
+ >>>
257
+ >>> # Install multiple workflows
258
+ >>> requests = [
259
+ ... VendorWorkflowInstallRequest(
260
+ ... id="WORKFLOW_ID_1",
261
+ ... flowAliasId="ALIAS_ID_1",
262
+ ... setting_id="SETTING_ID"
263
+ ... ),
264
+ ... VendorWorkflowInstallRequest(
265
+ ... id="WORKFLOW_ID_2",
266
+ ... flowAliasId="ALIAS_ID_2",
267
+ ... setting_id="SETTING_ID"
268
+ ... )
269
+ ... ]
270
+ >>> response = client.vendor_workflows.install_workflows(
271
+ ... vendor=VendorType.AWS,
272
+ ... workflows=requests
273
+ ... )
274
+ """
275
+ # Convert enum to string value if needed
276
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
277
+
278
+ # Convert list of Pydantic models to list of dicts
279
+ payload = [workflow.model_dump(by_alias=True) for workflow in workflows]
280
+
281
+ return self._api.post(
282
+ f"/v1/vendors/{vendor_str}/workflows",
283
+ VendorWorkflowInstallResponse,
284
+ json_data=payload,
285
+ wrap_response=False,
286
+ )
287
+
288
+ def uninstall_workflow(
289
+ self,
290
+ vendor: Union[VendorType, str],
291
+ ai_run_id: str,
292
+ ) -> VendorWorkflowUninstallResponse:
293
+ """Uninstall/deactivate a vendor workflow.
294
+
295
+ Args:
296
+ vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
297
+ ai_run_id: AI run ID returned from the workflow alias (aiRunId field)
298
+
299
+ Returns:
300
+ VendorWorkflowUninstallResponse with success status
301
+
302
+ Example:
303
+ >>> # Get workflow aliases to find the aiRunId
304
+ >>> aliases = client.vendor_workflows.get_workflow_aliases(
305
+ ... vendor=VendorType.AWS,
306
+ ... workflow_id="9HXLQ7J9YP",
307
+ ... setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c"
308
+ ... )
309
+ >>> # Find an installed alias with aiRunId
310
+ >>> for alias in aliases.data:
311
+ ... if alias.aiRunId:
312
+ ... ai_run_id = alias.aiRunId
313
+ ... break
314
+ >>>
315
+ >>> # Uninstall the workflow using the AI run ID
316
+ >>> response = client.vendor_workflows.uninstall_workflow(
317
+ ... vendor=VendorType.AWS,
318
+ ... ai_run_id="56fed66d-f66e-46e3-b420-bb3a8d93eed4"
319
+ ... )
320
+ >>> if response.success:
321
+ ... print("Workflow successfully uninstalled!")
322
+ """
323
+ # Convert enum to string value if needed
324
+ vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
325
+
326
+ return self._api.delete(
327
+ f"/v1/vendors/{vendor_str}/workflows/{ai_run_id}",
328
+ VendorWorkflowUninstallResponse,
329
+ wrap_response=False,
330
+ )
@@ -0,0 +1,41 @@
1
+ import requests
2
+
3
+ from typing import Dict, Any
4
+ from ..utils import ApiRequestHandler
5
+
6
+
7
+ class WebhookService:
8
+ """Webhook service implementation."""
9
+
10
+ def __init__(self, api_domain: str, token: str, verify_ssl: bool = True):
11
+ """Initialize the conversation service.
12
+
13
+ Args:
14
+ api_domain: Base URL for the API
15
+ token: Authentication token
16
+ verify_ssl: Whether to verify SSL certificates
17
+ """
18
+ self._api = ApiRequestHandler(api_domain, token, verify_ssl)
19
+
20
+ def trigger(
21
+ self, webhook_id: str, data: Dict[str, Any] = None
22
+ ) -> requests.Response:
23
+ """Trigger a webhook by sending a POST request with the provided data.
24
+
25
+ Args:
26
+ webhook_id: Webhook ID to trigger
27
+ data: Request body data. Defaults to {"test": "data"} if not provided
28
+
29
+ Returns:
30
+ Response object from requests library
31
+ """
32
+ if data is None:
33
+ data = {"test": "data"}
34
+
35
+ return self._api.post(
36
+ f"/v1/webhooks/{webhook_id}",
37
+ response_model=requests.Response,
38
+ json_data=data,
39
+ wrap_response=False,
40
+ raise_on_error=False,
41
+ )
@@ -130,8 +130,32 @@ class WorkflowService:
130
130
  """
131
131
  return self._api.delete(f"/v1/workflows/{workflow_id}", dict)
132
132
 
133
- def run(self, workflow_id: str, user_input: str = "") -> dict:
134
- return self.executions(workflow_id).create(user_input)
133
+ def run(
134
+ self,
135
+ workflow_id: str,
136
+ user_input: Optional[str] = None,
137
+ file_name: Optional[str] = None,
138
+ propagate_headers: bool = False,
139
+ headers: Optional[dict[str, str]] = None,
140
+ ) -> dict:
141
+ """Run a workflow with optional input parameters.
142
+
143
+ Args:
144
+ workflow_id: ID of the workflow to run
145
+ user_input: Optional user input for the workflow execution
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)
149
+
150
+ Returns:
151
+ dict: Created workflow execution details
152
+ """
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
+ )
135
159
 
136
160
  def executions(self, workflow_id: str) -> WorkflowExecutionService:
137
161
  """Get workflow execution service for the specified workflow.
@@ -4,6 +4,8 @@ from typing import List, Optional
4
4
 
5
5
  from ..models.common import PaginationParams
6
6
  from ..models.workflow import WorkflowExecution
7
+ from ..models.workflow_execution_payload import WorkflowExecutionCreateRequest
8
+ from ..models.workflow_thoughts import WorkflowExecutionThought
7
9
  from .workflow_execution_state import WorkflowExecutionStateService
8
10
  from ..utils import ApiRequestHandler
9
11
 
@@ -43,19 +45,34 @@ class WorkflowExecutionService:
43
45
  params=params,
44
46
  )
45
47
 
46
- def create(self, user_input: Optional[str] = "") -> dict:
48
+ def create(
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,
54
+ ) -> dict:
47
55
  """Create a new workflow execution.
48
56
 
49
57
  Args:
50
58
  user_input: Optional input data for the workflow execution.
51
- Defaults to None.
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)
52
62
 
53
63
  Returns:
54
64
  dict: Created workflow execution details
55
65
  """
56
- payload = {"user_input": user_input}
66
+ payload = WorkflowExecutionCreateRequest(
67
+ user_input=user_input,
68
+ file_name=file_name,
69
+ propagate_headers=propagate_headers,
70
+ )
57
71
  return self._api.post(
58
- f"/v1/workflows/{self._workflow_id}/executions", dict, json_data=payload
72
+ f"/v1/workflows/{self._workflow_id}/executions",
73
+ dict,
74
+ json_data=payload.model_dump(),
75
+ extra_headers=headers,
59
76
  )
60
77
 
61
78
  def get(self, execution_id: str) -> WorkflowExecution:
@@ -100,15 +117,46 @@ class WorkflowExecutionService:
100
117
  f"/v1/workflows/{self._workflow_id}/executions/{execution_id}/abort", dict
101
118
  )
102
119
 
103
- 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:
104
126
  """Resume an interrupted workflow execution.
105
127
 
106
128
  Args:
107
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)
108
132
 
109
133
  Returns:
110
134
  dict: Updated workflow execution details
111
135
  """
136
+ params = {"propagate_headers": propagate_headers}
137
+ # Empty body per API; passing empty dict to satisfy typing
112
138
  return self._api.put(
113
- 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,
144
+ )
145
+
146
+ def get_thoughts(
147
+ self, execution_id: str, thought_ids: List[str]
148
+ ) -> List[WorkflowExecutionThought]:
149
+ """Get detailed thoughts information for specific thought IDs.
150
+
151
+ Args:
152
+ execution_id: ID of the execution to get thoughts for
153
+ thought_ids: List of thought IDs to retrieve detailed information for
154
+
155
+ Returns:
156
+ List[WorkflowExecutionThought]: List of detailed thought objects
157
+ """
158
+ return self._api.post(
159
+ f"/v1/workflows/{self._workflow_id}/executions/{execution_id}/thoughts",
160
+ List[WorkflowExecutionThought],
161
+ json_data=thought_ids,
114
162
  )
codemie_sdk/utils/http.py CHANGED
@@ -1,12 +1,11 @@
1
1
  """HTTP utilities for CodeMie SDK."""
2
2
 
3
- from typing import TypeVar, Type, Optional, Any, Union, Dict, List, get_origin, get_args
4
- from pydantic import BaseModel
3
+ from typing import TypeVar, Type, Optional, Any, Dict, List, get_origin, get_args
5
4
  import requests
6
5
  import logging
7
6
  from functools import wraps
8
7
 
9
- T = TypeVar("T", bound=Union[BaseModel, List[BaseModel], dict])
8
+ T = TypeVar("T")
10
9
 
11
10
  logger = logging.getLogger(__name__)
12
11
 
@@ -45,6 +44,19 @@ class ApiRequestHandler:
45
44
  self._base_url = base_url.rstrip("/")
46
45
  self._token = token
47
46
  self._verify_ssl = verify_ssl
47
+ self._is_localhost = self._is_localhost_domain(base_url)
48
+
49
+ @staticmethod
50
+ def _is_localhost_domain(domain: str) -> bool:
51
+ """Check if the domain is a localhost variant."""
52
+ domain_lower = domain.lower()
53
+ localhost_patterns = [
54
+ "localhost",
55
+ "127.0.0.1",
56
+ "0.0.0.0",
57
+ "192.168",
58
+ ]
59
+ return any(pattern in domain_lower for pattern in localhost_patterns)
48
60
 
49
61
  def _get_headers(self, exclude_content_type: bool = False) -> dict:
50
62
  """Gets request headers with auth token.
@@ -56,14 +68,9 @@ class ApiRequestHandler:
56
68
  if not exclude_content_type:
57
69
  headers["Content-Type"] = "application/json"
58
70
 
59
- if (
60
- "0.0.0.0" in self._base_url
61
- or "127.0.0.1" in self._base_url
62
- or "localhost" in self._base_url
63
- ):
64
- headers["User-Id"] = "dev-codemie-user"
65
- else:
66
- headers["Authorization"] = f"Bearer {self._token}"
71
+ headers["Authorization"] = (
72
+ f"Bearer {self._token}" if not self._is_localhost else "dev-codemie-user"
73
+ )
67
74
  return headers
68
75
 
69
76
  def _parse_response(
@@ -131,7 +138,7 @@ class ApiRequestHandler:
131
138
  wrap_response: Whether response is wrapped in 'data' field
132
139
 
133
140
  Returns:
134
- Parsed response object or list of objects
141
+ Parsed response object or list of objects, or raw Response if response_model is requests.Response
135
142
  """
136
143
  if params:
137
144
  logger.debug(f"Request params: {params}")
@@ -143,6 +150,9 @@ class ApiRequestHandler:
143
150
  )
144
151
  response.raise_for_status()
145
152
 
153
+ if response_model is None or response_model is requests.Response:
154
+ return response
155
+
146
156
  return self._parse_response(response, response_model, wrap_response)
147
157
 
148
158
  @log_request
@@ -153,7 +163,9 @@ class ApiRequestHandler:
153
163
  json_data: Optional[Dict[str, Any]] = None,
154
164
  stream: bool = False,
155
165
  wrap_response: bool = True,
156
- ) -> Union[T, requests.Response]:
166
+ raise_on_error: bool = True,
167
+ extra_headers: Optional[Dict[str, str]] = None,
168
+ ) -> T:
157
169
  """Makes a POST request and parses the response.
158
170
 
159
171
  Args:
@@ -162,6 +174,8 @@ class ApiRequestHandler:
162
174
  json_data: JSON request body
163
175
  stream: Whether to return streaming response
164
176
  wrap_response: Whether response is wrapped in 'data' field
177
+ raise_on_error: Whether to raise exception on HTTP error status codes
178
+ extra_headers: Optional additional HTTP headers to include (e.g., X-* for MCP propagation)
165
179
 
166
180
  Returns:
167
181
  Parsed response object/list or streaming response
@@ -169,18 +183,26 @@ class ApiRequestHandler:
169
183
  if json_data:
170
184
  logger.debug(f"Request body: {json_data}")
171
185
 
186
+ headers = self._get_headers()
187
+ if extra_headers:
188
+ headers.update(extra_headers)
189
+
172
190
  response = requests.post(
173
191
  url=f"{self._base_url}{endpoint}",
174
- headers=self._get_headers(),
192
+ headers=headers,
175
193
  json=json_data,
176
194
  verify=self._verify_ssl,
177
195
  stream=stream,
178
196
  )
179
- response.raise_for_status()
197
+ if raise_on_error:
198
+ response.raise_for_status()
180
199
 
181
200
  if stream:
182
201
  return response
183
202
 
203
+ if response_model is None or response_model is requests.Response:
204
+ return response
205
+
184
206
  return self._parse_response(response, response_model, wrap_response)
185
207
 
186
208
  @log_request
@@ -229,6 +251,7 @@ class ApiRequestHandler:
229
251
  json_data: Dict[str, Any],
230
252
  params: Optional[Dict[str, Any]] = None,
231
253
  wrap_response: bool = True,
254
+ extra_headers: Optional[Dict[str, str]] = None,
232
255
  ) -> T:
233
256
  """Makes a PUT request and parses the response.
234
257
 
@@ -238,14 +261,18 @@ class ApiRequestHandler:
238
261
  json_data: JSON request body
239
262
  params: Query parameters
240
263
  wrap_response: Whether response is wrapped in 'data' field
264
+ extra_headers: Optional additional HTTP headers to include (e.g., X-* for MCP propagation)
241
265
 
242
266
  Returns:
243
267
  Parsed response object or list of objects
244
268
  """
245
269
  logger.debug(f"Request body: {json_data}")
270
+ headers = self._get_headers()
271
+ if extra_headers:
272
+ headers.update(extra_headers)
246
273
  response = requests.put(
247
274
  url=f"{self._base_url}{endpoint}",
248
- headers=self._get_headers(),
275
+ headers=headers,
249
276
  json=json_data,
250
277
  params=params,
251
278
  verify=self._verify_ssl,