unique_toolkit 0.5.6__py3-none-any.whl → 0.5.8__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.
@@ -0,0 +1,10 @@
1
+ import logging
2
+ from typing import Optional
3
+
4
+ from unique_toolkit.chat.state import ChatState
5
+
6
+
7
+ class BaseService:
8
+ def __init__(self, state: ChatState, logger: Optional[logging.Logger] = None):
9
+ self.state = state
10
+ self.logger = logger or logging.getLogger(__name__)
@@ -0,0 +1,70 @@
1
+ import asyncio
2
+ import logging
3
+ import threading
4
+ import time
5
+ from typing import (
6
+ Awaitable,
7
+ Optional,
8
+ Sequence,
9
+ TypeVar,
10
+ Union,
11
+ )
12
+
13
+ T = TypeVar("T")
14
+ Result = Union[T, BaseException]
15
+
16
+
17
+ async def run_async_tasks_parallel(
18
+ tasks: Sequence[Awaitable[T]],
19
+ max_tasks: Optional[int] = None,
20
+ logger: logging.Logger = logging.getLogger(__name__),
21
+ ) -> list[Result]:
22
+ """
23
+ Executes the a set of given async tasks and returns the results.
24
+
25
+ Args:
26
+ tasks (list[Awaitable[T]]): list of async callables to execute in parallel.
27
+ max_tasks (int): Maximum number of tasks for the asyncio Semaphore.
28
+
29
+ Returns:
30
+ list[Result]: list of results from the executed tasks.
31
+ """
32
+
33
+ max_tasks = max_tasks or len(tasks)
34
+
35
+ async def logging_wrapper(task: Awaitable[T], task_id: int) -> Result:
36
+ thread = threading.current_thread()
37
+ start_time = time.time()
38
+
39
+ logger.info(
40
+ f"Thread {thread.name} (ID: {thread.ident}) starting task {task_id}"
41
+ )
42
+
43
+ try:
44
+ result = await task
45
+ return result
46
+ except Exception as e:
47
+ logger.error(
48
+ f"Thread {thread.name} (ID: {thread.ident}) - Task {task_id} failed with error: {e}"
49
+ )
50
+ return e
51
+ finally:
52
+ end_time = time.time()
53
+ duration = end_time - start_time
54
+ logger.debug(
55
+ f"Thread {thread.name} (ID: {thread.ident}) - Task {task_id} finished in {duration:.2f} seconds"
56
+ )
57
+
58
+ sem = asyncio.Semaphore(max_tasks)
59
+
60
+ async def sem_task(task: Awaitable[T], task_id: int) -> Result:
61
+ async with sem:
62
+ return await logging_wrapper(task, task_id)
63
+
64
+ wrapped_tasks: list[Awaitable[Result]] = [
65
+ sem_task(task, i) for i, task in enumerate(tasks)
66
+ ]
67
+
68
+ results: list[Result] = await asyncio.gather(*wrapped_tasks, return_exceptions=True)
69
+
70
+ return results
@@ -6,14 +6,6 @@ from typing import Any, Callable, Coroutine, TypeVar
6
6
  T = TypeVar("T")
7
7
 
8
8
 
9
- def to_async(func: Callable[..., T]) -> Callable[..., Coroutine[Any, Any, T]]:
10
- @wraps(func)
11
- async def wrapper(*args, **kwargs) -> T:
12
- return await asyncio.to_thread(func, *args, **kwargs)
13
-
14
- return wrapper
15
-
16
-
17
9
  def async_warning(func):
18
10
  @wraps(func)
19
11
  async def wrapper(*args, **kwargs):
@@ -26,3 +18,22 @@ def async_warning(func):
26
18
  return await func(*args, **kwargs)
27
19
 
28
20
  return wrapper
21
+
22
+
23
+ @async_warning
24
+ def to_async(func: Callable[..., T]) -> Callable[..., Coroutine[Any, Any, T]]:
25
+ """
26
+ Decorator to convert a synchronous function to an asynchronous function using a thread pool executor.
27
+
28
+ Args:
29
+ func (Callable[..., T]): The synchronous function to convert.
30
+
31
+ Returns:
32
+ Callable[..., Coroutine[Any, Any, T]]: The asynchronous function.
33
+ """
34
+
35
+ @wraps(func)
36
+ async def wrapper(*args, **kwargs) -> T:
37
+ return await asyncio.to_thread(func, *args, **kwargs)
38
+
39
+ return wrapper
@@ -3,15 +3,16 @@ import re
3
3
  from typing import Optional
4
4
 
5
5
  import unique_sdk
6
+ from unique_sdk._list_object import ListObject
6
7
 
7
- from unique_toolkit.app.performance.async_wrapper import async_warning, to_async
8
+ from unique_toolkit._common._base_service import BaseService
8
9
  from unique_toolkit.chat.schemas import ChatMessage, ChatMessageRole
9
10
  from unique_toolkit.chat.state import ChatState
10
11
  from unique_toolkit.content.schemas import ContentReference
11
12
  from unique_toolkit.content.utils import count_tokens
12
13
 
13
14
 
14
- class ChatService:
15
+ class ChatService(BaseService):
15
16
  """
16
17
  Provides all functionalities to manage the chat session.
17
18
 
@@ -21,8 +22,10 @@ class ChatService:
21
22
  """
22
23
 
23
24
  def __init__(self, state: ChatState, logger: Optional[logging.Logger] = None):
24
- self.state = state
25
- self.logger = logger or logging.getLogger(__name__)
25
+ super().__init__(state, logger)
26
+
27
+ DEFAULT_PERCENT_OF_MAX_TOKENS = 0.15
28
+ DEFAULT_MAX_MESSAGES = 4
26
29
 
27
30
  def modify_assistant_message(
28
31
  self,
@@ -46,16 +49,24 @@ class ChatService:
46
49
  Raises:
47
50
  Exception: If the modification fails.
48
51
  """
49
- return self._trigger_modify_assistant_message(
50
- content=content,
51
- message_id=message_id,
52
- references=references,
53
- debug_info=debug_info,
54
- )
52
+ message_id = message_id or self.state.assistant_message_id
53
+
54
+ try:
55
+ message = unique_sdk.Message.modify(
56
+ user_id=self.state.user_id,
57
+ company_id=self.state.company_id,
58
+ id=message_id, # type: ignore
59
+ chatId=self.state.chat_id,
60
+ text=content,
61
+ references=self._map_references(references), # type: ignore
62
+ debugInfo=debug_info or {},
63
+ )
64
+ except Exception as e:
65
+ self.logger.error(f"Failed to modify assistant message: {e}")
66
+ raise e
67
+ return ChatMessage(**message)
55
68
 
56
- @to_async
57
- @async_warning
58
- def async_modify_assistant_message(
69
+ async def modify_assistant_message_async(
59
70
  self,
60
71
  content: str,
61
72
  references: list[ContentReference] = [],
@@ -77,12 +88,22 @@ class ChatService:
77
88
  Raises:
78
89
  Exception: If the modification fails.
79
90
  """
80
- return self._trigger_modify_assistant_message(
81
- content,
82
- message_id,
83
- references,
84
- debug_info,
85
- )
91
+ message_id = message_id or self.state.assistant_message_id
92
+
93
+ try:
94
+ message = await unique_sdk.Message.modify_async(
95
+ user_id=self.state.user_id,
96
+ company_id=self.state.company_id,
97
+ id=message_id, # type: ignore
98
+ chatId=self.state.chat_id,
99
+ text=content,
100
+ references=self._map_references(references), # type: ignore
101
+ debugInfo=debug_info or {},
102
+ )
103
+ except Exception as e:
104
+ self.logger.error(f"Failed to modify assistant message: {e}")
105
+ raise e
106
+ return ChatMessage(**message)
86
107
 
87
108
  def get_full_history(self) -> list[ChatMessage]:
88
109
  """
@@ -96,9 +117,7 @@ class ChatService:
96
117
  """
97
118
  return self._get_full_history()
98
119
 
99
- @to_async
100
- @async_warning
101
- def async_get_full_history(self) -> list[ChatMessage]:
120
+ async def get_full_history_async(self) -> list[ChatMessage]:
102
121
  """
103
122
  Loads the full chat history for the chat session asynchronously.
104
123
 
@@ -108,21 +127,21 @@ class ChatService:
108
127
  Raises:
109
128
  Exception: If the loading fails.
110
129
  """
111
- return self._get_full_history()
130
+ return await self._get_full_history_async()
112
131
 
113
132
  def get_full_and_selected_history(
114
133
  self,
115
134
  token_limit: int,
116
- percent_of_max_tokens: float,
117
- max_messages: int,
135
+ percent_of_max_tokens: float = DEFAULT_PERCENT_OF_MAX_TOKENS,
136
+ max_messages: int = DEFAULT_MAX_MESSAGES,
118
137
  ) -> tuple[list[ChatMessage], list[ChatMessage]]:
119
138
  """
120
139
  Loads the chat history for the chat session synchronously.
121
140
 
122
141
  Args:
123
142
  token_limit (int): The maximum number of tokens to load.
124
- percent_of_max_tokens (float): The percentage of the maximum tokens to load.
125
- max_messages (int): The maximum number of messages to load.
143
+ percent_of_max_tokens (float): The percentage of the maximum tokens to load. Defaults to 0.15.
144
+ max_messages (int): The maximum number of messages to load. Defaults to 4.
126
145
 
127
146
  Returns:
128
147
  tuple[list[ChatMessage], list[ChatMessage]]: The selected and full chat history.
@@ -130,27 +149,28 @@ class ChatService:
130
149
  Raises:
131
150
  Exception: If the loading fails.
132
151
  """
133
- return self._get_full_and_selected_history(
134
- token_limit=token_limit,
135
- percent_of_max_tokens=percent_of_max_tokens,
152
+ full_history = self._get_full_history()
153
+ selected_history = self._get_selection_from_history(
154
+ full_history=full_history,
155
+ max_tokens=int(round(token_limit * percent_of_max_tokens)),
136
156
  max_messages=max_messages,
137
157
  )
138
158
 
139
- @to_async
140
- @async_warning
141
- def async_get_full_and_selected_history(
159
+ return full_history, selected_history
160
+
161
+ async def get_full_and_selected_history_async(
142
162
  self,
143
163
  token_limit: int,
144
- percent_of_max_tokens: float,
145
- max_messages: int,
164
+ percent_of_max_tokens: float = DEFAULT_PERCENT_OF_MAX_TOKENS,
165
+ max_messages: int = DEFAULT_MAX_MESSAGES,
146
166
  ) -> tuple[list[ChatMessage], list[ChatMessage]]:
147
167
  """
148
168
  Loads the chat history for the chat session asynchronously.
149
169
 
150
170
  Args:
151
171
  token_limit (int): The maximum number of tokens to load.
152
- percent_of_max_tokens (float): The percentage of the maximum tokens to load.
153
- max_messages (int): The maximum number of messages to load.
172
+ percent_of_max_tokens (float): The percentage of the maximum tokens to load. Defaults to 0.15.
173
+ max_messages (int): The maximum number of messages to load. Defaults to 4.
154
174
 
155
175
  Returns:
156
176
  tuple[list[ChatMessage], list[ChatMessage]]: The selected and full chat history.
@@ -158,12 +178,15 @@ class ChatService:
158
178
  Raises:
159
179
  Exception: If the loading fails.
160
180
  """
161
- return self._get_full_and_selected_history(
162
- token_limit=token_limit,
163
- percent_of_max_tokens=percent_of_max_tokens,
181
+ full_history = await self._get_full_history_async()
182
+ selected_history = self._get_selection_from_history(
183
+ full_history=full_history,
184
+ max_tokens=int(round(token_limit * percent_of_max_tokens)),
164
185
  max_messages=max_messages,
165
186
  )
166
187
 
188
+ return full_history, selected_history
189
+
167
190
  def create_assistant_message(
168
191
  self,
169
192
  content: str,
@@ -184,15 +207,24 @@ class ChatService:
184
207
  Raises:
185
208
  Exception: If the creation fails.
186
209
  """
187
- return self._trigger_create_assistant_message(
188
- content=content,
189
- references=references,
190
- debug_info=debug_info,
191
- )
192
210
 
193
- @to_async
194
- @async_warning
195
- def async_create_assistant_message(
211
+ try:
212
+ message = unique_sdk.Message.create(
213
+ user_id=self.state.user_id,
214
+ company_id=self.state.company_id,
215
+ chatId=self.state.chat_id,
216
+ assistantId=self.state.assistant_id,
217
+ text=content,
218
+ role=ChatMessageRole.ASSISTANT.name,
219
+ references=self._map_references(references), # type: ignore
220
+ debugInfo=debug_info,
221
+ )
222
+ except Exception as e:
223
+ self.logger.error(f"Failed to create assistant message: {e}")
224
+ raise e
225
+ return ChatMessage(**message)
226
+
227
+ async def create_assistant_message_async(
196
228
  self,
197
229
  content: str,
198
230
  references: list[ContentReference] = [],
@@ -212,45 +244,8 @@ class ChatService:
212
244
  Raises:
213
245
  Exception: If the creation fails.
214
246
  """
215
-
216
- return self._trigger_create_assistant_message(
217
- content=content,
218
- references=references,
219
- debug_info=debug_info,
220
- )
221
-
222
- def _trigger_modify_assistant_message(
223
- self,
224
- content: str,
225
- message_id: Optional[str],
226
- references: list[ContentReference],
227
- debug_info: dict,
228
- ) -> ChatMessage:
229
- message_id = message_id or self.state.assistant_message_id
230
-
231
247
  try:
232
- message = unique_sdk.Message.modify(
233
- user_id=self.state.user_id,
234
- company_id=self.state.company_id,
235
- id=message_id, # type: ignore
236
- chatId=self.state.chat_id,
237
- text=content,
238
- references=self._map_references(references), # type: ignore
239
- debugInfo=debug_info or {},
240
- )
241
- except Exception as e:
242
- self.logger.error(f"Failed to modify assistant message: {e}")
243
- raise e
244
- return ChatMessage(**message)
245
-
246
- def _trigger_create_assistant_message(
247
- self,
248
- content: str,
249
- references: list[ContentReference],
250
- debug_info: dict,
251
- ) -> ChatMessage:
252
- try:
253
- message = unique_sdk.Message.create(
248
+ message = await unique_sdk.Message.create_async(
254
249
  user_id=self.state.user_id,
255
250
  company_id=self.state.company_id,
256
251
  chatId=self.state.chat_id,
@@ -278,25 +273,21 @@ class ChatService:
278
273
  for ref in references
279
274
  ]
280
275
 
281
- def _get_full_and_selected_history(
282
- self,
283
- token_limit,
284
- percent_of_max_tokens=0.15,
285
- max_messages=4,
286
- ):
287
- full_history = self._get_full_history()
288
- selected_history = self._get_selection_from_history(
289
- full_history,
290
- int(round(token_limit * percent_of_max_tokens)),
291
- max_messages,
292
- )
276
+ def _get_full_history(self):
277
+ messages = self._trigger_list_messages(self.state.chat_id)
278
+ messages = self._filter_valid_messages(messages)
293
279
 
294
- return full_history, selected_history
280
+ return self._map_to_chat_messages(messages)
295
281
 
296
- def _get_full_history(self):
297
- SYSTEM_MESSAGE_PREFIX = "[SYSTEM] "
282
+ async def _get_full_history_async(self):
283
+ messages = await self._trigger_list_messages_async(self.state.chat_id)
284
+ messages = self._filter_valid_messages(messages)
298
285
 
299
- messages = self._trigger_list_messages(self.state.chat_id)
286
+ return self._map_to_chat_messages(messages)
287
+
288
+ @staticmethod
289
+ def _filter_valid_messages(messages: ListObject[unique_sdk.Message]):
290
+ SYSTEM_MESSAGE_PREFIX = "[SYSTEM] "
300
291
 
301
292
  # Remove the last two messages
302
293
  messages = messages["data"][:-2] # type: ignore
@@ -309,7 +300,7 @@ class ChatService:
309
300
  else:
310
301
  filtered_messages.append(message)
311
302
 
312
- return self._map_to_chat_messages(filtered_messages)
303
+ return filtered_messages
313
304
 
314
305
  def _trigger_list_messages(self, chat_id: str):
315
306
  try:
@@ -323,6 +314,18 @@ class ChatService:
323
314
  self.logger.error(f"Failed to list chat history: {e}")
324
315
  raise e
325
316
 
317
+ async def _trigger_list_messages_async(self, chat_id: str):
318
+ try:
319
+ messages = await unique_sdk.Message.list_async(
320
+ user_id=self.state.user_id,
321
+ company_id=self.state.company_id,
322
+ chatId=chat_id,
323
+ )
324
+ return messages
325
+ except Exception as e:
326
+ self.logger.error(f"Failed to list chat history: {e}")
327
+ raise e
328
+
326
329
  @staticmethod
327
330
  def _map_to_chat_messages(messages: list[dict]):
328
331
  return [ChatMessage(**msg) for msg in messages]
@@ -13,8 +13,6 @@ class ChatState:
13
13
  company_id (str): The company ID.
14
14
  user_id (str): The user ID.
15
15
  chat_id (str): The chat ID.
16
- scope_ids (list[str] | None): The scope IDs.
17
- chat_only (bool): The chat only flag.
18
16
  user_message_text (str): The user message text.
19
17
  user_message_id (str): The user message ID.
20
18
  assistant_message_id (str): The assistant message ID.
@@ -24,8 +22,6 @@ class ChatState:
24
22
  user_id: str
25
23
  assistant_id: str
26
24
  chat_id: str
27
- scope_ids: list[str] | None = None
28
- chat_only: bool = False
29
25
  user_message_text: str | None = None
30
26
  user_message_id: str | None = None
31
27
  assistant_message_id: str | None = None
@@ -42,17 +38,11 @@ class ChatState:
42
38
  Returns:
43
39
  ChatManager: The ChatManager instance.
44
40
  """
45
- config = event.payload.configuration
46
-
47
- scope_ids = config.get("scopeIds") or None
48
- chat_only = config.get("scopeToChatOnUpload", False)
49
41
  return cls(
50
42
  user_id=event.user_id,
51
43
  chat_id=event.payload.chat_id,
52
44
  company_id=event.company_id,
53
45
  assistant_id=event.payload.assistant_id,
54
- scope_ids=scope_ids,
55
- chat_only=chat_only,
56
46
  user_message_text=event.payload.user_message.text,
57
47
  user_message_id=event.payload.user_message.id,
58
48
  assistant_message_id=event.payload.assistant_message.id,
@@ -3,7 +3,7 @@ from enum import StrEnum
3
3
  from typing import Optional
4
4
 
5
5
  from humps import camelize
6
- from pydantic import BaseModel, ConfigDict
6
+ from pydantic import BaseModel, ConfigDict, Field
7
7
 
8
8
  # set config to convert camelCase to snake_case
9
9
  model_config = ConfigDict(
@@ -92,5 +92,5 @@ class ContentUploadInput(BaseModel):
92
92
 
93
93
  class RerankerConfig(BaseModel):
94
94
  model_config = model_config
95
- deployment_name: str
95
+ deployment_name: str = Field(serialization_alias="deploymentName")
96
96
  options: dict | None = None