superwise-api 4.13.2__tar.gz → 4.15.0__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 (75) hide show
  1. {superwise_api-4.13.2 → superwise_api-4.15.0}/PKG-INFO +2 -1
  2. {superwise_api-4.13.2 → superwise_api-4.15.0}/pyproject.toml +2 -1
  3. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/client/api_client.py +71 -0
  4. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/client/exceptions.py +11 -0
  5. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/client/rest.py +88 -0
  6. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/agent.py +89 -98
  7. superwise_api-4.15.0/superwise_api/entities/base.py +215 -0
  8. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/dashboard.py +3 -89
  9. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/dataset.py +76 -92
  10. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/guardrails.py +153 -2
  11. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/knowledge.py +7 -3
  12. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/policy.py +119 -83
  13. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/agent/agent.py +14 -2
  14. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/agent/playground.py +32 -0
  15. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/dashboard/dashboard.py +5 -21
  16. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/dataset/dataset.py +7 -0
  17. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/dataset/dataset_schema.py +2 -0
  18. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/guardrails/guardrails.py +50 -0
  19. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/knowledge/knowledge.py +1 -0
  20. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/policy/policy.py +8 -2
  21. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/utils.py +29 -1
  22. superwise_api-4.13.2/superwise_api/entities/base.py +0 -77
  23. {superwise_api-4.13.2 → superwise_api-4.15.0}/README.md +0 -0
  24. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/__init__.py +0 -0
  25. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/client/__init__.py +0 -0
  26. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/client/api_response.py +0 -0
  27. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/client/configuration.py +0 -0
  28. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/client/models/page.py +0 -0
  29. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/client/models/threshold_settings.py +0 -0
  30. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/config.py +0 -0
  31. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/__init__.py +0 -0
  32. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/dashboard_item.py +0 -0
  33. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/dataset_source.py +0 -0
  34. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/destination.py +0 -0
  35. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/integration.py +0 -0
  36. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/mcp.py +0 -0
  37. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/model.py +0 -0
  38. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/model_provider.py +0 -0
  39. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/source.py +0 -0
  40. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/entities/tag.py +0 -0
  41. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/errors.py +0 -0
  42. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/__init__.py +0 -0
  43. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/agent/__init__.py +0 -0
  44. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/agent/feedback.py +0 -0
  45. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/agent/flowise.py +0 -0
  46. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/context/__init__.py +0 -0
  47. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/context/context.py +0 -0
  48. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/dashboard/__init__.py +0 -0
  49. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/dashboard/query.py +0 -0
  50. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/dashboard_item/__init__.py +0 -0
  51. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/dashboard_item/dashboard_item.py +0 -0
  52. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/dataset/__init__.py +0 -0
  53. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/dataset_source/__init__.py +0 -0
  54. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/dataset_source/dataset_source.py +0 -0
  55. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/destination/__init__.py +0 -0
  56. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/destination/destination.py +0 -0
  57. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/guardrails/__init__.py +0 -0
  58. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/integration/__init__.py +0 -0
  59. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/integration/integration.py +0 -0
  60. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/knowledge/__init__.py +0 -0
  61. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/mcp/__init__.py +0 -0
  62. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/mcp/mcp.py +0 -0
  63. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/model/__init__.py +0 -0
  64. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/model/model.py +0 -0
  65. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/model_provider/__init__.py +0 -0
  66. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/model_provider/model_provider.py +0 -0
  67. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/policy/__init__.py +0 -0
  68. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/policy/query.py +0 -0
  69. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/source/__init__.py +0 -0
  70. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/source/source.py +0 -0
  71. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/tag/__init__.py +0 -0
  72. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/tag/tag.py +0 -0
  73. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/tool/__init__.py +0 -0
  74. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/models/tool/tool.py +0 -0
  75. {superwise_api-4.13.2 → superwise_api-4.15.0}/superwise_api/superwise_client.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: superwise-api
3
- Version: 4.13.2
3
+ Version: 4.15.0
4
4
  Summary: Superwise SDK for Python
5
5
  License: MIT
6
6
  Keywords: OpenAPI,SuperwiseSDK
@@ -12,6 +12,7 @@ Classifier: Programming Language :: Python :: 3.10
12
12
  Classifier: Programming Language :: Python :: 3.11
13
13
  Classifier: Programming Language :: Python :: 3.12
14
14
  Classifier: Programming Language :: Python :: 3.13
15
+ Requires-Dist: Pillow (>=10.0.0,<11.0.0)
15
16
  Requires-Dist: aenum (>=3.1.11)
16
17
  Requires-Dist: pydantic (>=2.10.4,<3.0.0)
17
18
  Requires-Dist: pydantic-settings (>=2.7.0,<3.0.0)
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  packages = [{ include = "superwise_api", from = "." }]
3
- version = "4.13.2"
3
+ version = "4.15.0"
4
4
  include = ["superwise_api/client/py.typed"]
5
5
 
6
6
  [tool.poetry.dependencies]
@@ -10,6 +10,7 @@ pydantic = "^2.10.4"
10
10
  aenum = ">=3.1.11"
11
11
  requests = "^2.32.0"
12
12
  pydantic-settings = "^2.7.0"
13
+ Pillow = "^10.0.0"
13
14
 
14
15
 
15
16
  [tool.poetry.group.dev.dependencies]
@@ -18,7 +18,9 @@ from superwise_api.client.configuration import Configuration
18
18
  from superwise_api.client.exceptions import ApiException
19
19
  from superwise_api.client.exceptions import ApiTypeError
20
20
  from superwise_api.client.exceptions import ApiValueError
21
+ from superwise_api.client.exceptions import StreamingError
21
22
  from superwise_api.client.models.page import Page
23
+ from superwise_api.models.agent.playground import StreamEventType
22
24
 
23
25
 
24
26
  class ApiClient:
@@ -1145,6 +1147,75 @@ class ApiClient:
1145
1147
  _request_auth=_params.get("_request_auth"),
1146
1148
  )
1147
1149
 
1150
+ def post_stream(
1151
+ self,
1152
+ resource_path: str,
1153
+ response_model,
1154
+ data: Optional[dict] = None,
1155
+ **kwargs,
1156
+ ):
1157
+ """Performs a streaming POST request and yields parsed response objects.
1158
+
1159
+ :param resource_path: Path to method endpoint.
1160
+ :param response_model: Model class to deserialize each line to.
1161
+ :param data: Request body data.
1162
+ :param kwargs: Additional arguments (_headers, _request_timeout, _request_auth).
1163
+ :yields: Parsed response_model instances.
1164
+ :raises StreamingError: If an error event is received in the stream.
1165
+ """
1166
+ _params = kwargs
1167
+
1168
+ # Process header parameters
1169
+ _header_params = dict(_params.get("_headers", {}))
1170
+ _header_params.update(self.default_headers)
1171
+ if self.cookie:
1172
+ _header_params["Cookie"] = self.cookie
1173
+ _header_params = self.sanitize_for_serialization(_header_params)
1174
+ _header_params = dict(self.parameters_to_tuples(_header_params, {}))
1175
+
1176
+ # Set content type and accept headers
1177
+ _header_params["Accept"] = "application/x-ndjson"
1178
+ _header_params["Content-Type"] = "application/json"
1179
+
1180
+ # Handle authentication
1181
+ _request_auth = _params.get("_request_auth")
1182
+ auth_settings = ["implicit"]
1183
+ self.update_params_for_auth(
1184
+ _header_params, [], auth_settings, resource_path, "POST", data, request_auth=_request_auth
1185
+ )
1186
+
1187
+ # Sanitize body
1188
+ if data:
1189
+ data = self.sanitize_for_serialization(data)
1190
+
1191
+ # Build URL
1192
+ url = self.configuration.host + resource_path
1193
+
1194
+ # Make streaming request
1195
+ streaming_response = self.rest_client.post_stream(
1196
+ url,
1197
+ headers=_header_params,
1198
+ body=data,
1199
+ _request_timeout=_params.get("_request_timeout"),
1200
+ )
1201
+
1202
+ try:
1203
+ for line in streaming_response.iter_lines():
1204
+ try:
1205
+ parsed = json.loads(line)
1206
+ response_obj = response_model.from_dict(parsed)
1207
+
1208
+ # Raise exception for error events
1209
+ if response_obj.event == StreamEventType.ERROR:
1210
+ raise StreamingError(response_obj.text)
1211
+
1212
+ yield response_obj
1213
+ except json.JSONDecodeError:
1214
+ # Skip malformed lines
1215
+ continue
1216
+ finally:
1217
+ streaming_response.release_conn()
1218
+
1148
1219
  def get(
1149
1220
  self,
1150
1221
  resource_path: str,
@@ -134,6 +134,17 @@ class ServiceException(ApiException):
134
134
  super(ServiceException, self).__init__(status, reason, http_resp)
135
135
 
136
136
 
137
+ class StreamingError(OpenApiException):
138
+ """Exception raised when a streaming response contains an error event."""
139
+
140
+ def __init__(self, message: str) -> None:
141
+ self.message = message
142
+ super().__init__(message)
143
+
144
+ def __str__(self):
145
+ return f"Streaming error: {self.message}"
146
+
147
+
137
148
  def render_path(path_to_item):
138
149
  """Returns a string representation of a path"""
139
150
  result = ""
@@ -46,6 +46,33 @@ class RESTResponse(io.IOBase):
46
46
  return self.urllib3_response.headers.get(name, default)
47
47
 
48
48
 
49
+ class StreamingRESTResponse:
50
+ """A streaming response wrapper that yields lines from NDJSON responses."""
51
+
52
+ def __init__(self, resp) -> None:
53
+ self.urllib3_response = resp
54
+ self.status = resp.status
55
+ self.reason = resp.reason
56
+
57
+ def iter_lines(self):
58
+ """Iterate over response lines for NDJSON streaming."""
59
+ buffer = b""
60
+ for chunk in self.urllib3_response.stream(amt=1024):
61
+ buffer += chunk
62
+ while b"\n" in buffer:
63
+ line, buffer = buffer.split(b"\n", 1)
64
+ line = line.strip()
65
+ if line:
66
+ yield line.decode("utf-8")
67
+ # Handle any remaining data in buffer
68
+ if buffer.strip():
69
+ yield buffer.strip().decode("utf-8")
70
+
71
+ def release_conn(self):
72
+ """Release the connection back to the pool."""
73
+ self.urllib3_response.release_conn()
74
+
75
+
49
76
  class RESTClientObject:
50
77
  def __init__(self, configuration, pools_size=4, maxsize=None) -> None:
51
78
  # urllib3.PoolManager will pass all kw parameters to connectionpool
@@ -371,3 +398,64 @@ class RESTClientObject:
371
398
  _request_timeout=_request_timeout,
372
399
  body=body,
373
400
  )
401
+
402
+ def post_stream(
403
+ self,
404
+ url,
405
+ headers=None,
406
+ body=None,
407
+ _request_timeout=None,
408
+ ):
409
+ """Perform a streaming POST request.
410
+
411
+ :param url: http request url
412
+ :param headers: http request headers
413
+ :param body: request json body
414
+ :param _request_timeout: timeout setting for this request
415
+ :return: StreamingRESTResponse object for iterating over lines
416
+ """
417
+ headers = headers or {}
418
+
419
+ timeout = None
420
+ if _request_timeout:
421
+ if isinstance(_request_timeout, (int, float)):
422
+ timeout = urllib3.Timeout(total=_request_timeout)
423
+ elif isinstance(_request_timeout, tuple) and len(_request_timeout) == 2:
424
+ timeout = urllib3.Timeout(connect=_request_timeout[0], read=_request_timeout[1])
425
+
426
+ try:
427
+ request_body = None
428
+ if body is not None:
429
+ request_body = json.dumps(body)
430
+
431
+ r = self.pool_manager.request(
432
+ "POST",
433
+ url,
434
+ body=request_body,
435
+ preload_content=False,
436
+ timeout=timeout,
437
+ headers=headers,
438
+ )
439
+ except urllib3.exceptions.SSLError as e:
440
+ msg = "{0}\n{1}".format(type(e).__name__, str(e))
441
+ raise ApiException(status=0, reason=msg)
442
+
443
+ if not 200 <= r.status <= 299:
444
+ # Read the error response body and wrap in RESTResponse
445
+ r._body = r.read()
446
+ r.release_conn()
447
+ rest_response = RESTResponse(r)
448
+
449
+ if r.status == 400:
450
+ raise BadRequestException(http_resp=rest_response)
451
+ if r.status == 401:
452
+ raise UnauthorizedException(http_resp=rest_response)
453
+ if r.status == 403:
454
+ raise ForbiddenException(http_resp=rest_response)
455
+ if r.status == 404:
456
+ raise NotFoundException(http_resp=rest_response)
457
+ if 500 <= r.status <= 599:
458
+ raise ServiceException(http_resp=rest_response)
459
+ raise ApiException(http_resp=rest_response)
460
+
461
+ return StreamingRESTResponse(r)
@@ -1,25 +1,24 @@
1
+ from typing import Iterator
1
2
  from typing import Optional
2
3
  from typing import Sequence
3
- from typing import Any
4
- from typing import Literal
5
4
 
6
5
  from pydantic import conint
7
6
 
8
7
  from superwise_api.client.api_client import ApiClient
9
8
  from superwise_api.client.models.page import Page
10
- from superwise_api.entities.base import BaseApi
11
- from superwise_api.models.agent.agent import AgentConfig, Version, ExtendedAgent
9
+ from superwise_api.entities.base import BaseApi, CoreEntitiyApi
10
+ from superwise_api.models.agent.agent import AgentConfig, Version, GuardrailConfig
12
11
  from superwise_api.models.agent.agent import Agent
13
12
  from superwise_api.models.agent.agent import ModelLLM
14
13
  from superwise_api.models.agent.agent import ToolDef
15
14
  from superwise_api.models.agent.flowise import FlowiseCredentialUserInput
16
15
  from superwise_api.models.agent.playground import AskResponsePayload
16
+ from superwise_api.models.agent.playground import StreamedResponse
17
17
  from superwise_api.models.dataset.dataset import Dataset
18
18
  from superwise_api.models.agent.feedback import EventFeedbackData
19
- from superwise_api.models.utils import SearchParams
20
19
 
21
20
 
22
- class AgentApi(BaseApi):
21
+ class AgentApi(CoreEntitiyApi):
23
22
  """
24
23
  This class provides methods to interact with the Agent API.
25
24
 
@@ -49,7 +48,6 @@ class AgentApi(BaseApi):
49
48
  description: str = None,
50
49
  authentication_enabled: bool = False,
51
50
  observability_enabled: bool = True,
52
- block_guardrails_violations: bool = True,
53
51
  guardrails_violation_message: str | None = None,
54
52
  **kwargs,
55
53
  ) -> Agent:
@@ -70,7 +68,6 @@ class AgentApi(BaseApi):
70
68
  "name": name,
71
69
  "authentication_enabled": authentication_enabled,
72
70
  "observability_enabled": observability_enabled,
73
- "block_guardrails_violations": block_guardrails_violations,
74
71
  }
75
72
  if description:
76
73
  payload["description"] = description
@@ -89,7 +86,7 @@ class AgentApi(BaseApi):
89
86
  agent_id: str,
90
87
  name: str,
91
88
  agent_config: AgentConfig,
92
- guardrails: list[str] = None,
89
+ guardrails: list[GuardrailConfig] | list[str] = None,
93
90
  description: str = None,
94
91
  **kwargs,
95
92
  ) -> Version:
@@ -101,16 +98,21 @@ class AgentApi(BaseApi):
101
98
  name (str): The name of the version.
102
99
  description (str): The version's description.
103
100
  agent_config (AgentConfig): agent configuration for this version.
104
- guardrails: set(str): set of guardrails ids to apply to this version.
101
+ guardrails: list(GuardrailConfig): list of guardrails ids and blocking configuration.
105
102
 
106
103
  Returns:
107
104
  Version: The created version.
108
105
  """
109
106
 
107
+ processed_guardrails = [
108
+ GuardrailConfig(id=guardrail) if not isinstance(guardrail, GuardrailConfig) else guardrail
109
+ for guardrail in (guardrails or [])
110
+ ]
111
+
110
112
  payload = {
111
113
  "name": name,
112
114
  "agent_config": agent_config,
113
- "guardrails": guardrails if guardrails else [],
115
+ "guardrails": [guardrail.model_dump() for guardrail in processed_guardrails],
114
116
  }
115
117
  if description:
116
118
  payload["description"] = description
@@ -129,7 +131,6 @@ class AgentApi(BaseApi):
129
131
  description: str | None = None,
130
132
  authentication_enabled: bool | None = None,
131
133
  observability_enabled: bool | None = None,
132
- block_guardrails_violations: bool | None = None,
133
134
  guardrails_violation_message: str | None = None,
134
135
  dataset_id: str | None = None,
135
136
  **kwargs,
@@ -143,7 +144,6 @@ class AgentApi(BaseApi):
143
144
  description (str, optional): Description for the agent.
144
145
  authentication_enabled (bool, optional): Whether the agent requires an api token for access or not.
145
146
  observability_enabled (bool, optional): Whether the agent logs conversation to the db or not.
146
- block_guardrails_violations (bool, optional): Whether to block guardrails violations.
147
147
  guardrails_violation_message (str, optional): The message to show when guardrails are violated.
148
148
  dataset_id (str, optional): The id of the dataset to associate with the agent.
149
149
 
@@ -156,7 +156,6 @@ class AgentApi(BaseApi):
156
156
  "description": description,
157
157
  "authentication_enabled": authentication_enabled,
158
158
  "observability_enabled": observability_enabled,
159
- "block_guardrails_violations": block_guardrails_violations,
160
159
  "guardrails_violation_message": guardrails_violation_message,
161
160
  "dataset_id": dataset_id,
162
161
  }
@@ -398,6 +397,80 @@ class AgentApi(BaseApi):
398
397
  **kwargs,
399
398
  )
400
399
 
400
+ def ask_playground_stream(
401
+ self,
402
+ input: str,
403
+ agent_config: AgentConfig,
404
+ chat_history: Optional[Sequence[dict]] = None,
405
+ guardrails: Optional[list[str]] = None,
406
+ **kwargs,
407
+ ) -> Iterator[StreamedResponse]:
408
+ """
409
+ Performs streaming ask request in playground mode.
410
+
411
+ Args:
412
+ input (str): The input to the model.
413
+ agent_config (AgentConfig): The type of the agent and connected tools/context.
414
+ chat_history (Sequence[dict], optional): The chat history.
415
+ guardrails (list[str], optional): List of guardrail IDs to apply.
416
+
417
+ Yields:
418
+ StreamedResponse: Individual streaming events (system, guardrails, llm, tool, response).
419
+
420
+ Raises:
421
+ StreamingError: If an error event is received in the stream.
422
+ """
423
+ payload = {
424
+ "config": agent_config,
425
+ "input": input,
426
+ "chat_history": chat_history or [],
427
+ }
428
+ if guardrails:
429
+ payload["guardrails"] = guardrails
430
+ return self.api_client.post_stream(
431
+ resource_path=self._playground_resource_path + "/ask/stream",
432
+ response_model=StreamedResponse,
433
+ data=payload,
434
+ **kwargs,
435
+ )
436
+
437
+ def ask_worker_stream(
438
+ self,
439
+ agent_id: str,
440
+ input: str,
441
+ api_token: str | None = None,
442
+ chat_history: Optional[Sequence[dict]] = None,
443
+ context_id: str | None = None,
444
+ metadata: dict | None = None,
445
+ **kwargs,
446
+ ) -> Iterator[StreamedResponse]:
447
+ """
448
+ Performs streaming ask request to the specified worker.
449
+
450
+ Args:
451
+ agent_id (str): The agent asked.
452
+ input (str): The input to the model.
453
+ api_token (str): The API token of the agent.
454
+ chat_history (Sequence[dict], optional): The chat history.
455
+ context_id (str, optional): The context ID for the conversation.
456
+ metadata (dict, optional): Additional metadata.
457
+
458
+ Yields:
459
+ StreamedResponse: Individual streaming events (system, guardrails, llm, tool, response).
460
+
461
+ Raises:
462
+ StreamingError: If an error event is received in the stream.
463
+ """
464
+ payload = {"input": input, "chat_history": chat_history or [], "context_id": context_id, "metadata": metadata}
465
+ headers = {"x-api-token": api_token} if api_token else {}
466
+ return self.api_client.post_stream(
467
+ resource_path=f"/v1/app-worker/{agent_id}/v1/ask/stream",
468
+ response_model=StreamedResponse,
469
+ data=payload,
470
+ _headers=headers,
471
+ **kwargs,
472
+ )
473
+
401
474
  def get_flowise_credential_schema(self, url: str, api_key: str, flow_id, **kwargs) -> FlowiseCredentialUserInput:
402
475
  """
403
476
  Get credential schema.
@@ -527,87 +600,5 @@ class AgentApi(BaseApi):
527
600
  **kwargs,
528
601
  )
529
602
 
530
- def add_tags(
531
- self,
532
- agent_id: str,
533
- tag_ids: list[str],
534
- **kwargs,
535
- ):
536
- """
537
- Add tags to an agent.
538
-
539
- Args:
540
- agent_id (str): The id of the agent.
541
- tag_ids list(str): List of tag ids to add the agent
542
- Returns:
543
- None
544
- """
545
- return self.api_client.post(
546
- resource_path=self._resource_path + f"/{agent_id}/tags",
547
- data=tag_ids,
548
- **kwargs,
549
- )
550
-
551
- def remove_tags(
552
- self,
553
- agent_id: str,
554
- tag_ids: list[str],
555
- **kwargs,
556
- ):
557
- """
558
- Delete tags from an agent.
559
-
560
- Args:
561
- agent_id (str): The id of the agent.
562
- tag_ids list(str): List of tag ids to remove from the agent
563
- Returns:
564
- None
565
- """
566
- return self.api_client.delete(
567
- resource_path=self._resource_path + "/{agent_id}/tags",
568
- model_name=self._model_name,
569
- entity_id=agent_id,
570
- path_params={"agent_id": agent_id},
571
- query_params=[("ids", id) for id in tag_ids],
572
- **kwargs,
573
- )
574
-
575
- def search(
576
- self,
577
- filters: list[Any] | None = None,
578
- search: str | None = None,
579
- sort_by: str | None = None,
580
- sort_direction: Literal["asc", "desc"] = "desc",
581
- **kwargs,
582
- ) -> Page:
583
- """
584
- Searches for agents based on a prefix.
585
-
586
- Args:
587
- filters (list[Any]): Filter on db columns, list of tuples.
588
- e.g. [[["id", "eq", "5c05dc9f-f04a-4ce8-9d57-2ec63ee76aac"], "and", ["description", "ilike", "Construction"]], "or", ["name", "ilike", "active"]]
589
- search (str): Free text search on searchable fields
590
- sort_by (str): Field to sort by
591
- sort_direction (Literal["asc", "desc"]): Sort direction (ascending or descending)
592
- query_params can be passed as part of the kwargs for pagination
593
-
594
- Returns:
595
- Page: A page of enriched agents.
596
- """
597
- data = SearchParams(
598
- filters=filters,
599
- search=search,
600
- sort_by=sort_by,
601
- sort_direction=sort_direction,
602
- ).model_dump(exclude_none=True)
603
- response_types_map = {
604
- "200": Page.set_model(ExtendedAgent),
605
- "422": None,
606
- "500": None,
607
- }
608
- return self.api_client.post(
609
- resource_path="/v1/enriched-agents/search",
610
- data=data,
611
- response_types_map=response_types_map,
612
- **kwargs,
613
- )
603
+ def unpublish(self, agent_id: str, **kwargs):
604
+ return self.api_client.post(f"{self._resource_path}/{agent_id}/unpublish", **kwargs)