image_get fix, history bulk compression fix
frdel committed
May 20, 2025 at 10:57 UTC
5db74202d632306a883ccce7339c5bdba0d16c5a
4 files changed
+39
-43
python/api/image_get.py
+19
-1
@@ -1,5 +1,7 @@
1
import os
2
+import re
3
from python.helpers.api import ApiHandler
4
+from python.helpers import files
5
from flask import Request, Response, send_file
6
7
@@ -9,7 +11,23 @@ class ImageGet(ApiHandler):
11
path = input.get("path", request.args.get("path", ""))
12
if not path:
13
raise ValueError("No path provided")
12
-
14
+
15
+ # check if path is within base directory
16
+ if not files.is_in_base_dir(path):
17
+ raise ValueError("Path is outside of allowed directory")
18
+
19
+ # check if file has an image extension
20
+ # list of allowed image extensions
21
+ allowed_extensions = [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp", ".svg"]
22
+ # get file extension
23
+ file_ext = os.path.splitext(path)[1].lower()
24
+ if file_ext not in allowed_extensions:
25
+ raise ValueError(f"File type not allowed. Allowed types: {', '.join(allowed_extensions)}")
26
+
27
+ # check if file exists
28
+ if not os.path.exists(path):
29
+ raise ValueError("File not found")
30
+
31
# send file
32
return send_file(path)
33
python/extensions/message_loop_prompts/_30_include_attachments._py
deleted
-40
@@ -1,40 +0,0 @@
1
-# python/extensions/monologue_start/include_attachments.py
2
-from python.helpers.extension import Extension
3
-from python.helpers.attachment_manager import AttachmentManager
4
-from agent import Agent, LoopData
5
-import os
6
-
7
-class IncludeAttachments(Extension):
8
- async def execute(self, loop_data: LoopData = LoopData(), **kwargs):
9
- attachments = self.agent.get_data('attachments') or []
10
- if attachments:
11
- loop_data.attachments = []
12
- file_manager = AttachmentManager(os.path.join(os.getcwd(), 'work_dir'))
13
-
14
- for attachment in attachments:
15
- if os.path.exists(attachment):
16
- filename = os.path.basename(attachment)
17
- file_type = file_manager.get_file_type(filename)
18
-
19
- attachment_html = f'<div class="attachment-item attachment-{file_type}">'
20
- if file_type == 'image':
21
- preview = file_manager.generate_image_preview(attachment)
22
- if preview:
23
- attachment_html += f'<img src="data:image/jpeg;base64,{preview}" alt="{filename}" class="attachment-preview"/>'
24
- else:
25
- # Add placeholder for non-image files
26
- attachment_html += f'<div class="attachment-placeholder">{file_type.upper()}</div>'
27
-
28
- # Add filename and extension badge
29
- ext = file_manager.get_file_extension(filename)
30
- attachment_html += f'''
31
- <div class="attachment-info">
32
- <span class="attachment-name">{filename}</span>
33
- <span class="attachment-badge">{ext}</span>
34
- </div>
35
- </div>'''
36
-
37
- loop_data.attachments.append(attachment_html)
38
-
39
- # Clear attachments after processing
40
- self.agent.set_data('attachments', [])
\ No newline at end of file
python/helpers/files.py
+8
@@ -249,6 +249,14 @@ def get_base_dir():
249
base_dir = os.path.dirname(os.path.abspath(os.path.join(__file__, "../../")))
250
return base_dir
251
252
+def is_in_base_dir(path: str):
253
+ # check if the given path is within the base directory
254
+ base_dir = get_base_dir()
255
+ # normalize paths to handle relative paths and symlinks
256
+ abs_path = os.path.abspath(path)
257
+ # check if the absolute path starts with the base directory
258
+ return os.path.commonpath([abs_path, base_dir]) == base_dir
259
+
260
261
def get_subdirectories(relative_path: str, include: str | list[str] = "*", exclude: str | list[str] | None = None):
262
abs_path = get_abs_path(relative_path)
python/helpers/history.py
+12
-2
@@ -248,6 +248,12 @@ class Bulk(Record):
248
self.summary: str = ""
249
self.records: list[Record] = []
250
251
+ def get_tokens(self):
252
+ if self.summary:
253
+ return tokens.approximate_tokens(self.summary)
254
+ else:
255
+ return sum([r.get_tokens() for r in self.records])
256
+
257
def output(
258
self, human_label: str = "user", ai_label: str = "ai"
259
) -> list[OutputMessage]:
@@ -402,7 +408,8 @@ class History(Record):
408
await bulk.summarize()
409
self.bulks.append(bulk)
410
self.topics.remove(topic)
405
- return True
411
+ return True
412
+ return False
413
414
async def compress_bulks(self):
415
# merge bulks if possible
@@ -410,11 +417,14 @@ class History(Record):
417
# remove oldest bulk if necessary
418
if not compressed:
419
self.bulks.pop(0)
420
+ return True
421
return compressed
422
423
async def merge_bulks_by(self, count: int):
416
- if len(self.bulks) > 0:
424
+ # if bulks is empty, return False
425
+ if len(self.bulks) == 0:
426
return False
427
+ # merge bulks in groups of count, even if there are fewer than count
428
bulks = await asyncio.gather(
429
*[
430
self.merge_bulks(self.bulks[i : i + count])