codemie-sdk-python 0.1.80__py3-none-any.whl → 0.1.81__py3-none-any.whl

codemie_sdk/models/assistant.py

@@ -3,7 +3,7 @@
 import uuid
 from datetime import datetime
 from enum import Enum
-from typing import List, Optional, Any, Union, Dict
+from typing import List, Optional, Any, Union, Dict, Type
 
 from pydantic import BaseModel, Field, ConfigDict, model_validator
 
@@ -227,12 +227,19 @@ class AssistantChatRequest(BaseModel):
         default=None, description="Provide additional metadata"
     )
     tools_config: Optional[List[ToolConfig]] = None
+    output_schema: Optional[dict | Type[BaseModel]] = Field(
+        default=None,
+        description="Structured output schema for the agent. \
+            If specified, `generated` field in response will have the same type",
+    )
 
 
 class BaseModelResponse(BaseModel):
     """Model for chat response from assistant."""
 
-    generated: str = Field(description="Generated response error_message")
+    generated: str | dict | BaseModel = Field(
+        description="Generated response. If output_schema in request is specified, corresponds with its type"
+    )
     time_elapsed: Optional[float] = Field(
         default=None, alias="timeElapsed", description="Time taken for generation"
     )

codemie_sdk/services/assistant.py

@@ -3,6 +3,8 @@
 import json
 from pathlib import Path
 from typing import List, Union, Optional, Dict, Any, Literal
+from pydantic import BaseModel
+from copy import deepcopy
 
 import requests
 
@@ -180,12 +182,23 @@ class AssistantService:
         Returns:
             Chat response or streaming response
         """
-        return self._api.post(
+        pydantic_schema = None
+        if issubclass(request.output_schema, BaseModel):
+            pydantic_schema = deepcopy(request.output_schema)
+            request.output_schema = request.output_schema.model_json_schema()
+
+        response = self._api.post(
             f"/v1/assistants/{assistant_id}/model",
             BaseModelResponse,
             json_data=request.model_dump(exclude_none=True, by_alias=True),
             stream=request.stream,
         )
+        if not request.stream and pydantic_schema:
+            # we do conversion to the BaseModel here because self._parse_response don't see actual request model,
+            # where reflected desired output format for structured output
+            response.generated = pydantic_schema.model_validate(response.generated)
+
+        return response
 
     def upload_file_to_chat(self, file_path: Path):
         """Upload a file to assistant chat and return the response containing file_url."""

{codemie_sdk_python-0.1.80.dist-info → codemie_sdk_python-0.1.81.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: codemie-sdk-python
-Version: 0.1.80
+Version: 0.1.81
 Summary: CodeMie SDK for Python
 Author: Vadym Vlasenko
 Author-email: vadym_vlasenko@epam.com
@@ -182,7 +182,50 @@ chat_request = AssistantChatRequest(
 response = client.assistant.chat("assistant-id", chat_request)
 ```
 
-7. **Work with Prebuilt Assistants**
+7. **Utilize structured outputs with Assistant**
+```python
+from pydantic import BaseModel
+
+class OutputSchema(BaseModel):
+    requirements: list[str]
+
+chat_request = AssistantChatRequest(
+    text="Your message here",
+    stream=False, 
+    output_schema=OutputSchema
+    # Additional parameters
+)
+
+response = client.assistants.chat("id", chat_request)
+# response.generated is a Pydantic object
+```
+Or using JSON schema in dict format
+```python
+output_schema = {
+    "properties": {
+        "requirements": {
+            "items": {"type": "string"},
+            "title": "Requirements",
+            "type": "array",
+        }
+    },
+    "required": ["requirements"],
+    "title": "OutputSchema",
+    "type": "object",
+}
+
+chat_request = AssistantChatRequest(
+    text="Your message here",
+    stream=False, 
+    output_schema=output_schema 
+    # Additional parameters
+)
+
+response = client.assistants.chat("id", chat_request)
+# response.generated is a dict corresponded with JSON schema
+```
+
+8. **Work with Prebuilt Assistants**
 ```python
 # List prebuilt assistants
 prebuilt = client.assistant.get_prebuilt()
@@ -191,7 +234,7 @@ prebuilt = client.assistant.get_prebuilt()
 prebuilt_assistant = client.assistant.get_prebuilt_by_slug("assistant-slug")
 ```
 
-8. **Get Available Tools**
+9. **Get Available Tools**
 ```python
 tools = client.assistant.get_tools()
 ```

{codemie_sdk_python-0.1.80.dist-info → codemie_sdk_python-0.1.81.dist-info}/RECORD

@@ -5,7 +5,7 @@ codemie_sdk/client/__init__.py,sha256=yf6C39MmrJ6gK9ZHMhBeynKwUUYVSUTQbKxU8-4qpK
 codemie_sdk/client/client.py,sha256=dVcWqd-ruWd8GOZ3_33C3LnHgUyyxIoCKxKYPYbBPq8,4373
 codemie_sdk/exceptions.py,sha256=XoVPyognx-JmyVxLHkZPAcX1CMi1OoT1diBFJLU54so,1183
 codemie_sdk/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-codemie_sdk/models/assistant.py,sha256=ZEhn5xH0EsmmZMI2U8kaGryPfFgg-YzlIV9XmOFiUMs,8524
+codemie_sdk/models/assistant.py,sha256=qIRJgM0cR_xjFbDLQ8yPypRXD-tJUqt2Ny47WffrXOo,8852
 codemie_sdk/models/common.py,sha256=gmZ-ps8TbaieNKr0kUKoQEjhVrHD2CAYomOpZQRatH8,1195
 codemie_sdk/models/datasource.py,sha256=u6X1xXZ3ZR9CWH0B0bSw3IjC7KyEwtRKHIc6ZuhJwsg,10813
 codemie_sdk/models/integration.py,sha256=c2D7QsM9HJ9zpupkjKz1v60rtBMGgTyJLYXUyAzS5no,2240
@@ -14,7 +14,7 @@ codemie_sdk/models/task.py,sha256=J4ZFRY3s8qBGrqB5NLQF0rMbInLh4s7OEZ0ZfmnW0Ho,14
 codemie_sdk/models/user.py,sha256=Q0rjimZh-IbeaPfq6b6fk6ZaCtwLqWHEIlU863suCS4,1777
 codemie_sdk/models/workflow.py,sha256=h3ZM-2bBucRd9CzZUYDTU0rGR1ihEN7cxkHgQmJfJVc,2394
 codemie_sdk/models/workflow_state.py,sha256=CMYFQZ7sy4QxmnWmc83TFfqP7TG_3rW5MdH5fxsS9kY,1251
-codemie_sdk/services/assistant.py,sha256=0OFx7enE65I2uiLpbIebnCXhap2QdkFl8CFBUPzf9Jc,7022
+codemie_sdk/services/assistant.py,sha256=4El_J_fPsj5OR4bN7EnmsNY9n02TLxt4uSw8s4b7J_s,7660
 codemie_sdk/services/datasource.py,sha256=gIP9D-sq4733x9f0mWL87bMI1pWYxIsv5eYXkAs0O6Q,7150
 codemie_sdk/services/integration.py,sha256=CrUfO-4_JwtCt8Qig3RP2uTi1uaTkedyTaPdkFY509g,5410
 codemie_sdk/services/llm.py,sha256=0-e4_7RvLHs2giCyoQ5U4KDTh6p5VXgPKNxnDP9ZDFU,1100
@@ -25,6 +25,6 @@ codemie_sdk/services/workflow_execution.py,sha256=aGoT3rdTmh5-doAsrmBBjLEuOfvL5a
 codemie_sdk/services/workflow_execution_state.py,sha256=tXoaa8yT09xgYEUNiHhVULe76TwGwVgZupMIUyyLxdo,2070
 codemie_sdk/utils/__init__.py,sha256=BXAJJfAzO89-kMYvWWo9wSNhSbGgF3vB1In9sePFhMM,109
 codemie_sdk/utils/http.py,sha256=YzE53y1r7PBUKcEQt-gGYduRY06I8xzqVRzlj9f_voA,10084
-codemie_sdk_python-0.1.80.dist-info/METADATA,sha256=hUotIIEzYd8HHbOi2ZwgiaRlCpyO259rRZ4c5hkHtJk,22155
-codemie_sdk_python-0.1.80.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-codemie_sdk_python-0.1.80.dist-info/RECORD,,
+codemie_sdk_python-0.1.81.dist-info/METADATA,sha256=-Dx6Pf-_kjlOIpA0kccg5dLStF8PP6hhndQH3UHyV0A,23146
+codemie_sdk_python-0.1.81.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+codemie_sdk_python-0.1.81.dist-info/RECORD,,