Fix e2b websocket issue

pyproject.toml

@@ -35,11 +35,12 @@ rich = "^13.7.1"
 langsmith = "^0.1.58"
 ipykernel = "^6.29.4"
 e2b = "^0.17.1"
-e2b-code-interpreter = "0.0.11a2"
+e2b-code-interpreter = "0.0.11a17"
 tenacity = "^8.3.0"
 pillow-heif = "^0.16.0"
 pytube = "15.0.0"
 anthropic = "^0.31.0"
+pydantic = "2.7.4"
 
 [tool.poetry.group.dev.dependencies]
 autoflake = "1.*"

tests/unit/tools/test_tools.py

@@ -0,0 +1,31 @@
+# Generated by CodiumAI
+from pathlib import Path
+
+import numpy as np
+
+from vision_agent.tools.tools import save_video
+
+
+class TestSaveVideo:
+ def test_saves_frames_without_output_path(self):
+ frames = [
+ np.random.randint(0, 256, (480, 640, 3), dtype=np.uint8) for _ in range(10)
+ ]
+ output_path = save_video(frames)
+ assert Path(output_path).exists()
+
+ def test_saves_frames_with_output_path(self, tmp_path):
+ frames = [
+ np.random.randint(0, 256, (480, 640, 3), dtype=np.uint8) for _ in range(10)
+ ]
+ video_output_path = str(tmp_path / "output.mp4")
+ output_path = save_video(frames, video_output_path)
+
+ assert output_path == video_output_path
+ assert Path(output_path).exists()
+
+ # Handles an empty list of frames gracefully
+ def test_handles_empty_frames_list(self):
+ frames = []
+ output_path = save_video(frames)
+ assert Path(output_path).exists()

vision_agent/tools/tools.py

@@ -11,7 +11,6 @@
 import numpy as np
 import pandas as pd
 import requests
-from moviepy.editor import ImageSequenceClip
 from PIL import Image, ImageDraw, ImageFont
 from pillow_heif import register_heif_opener # type: ignore
 from pytube import YouTube # type: ignore
@@ -1042,15 +1041,21 @@ def save_video(
  if fps <= 0:
  _LOGGER.warning(f"Invalid fps value: {fps}. Setting fps to 4 (default value).")
  fps = 4
- with ImageSequenceClip(frames, fps=fps) as video:
- if output_video_path:
- f = open(output_video_path, "wb")
- else:
- f = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) # type: ignore
- video.write_videofile(f.name, codec="libx264")
- f.close()
- _save_video_to_result(f.name)
- return f.name
+
+ if not output_video_path:
+ output_video_path = tempfile.NamedTemporaryFile(
+ suffix=".mp4", delete=False
+ ).name
+
+ height, width, layers = frames[0].shape if frames else (0, 0, 0)
+ fourcc = cv2.VideoWriter_fourcc(*"mp4v") # type: ignore
+ video = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
+ for frame in frames:
+ video.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+ video.release()
+
+ _save_video_to_result(output_video_path)
+ return output_video_path
 
 
 def _save_video_to_result(video_uri: str) -> None: