edison-client 0.10.1.dev352__tar.gz → 0.11.1.dev194__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 (55) hide show
  1. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/AGENTS.md +2 -0
  2. {edison_client-0.10.1.dev352/src/edison_client.egg-info → edison_client-0.11.1.dev194}/PKG-INFO +2 -2
  3. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/pyproject.toml +1 -1
  4. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/clients/chat_methods.py +102 -0
  5. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/clients/data_storage_methods.py +79 -34
  6. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/clients/job_client.py +19 -3
  7. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/models/app.py +28 -1
  8. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/models/chat.py +21 -3
  9. edison_client-0.11.1.dev194/src/edison_client/version.py +24 -0
  10. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194/src/edison_client.egg-info}/PKG-INFO +2 -2
  11. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client.egg-info/SOURCES.txt +1 -0
  12. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client.egg-info/requires.txt +1 -1
  13. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/client_e2e_uploaded_file_trajectory_test.py +3 -3
  14. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_chat_methods.py +73 -3
  15. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_data_storage_e2e.py +1 -1
  16. edison_client-0.11.1.dev194/tests/test_docker_container_config.py +28 -0
  17. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_rest.py +13 -13
  18. edison_client-0.10.1.dev352/src/edison_client/version.py +0 -34
  19. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/.gitignore +0 -0
  20. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/LICENSE +0 -0
  21. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/README.md +0 -0
  22. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/data_storage.md +0 -0
  23. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/docs/__init__.py +0 -0
  24. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/docs/client_notebook.ipynb +0 -0
  25. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/docs/data_storage_service.puml +0 -0
  26. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/docs/dss_x_finch.png +0 -0
  27. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/setup.cfg +0 -0
  28. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/__init__.py +0 -0
  29. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/clients/__init__.py +0 -0
  30. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/clients/rest_client.py +0 -0
  31. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/models/__init__.py +0 -0
  32. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/models/client.py +0 -0
  33. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/models/data_storage_methods.py +0 -0
  34. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/models/job_event.py +0 -0
  35. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/models/rest.py +0 -0
  36. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/py.typed +0 -0
  37. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/utils/__init__.py +0 -0
  38. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/utils/auth.py +0 -0
  39. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/utils/general.py +0 -0
  40. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/utils/module_utils.py +3 -3
  41. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/utils/monitoring.py +0 -0
  42. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client/utils/world_model_tools.py +0 -0
  43. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client.egg-info/dependency_links.txt +0 -0
  44. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/src/edison_client.egg-info/top_level.txt +0 -0
  45. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/conftest.py +0 -0
  46. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_client.py +0 -0
  47. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_client_close.py +0 -0
  48. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_data/test_file.txt +0 -0
  49. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_data/test_information.txt +0 -0
  50. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_data/test_manifest.yaml +0 -0
  51. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_data_storage_methods.py +0 -0
  52. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_dss_restclient_api_e2e.py +0 -0
  53. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_min_replicas_validation.py +0 -0
  54. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_similarity_search_e2e.py +0 -0
  55. {edison_client-0.10.1.dev352 → edison_client-0.11.1.dev194}/tests/test_world_model_e2e.py +0 -0
@@ -138,8 +138,10 @@ async def main():
138
138
 
139
139
  - `clients/rest_client.py`: Main EdisonClient implementation with sync/async methods
140
140
  - `clients/job_client.py`: Job-specific client functionality
141
+ - `clients/chat_methods.py`: Chat session SDK methods
141
142
  - `clients/data_storage_methods.py`: Data storage integration
142
143
  - `models/app.py`: TaskRequest, TaskResponse, RuntimeConfig models
144
+ - `models/chat.py`: Chat session request/response models
143
145
  - `utils/auth.py`: Authentication utilities
144
146
 
145
147
  **Documentation:**
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: edison-client
3
- Version: 0.10.1.dev352
3
+ Version: 0.11.1.dev194
4
4
  Summary: A client for interacting with endpoints of the Edison Scientific service.
5
5
  Author-email: Edison Scientific technical staff <hello@futurehouse.org>
6
6
  License: Apache License
@@ -221,7 +221,7 @@ Requires-Dist: google-resumable-media
221
221
  Requires-Dist: httpx
222
222
  Requires-Dist: httpx-aiohttp
223
223
  Requires-Dist: ldp>=0.22.0
224
- Requires-Dist: litellm
224
+ Requires-Dist: litellm<=1.82.6
225
225
  Requires-Dist: openai>=2.7
226
226
  Requires-Dist: pydantic
227
227
  Requires-Dist: python-dotenv
@@ -25,7 +25,7 @@ dependencies = [
25
25
  "httpx",
26
26
  "httpx-aiohttp",
27
27
  "ldp>=0.22.0",
28
- "litellm",
28
+ "litellm<=1.82.6", # Pinned: supply chain attack in 1.82.7 and 1.82.8
29
29
  "openai>=2.7", # Pin to keep recent
30
30
  "pydantic",
31
31
  "python-dotenv",
@@ -17,6 +17,8 @@ from edison_client.models.chat import (
17
17
  ChatRecoverPayload,
18
18
  ChatRecoverResponse,
19
19
  ConversationDetail,
20
+ QueueMessagePayload,
21
+ QueueMessageResponse,
20
22
  StoreMessagePayload,
21
23
  )
22
24
  from edison_client.utils.general import (
@@ -501,3 +503,103 @@ class ChatMethods:
501
503
  if isinstance(e, _BASE_CONNECTION_ERRORS):
502
504
  raise
503
505
  raise ChatError(f"An unexpected error occurred: {e!r}") from e
506
+
507
+ # ── queue_chat_message ─────────────────────────────────────────────
508
+
509
+ @retry(
510
+ stop=stop_after_attempt(3),
511
+ wait=wait_exponential(multiplier=1, max=10),
512
+ retry=retry_if_connection_error,
513
+ before_sleep=before_sleep_log(logger, logging.WARNING),
514
+ )
515
+ def queue_chat_message(
516
+ self,
517
+ session_id: UUID,
518
+ project_id: UUID,
519
+ message: str,
520
+ ) -> QueueMessageResponse:
521
+ """Queue a message for injection into a running chat session.
522
+
523
+ The message is placed in the pending queue and picked up by the
524
+ rollout's ``before_transition`` hook on its next step.
525
+
526
+ Args:
527
+ session_id: Target chat session ID.
528
+ project_id: Project the session belongs to.
529
+ message: User message content.
530
+
531
+ Returns:
532
+ QueueMessageResponse with the pending message ID and target ID.
533
+
534
+ Raises:
535
+ ChatError: If there's an error queueing the message.
536
+ """
537
+ try:
538
+ payload = QueueMessagePayload(
539
+ target_id=session_id,
540
+ project_id=project_id,
541
+ message=message,
542
+ )
543
+ response = self.client.post(
544
+ "/v0.1/conversations/queue",
545
+ json=payload.model_dump(mode="json"),
546
+ )
547
+ response.raise_for_status()
548
+ return QueueMessageResponse.model_validate(response.json())
549
+ except HTTPStatusError as e:
550
+ self._handle_chat_http_errors(e, "queueing chat message")
551
+ except ChatError:
552
+ raise
553
+ except Exception as e:
554
+ if isinstance(e, _BASE_CONNECTION_ERRORS):
555
+ raise
556
+ raise ChatError(f"An unexpected error occurred: {e!r}") from e
557
+
558
+ @retry(
559
+ stop=stop_after_attempt(3),
560
+ wait=wait_exponential(multiplier=1, max=10),
561
+ retry=retry_if_connection_error,
562
+ before_sleep=before_sleep_log(logger, logging.WARNING),
563
+ )
564
+ async def aqueue_chat_message(
565
+ self,
566
+ session_id: UUID,
567
+ project_id: UUID,
568
+ message: str,
569
+ ) -> QueueMessageResponse:
570
+ """Queue a message for injection into a running chat session (async).
571
+
572
+ The message is placed in the pending queue and picked up by the
573
+ rollout's ``before_transition`` hook on its next step.
574
+
575
+ Args:
576
+ session_id: Target chat session ID.
577
+ project_id: Project the session belongs to.
578
+ message: User message content.
579
+
580
+ Returns:
581
+ QueueMessageResponse with the pending message ID and target ID.
582
+
583
+ Raises:
584
+ ChatError: If there's an error queueing the message.
585
+ """
586
+ try:
587
+ payload = QueueMessagePayload(
588
+ target_id=session_id,
589
+ project_id=project_id,
590
+ message=message,
591
+ )
592
+ response = await self.async_client.post(
593
+ "/v0.1/conversations/queue",
594
+ json=payload.model_dump(mode="json"),
595
+ )
596
+ response.raise_for_status()
597
+ return QueueMessageResponse.model_validate(response.json())
598
+ except HTTPStatusError as e:
599
+ self._handle_chat_http_errors(e, "queueing chat message")
600
+ except ChatError:
601
+ raise
602
+ except Exception as e:
603
+ if isinstance(e, _BASE_CONNECTION_ERRORS):
604
+ raise
605
+ raise ChatError(f"An unexpected error occurred: {e!r}") from e
@@ -4,6 +4,7 @@ import json
4
4
  import logging
5
5
  import shutil
6
6
  import tempfile
7
+ import time
7
8
  import zipfile
8
9
  from os import PathLike
9
10
  from pathlib import Path
@@ -643,59 +644,103 @@ class DataStorageMethods: # noqa: PLR0904
643
644
  except Exception as e:
644
645
  raise DataStorageError(f"Failed to download from GCS: {e}") from e
645
646
 
646
- def _download_from_gcs(self, signed_url: str, file_name: str | None = None) -> Path:
647
+ def _download_from_gcs(
648
+ self,
649
+ signed_url: str,
650
+ file_name: str | None = None,
651
+ max_retries: int = 5,
652
+ retry_wait_base: int = 2,
653
+ ) -> Path:
647
654
  """Download file from GCS using signed URL and handle unzipping if needed (sync version).
648
655
 
656
+ Resumes interrupted downloads using HTTP Range headers rather than restarting
657
+ from byte 0. Retries up to max_retries times with exponential backoff.
658
+
649
659
  Args:
650
660
  signed_url: The signed URL to download from
651
661
  file_name: The name of the file to download
662
+ max_retries: Maximum number of resume attempts on connection failure
663
+ retry_wait_base: Base seconds for exponential backoff between retries
652
664
  Returns:
653
665
  Path to the downloaded file (or unzipped directory if it was a zip)
654
666
  """
655
667
  file_name = file_name or "downloaded_file"
656
-
668
+ logger.info("Downloading file from GCS with signed URL: %s", signed_url)
657
669
  try:
658
670
  with tempfile.TemporaryDirectory() as temp_dir_str:
659
671
  temp_dir = Path(temp_dir_str)
660
672
  temp_file = temp_dir / file_name
673
+ offset = 0
674
+
675
+ for attempt in range(max_retries + 1):
676
+ headers = {}
677
+ if offset > 0:
678
+ headers["Range"] = f"bytes={offset}-"
679
+
680
+ try:
681
+ with requests_lib.get(
682
+ signed_url, stream=True, timeout=30, headers=headers
683
+ ) as response:
684
+ response.raise_for_status()
685
+
686
+ if offset == 0:
687
+ content_disposition = response.headers.get(
688
+ "content-disposition", ""
689
+ )
690
+ filename = file_name
691
+ if "filename=" in content_disposition:
692
+ filename = content_disposition.split("filename=")[
693
+ -1
694
+ ].strip('"')
695
+ if filename != file_name:
696
+ temp_file = temp_dir / filename
697
+
698
+ mode = "ab" if offset > 0 else "wb"
699
+ with open(temp_file, mode) as f:
700
+ for chunk in response.iter_content(chunk_size=8192):
701
+ if chunk:
702
+ f.write(chunk)
703
+ offset += len(chunk)
704
+
705
+ break # download complete
706
+
707
+ except (
708
+ requests_lib.exceptions.ChunkedEncodingError,
709
+ requests_lib.exceptions.ConnectionError,
710
+ requests_lib.exceptions.Timeout,
711
+ ) as e:
712
+ if attempt == max_retries:
713
+ raise
714
+ wait = retry_wait_base**attempt
715
+ logger.warning(
716
+ f"Download interrupted at {offset:,} bytes "
717
+ f"(attempt {attempt + 1}/{max_retries}), "
718
+ f"resuming in {wait}s: {e}"
719
+ )
720
+ time.sleep(wait)
661
721
 
662
- with requests_lib.get(signed_url, stream=True, timeout=30) as response:
663
- response.raise_for_status()
664
-
665
- content_disposition = response.headers.get(
666
- "content-disposition", ""
667
- )
668
- filename = file_name
669
- if "filename=" in content_disposition:
670
- filename = content_disposition.split("filename=")[-1].strip('"')
671
-
672
- if filename != file_name:
673
- temp_file = temp_dir / filename
674
-
675
- with open(temp_file, "wb") as f:
676
- f.writelines(response.iter_content(chunk_size=8192))
722
+ logger.debug(
723
+ f"Downloaded file to {temp_file} (size: {temp_file.stat().st_size:,} bytes)"
724
+ )
677
725
 
678
- logger.debug(
679
- f"Downloaded file to {temp_file} (size: {temp_file.stat().st_size:,} bytes)"
680
- )
726
+ if self._is_zip_file(temp_file):
727
+ logger.debug(f"File {temp_file} is a zip file, extracting...")
728
+ extracted_path = self._extract_zip_file(temp_file, temp_dir)
681
729
 
682
- if self._is_zip_file(temp_file):
683
- logger.debug(f"File {temp_file} is a zip file, extracting...")
684
- extracted_path = self._extract_zip_file(temp_file, temp_dir)
730
+ final_temp_dir = Path(tempfile.mkdtemp())
731
+ final_path = final_temp_dir / extracted_path.name
685
732
 
686
- final_temp_dir = Path(tempfile.mkdtemp())
687
- final_path = final_temp_dir / extracted_path.name
733
+ if extracted_path.is_dir():
734
+ shutil.copytree(extracted_path, final_path)
735
+ else:
736
+ shutil.copy2(extracted_path, final_path)
688
737
 
689
- if extracted_path.is_dir():
690
- shutil.copytree(extracted_path, final_path)
691
- else:
692
- shutil.copy2(extracted_path, final_path)
738
+ return final_path
693
739
 
694
- return final_path
695
- final_temp_dir = Path(tempfile.mkdtemp())
696
- final_file = final_temp_dir / temp_file.name
697
- shutil.copy2(temp_file, final_file)
698
- return final_file
740
+ final_temp_dir = Path(tempfile.mkdtemp())
741
+ final_file = final_temp_dir / temp_file.name
742
+ shutil.copy2(temp_file, final_file)
743
+ return final_file
699
744
 
700
745
  except Exception as e:
701
746
  raise DataStorageError(f"Failed to download from GCS: {e}") from e
@@ -1,3 +1,6 @@
1
+ import asyncio
2
+ import gzip
3
+ import json
1
4
  import logging
2
5
  from typing import ClassVar
3
6
  from uuid import UUID, uuid4
@@ -27,6 +30,7 @@ logger = logging.getLogger(__name__)
27
30
 
28
31
  class JobClient:
29
32
  REQUEST_TIMEOUT: ClassVar[float] = 30.0 # sec
33
+ LARGE_PAYLOAD_TIMEOUT: ClassVar[float] = 120.0 # sec, for streaming uploads
30
34
  MAX_RETRY_ATTEMPTS: ClassVar[int] = 3
31
35
  RETRY_MULTIPLIER: ClassVar[int] = 1
32
36
  MAX_RETRY_WAIT: ClassVar[int] = 10
@@ -141,13 +145,25 @@ class JobClient:
141
145
  trajectory_timestep=self.current_timestep,
142
146
  )
143
147
 
148
+ def _serialize_and_compress() -> tuple[bytes, int]:
149
+ json_bytes = json.dumps(data.model_dump(mode="json")).encode("utf-8")
150
+ return gzip.compress(json_bytes), len(json_bytes)
151
+
152
+ compressed, raw_size = await asyncio.to_thread(_serialize_and_compress)
153
+ logger.debug(
154
+ f"Compressed agent state payload: {raw_size} -> {len(compressed)} bytes "
155
+ f"({len(compressed) / raw_size:.1%})",
156
+ )
157
+
144
158
  try:
145
159
  async with httpx_aiohttp.HttpxAiohttpClient(
146
- timeout=self.REQUEST_TIMEOUT
160
+ timeout=self.LARGE_PAYLOAD_TIMEOUT
147
161
  ) as client:
148
162
  url = f"{self.base_uri}/v0.1/trajectories/{self.trajectory_id}/agent-state"
149
163
  headers = {
150
164
  "Authorization": f"Bearer {self.oauth_jwt}",
165
+ "Content-Type": "application/json",
166
+ "Content-Encoding": "gzip",
151
167
  "x-trajectory-id": self.trajectory_id,
152
168
  }
153
169
 
@@ -168,7 +184,7 @@ class JobClient:
168
184
 
169
185
  response = await client.post(
170
186
  url=url,
171
- json=data.model_dump(mode="json"),
187
+ content=compressed,
172
188
  headers=headers,
173
189
  )
174
190
  response.raise_for_status()
@@ -182,7 +198,7 @@ class JobClient:
182
198
  )
183
199
  except httpx.TimeoutException:
184
200
  logger.exception(
185
- f"Timeout while storing agent state after {self.REQUEST_TIMEOUT}s",
201
+ f"Timeout while storing agent state after {self.LARGE_PAYLOAD_TIMEOUT}s",
186
202
  )
187
203
  raise
188
204
  except httpx.NetworkError:
@@ -63,7 +63,9 @@ class JobNames(StrEnum):
63
63
  """Get human-readable descriptions for each job type."""
64
64
  return {
65
65
  cls.LITERATURE: "Paper analysis and literature review",
66
- cls.LITERATURE_HIGH: "High-reasoning mode paper analysis and literature review",
66
+ cls.LITERATURE_HIGH: (
67
+ "High-reasoning mode paper analysis and literature review"
68
+ ),
67
69
  cls.CROW: "Paper analysis and literature review",
68
70
  cls.FALCON: "Paper analysis and literature review",
69
71
  cls.ANALYSIS: "Data analysis with Python/notebooks",
@@ -413,6 +415,10 @@ class DockerContainerConfiguration(BaseModel):
413
415
  default="10Gi",
414
416
  description="Disk space for the agent workspace volume (e.g., '10Gi')",
415
417
  )
418
+ tmp_size: str = Field(
419
+ default="5Gi",
420
+ description="Size limit for the /tmp emptyDir volume (e.g., '5Gi')",
421
+ )
416
422
  gpu_count: int | None = Field(
417
423
  default=None,
418
424
  description="Number of NVIDIA GPUs to allocate. Requires CELERY backend.",
@@ -444,6 +450,19 @@ class DockerContainerConfiguration(BaseModel):
444
450
  raise ValueError("Disk space must not exceed 100Gi")
445
451
  return v
446
452
 
453
+ @field_validator("tmp_size")
454
+ @classmethod
455
+ def validate_tmp_size(cls, v: str) -> str:
456
+ match = re.match(r"^(\d+)Gi$", v)
457
+ if not match:
458
+ raise ValueError("tmp_size must be in Gi format (e.g., '5Gi')")
459
+ value = int(match.group(1))
460
+ if value < 1:
461
+ raise ValueError("tmp_size must be at least 1Gi")
462
+ if value > 50: # noqa: PLR2004
463
+ raise ValueError("tmp_size must not exceed 50Gi")
464
+ return v
465
+
447
466
  @field_validator("memory")
448
467
  @classmethod
449
468
  def validate_memory(cls, v: str) -> str:
@@ -977,6 +996,14 @@ class TaskRequest(BaseModel):
977
996
  default=None,
978
997
  description="Project ID for the /conversations/queue completion notification",
979
998
  )
999
+ tags: list[str] | None = Field(
1000
+ default=None,
1001
+ description="Grouping association tags for a task.",
1002
+ )
1003
+ workspace_id: str | None = Field(
1004
+ default=None,
1005
+ description="JuiceFS workspace identifier for persistent storage across sandbox pods.",
1006
+ )
980
1007
 
981
1008
  @model_validator(mode="after")
982
1009
  def validate_caller_fields(self) -> Self:
@@ -24,7 +24,9 @@ class ChatMessagePayload(BaseModel):
24
24
  message: str = Field(description="User message content")
25
25
  ttl_seconds: int | None = Field(
26
26
  default=None,
27
- description="Optional sandbox TTL override in seconds (60-86400)",
27
+ ge=60,
28
+ le=86400,
29
+ description="Optional sandbox TTL override in seconds",
28
30
  )
29
31
  job_name: str | None = Field(
30
32
  default=None, description="Job name identifying which agent to chat with"
@@ -53,7 +55,7 @@ class ChatRecoverPayload(BaseModel):
53
55
  project_id: UUID = Field(description="Project ID for recovery")
54
56
  job_name: str | None = Field(
55
57
  default=None,
56
- description="Job name for the agent (used to resolve the sandbox template on recovery)",
58
+ description="Job name for the agent (used to resolve the sandbox template on recovery) -- optional will infer recovery template from the default service template (dummy-env)",
57
59
  )
58
60
 
59
61
 
@@ -93,6 +95,7 @@ class ChatConversation(BaseModel):
93
95
  created_at: datetime = Field(
94
96
  description="Timestamp when the conversation was created"
95
97
  )
98
+ # stolen from the service side code, but we should consider a more robust schema
96
99
  messages: list[Any] = Field(description="The messages in this conversation")
97
100
 
98
101
 
@@ -100,7 +103,7 @@ class ChatConversationList(BaseModel):
100
103
  """Response listing conversations."""
101
104
 
102
105
  conversations: list[ChatConversation] = Field(
103
- description="List of conversations grouped by session_id"
106
+ description="List of conversations associated with the provided session_id if any"
104
107
  )
105
108
 
106
109
 
@@ -122,3 +125,18 @@ class StoreMessagePayload(BaseModel):
122
125
  metadata: dict | None = Field(
123
126
  default=None, description="Optional metadata to attach to the message"
124
127
  )
128
+
129
+
130
+ class QueueMessagePayload(BaseModel):
131
+ """Payload for queueing a message to a running rollout's pending queue."""
132
+
133
+ target_id: UUID = Field(description="Session ID (chat) or trajectory ID (run)")
134
+ project_id: UUID = Field(description="Project the target belongs to")
135
+ message: str = Field(description="User message content")
136
+
137
+
138
+ class QueueMessageResponse(BaseModel):
139
+ """Response from the queue endpoint."""
140
+
141
+ id: UUID = Field(description="Pending message ID")
142
+ target_id: UUID = Field(description="Session or trajectory ID")
@@ -0,0 +1,24 @@
1
+ # file generated by vcs-versioning
2
+ # don't change, don't track in version control
3
+ from __future__ import annotations
4
+
5
+ __all__ = [
6
+ "__version__",
7
+ "__version_tuple__",
8
+ "version",
9
+ "version_tuple",
10
+ "__commit_id__",
11
+ "commit_id",
12
+ ]
13
+
14
+ version: str
15
+ __version__: str
16
+ __version_tuple__: tuple[int | str, ...]
17
+ version_tuple: tuple[int | str, ...]
18
+ commit_id: str | None
19
+ __commit_id__: str | None
20
+
21
+ __version__ = version = '0.11.1.dev194'
22
+ __version_tuple__ = version_tuple = (0, 11, 1, 'dev194')
23
+
24
+ __commit_id__ = commit_id = 'gb74b18d01'
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: edison-client
3
- Version: 0.10.1.dev352
3
+ Version: 0.11.1.dev194
4
4
  Summary: A client for interacting with endpoints of the Edison Scientific service.
5
5
  Author-email: Edison Scientific technical staff <hello@futurehouse.org>
6
6
  License: Apache License
@@ -221,7 +221,7 @@ Requires-Dist: google-resumable-media
221
221
  Requires-Dist: httpx
222
222
  Requires-Dist: httpx-aiohttp
223
223
  Requires-Dist: ldp>=0.22.0
224
- Requires-Dist: litellm
224
+ Requires-Dist: litellm<=1.82.6
225
225
  Requires-Dist: openai>=2.7
226
226
  Requires-Dist: pydantic
227
227
  Requires-Dist: python-dotenv
@@ -41,6 +41,7 @@ tests/test_client.py
41
41
  tests/test_client_close.py
42
42
  tests/test_data_storage_e2e.py
43
43
  tests/test_data_storage_methods.py
44
+ tests/test_docker_container_config.py
44
45
  tests/test_dss_restclient_api_e2e.py
45
46
  tests/test_min_replicas_validation.py
46
47
  tests/test_rest.py
@@ -5,7 +5,7 @@ google-resumable-media
5
5
  httpx
6
6
  httpx-aiohttp
7
7
  ldp>=0.22.0
8
- litellm
8
+ litellm<=1.82.6
9
9
  openai>=2.7
10
10
  pydantic
11
11
  python-dotenv
@@ -12,7 +12,7 @@ from edison_client.models.data_storage_methods import RawFetchResponse
12
12
  pytestmark = [pytest.mark.live, pytest.mark.live_agent]
13
13
 
14
14
 
15
- @pytest.mark.timeout(660)
15
+ @pytest.mark.timeout(180)
16
16
  def test_upload_and_run_analysis_task(admin_client: RestClient):
17
17
  """Test uploading a file with prompt, running analysis task, and downloading outputs."""
18
18
  prompt_content = "draw a blue square"
@@ -40,7 +40,7 @@ def test_upload_and_run_analysis_task(admin_client: RestClient):
40
40
  name=JobNames.ANALYSIS, # Use the data analysis crow
41
41
  query="Open the attached file and follow the instructions in it. Create an image as specified.",
42
42
  runtime_config=RuntimeConfig(
43
- timeout=200,
43
+ timeout=180,
44
44
  ),
45
45
  )
46
46
 
@@ -49,7 +49,7 @@ def test_upload_and_run_analysis_task(admin_client: RestClient):
49
49
  files=[data_entry_uri], # Simple file attachment!
50
50
  verbose=True,
51
51
  progress_bar=True,
52
- timeout=600,
52
+ timeout=180,
53
53
  )
54
54
 
55
55
  assert len(results) == 1
@@ -1,5 +1,6 @@
1
1
  import asyncio
2
2
  import time
3
+ from collections.abc import Generator
3
4
  from uuid import UUID, uuid4
4
5
 
5
6
  import pytest
@@ -10,6 +11,7 @@ from edison_client.models.chat import (
10
11
  ChatConversationList,
11
12
  ChatMessageResponse,
12
13
  ConversationDetail,
14
+ QueueMessageResponse,
13
15
  )
14
16
 
15
17
  CHAT_JOB_NAME = "job-futurehouse-dummy-env"
@@ -20,9 +22,11 @@ RETRY_ATTEMPTS = 3
20
22
 
21
23
 
22
24
  @pytest.fixture(name="chat_project_id")
23
- def fixture_chat_project_id(admin_client: RestClient) -> UUID:
24
- """Create a fresh project for chat tests."""
25
- return admin_client.create_project(f"chat-e2e-{uuid4()}")
25
+ def fixture_chat_project_id(admin_client: RestClient) -> Generator[UUID, None, None]:
26
+ """Create a fresh project for chat tests, clean up afterwards."""
27
+ project_id = admin_client.create_project(f"chat-e2e-{uuid4()}")
28
+ yield project_id
29
+ admin_client.delete_project(project_id)
26
30
 
27
31
 
28
32
  class TestSyncChatLifecycle:
@@ -133,6 +137,37 @@ class TestSyncChatLifecycle:
133
137
  )
134
138
  assert "not found" in str(exc_info.value).lower()
135
139
 
140
+ def test_queue_chat_message_bad_project(self, admin_client: RestClient):
141
+ """Queueing a message with a non-existent project_id returns 404."""
142
+ with pytest.raises(ChatError, match="404") as exc_info:
143
+ admin_client.queue_chat_message(
144
+ session_id=uuid4(),
145
+ project_id=uuid4(),
146
+ message="Should fail",
147
+ )
148
+ assert "not found" in str(exc_info.value).lower()
149
+
150
+ @pytest.mark.live_agent
151
+ @pytest.mark.timeout(TEST_TIMEOUT)
152
+ @pytest.mark.flaky(reruns=RETRY_ATTEMPTS)
153
+ def test_queue_chat_message(self, admin_client: RestClient, chat_project_id: UUID):
154
+ """Queue a message into a running chat session (sync)."""
155
+ send_result = admin_client.send_chat_message(
156
+ project_id=chat_project_id,
157
+ message="Setting up for queue test",
158
+ job_name=CHAT_JOB_NAME,
159
+ )
160
+ session_id = send_result.session_id
161
+
162
+ result = admin_client.queue_chat_message(
163
+ session_id=session_id,
164
+ project_id=chat_project_id,
165
+ message="Injected via queue",
166
+ )
167
+ assert isinstance(result, QueueMessageResponse)
168
+ assert result.id is not None
169
+ assert result.target_id == session_id
170
+
136
171
 
137
172
  class TestAsyncChatLifecycle:
138
173
  """Async e2e: asend_chat_message → aget_conversations → aget_conversation → astore_conversation_message."""
@@ -246,3 +281,38 @@ class TestAsyncChatLifecycle:
246
281
  message={"role": "user", "content": "orphan"},
247
282
  )
248
283
  assert "not found" in str(exc_info.value).lower()
284
+
285
+ @pytest.mark.asyncio
286
+ async def test_aqueue_chat_message_bad_project(self, admin_client: RestClient):
287
+ """Queueing a message with a non-existent project_id returns 404 (async)."""
288
+ with pytest.raises(ChatError, match="404") as exc_info:
289
+ await admin_client.aqueue_chat_message(
290
+ session_id=uuid4(),
291
+ project_id=uuid4(),
292
+ message="Should fail",
293
+ )
294
+ assert "not found" in str(exc_info.value).lower()
295
+
296
+ @pytest.mark.live_agent
297
+ @pytest.mark.timeout(TEST_TIMEOUT)
298
+ @pytest.mark.flaky(reruns=RETRY_ATTEMPTS)
299
+ @pytest.mark.asyncio
300
+ async def test_aqueue_chat_message(
301
+ self, admin_client: RestClient, chat_project_id: UUID
302
+ ):
303
+ """Queue a message into a running chat session (async)."""
304
+ send_result = await admin_client.asend_chat_message(
305
+ project_id=chat_project_id,
306
+ message="Setting up for async queue test",
307
+ job_name=CHAT_JOB_NAME,
308
+ )
309
+ session_id = send_result.session_id
310
+
311
+ result = await admin_client.aqueue_chat_message(
312
+ session_id=session_id,
313
+ project_id=chat_project_id,
314
+ message="Injected via async queue",
315
+ )
316
+ assert isinstance(result, QueueMessageResponse)
317
+ assert result.id is not None
318
+ assert result.target_id == session_id
@@ -626,7 +626,7 @@ class TestDataStorageSearch:
626
626
  dataset_id_test_criteria = SearchCriterion(
627
627
  field="dataset_id",
628
628
  operator=SearchOperator.EQUALS,
629
- value="8bb27d71-8ec6-4845-a09e-4eceefb38cba",
629
+ value="328375e2-591e-4cc9-b17b-cccf040c343a",
630
630
  )
631
631
 
632
632
  description_test_criteria = SearchCriterion(
@@ -0,0 +1,28 @@
1
+ import pytest
2
+ from pydantic import ValidationError
3
+
4
+ from edison_client.models.app import DockerContainerConfiguration
5
+
6
+
7
+ def test_tmp_size_default():
8
+ config = DockerContainerConfiguration(cpu="1", memory="2Gi")
9
+ assert config.tmp_size == "5Gi"
10
+
11
+
12
+ def test_tmp_size_custom():
13
+ config = DockerContainerConfiguration(cpu="4", memory="8Gi", tmp_size="20Gi")
14
+ assert config.tmp_size == "20Gi"
15
+
16
+
17
+ @pytest.mark.parametrize(
18
+ ("tmp_size", "error_match"),
19
+ [
20
+ ("0Gi", "tmp_size must be at least 1Gi"),
21
+ ("51Gi", "tmp_size must not exceed 50Gi"),
22
+ ("5MB", "tmp_size must be in Gi format"),
23
+ ("abc", "tmp_size must be in Gi format"),
24
+ ],
25
+ )
26
+ def test_tmp_size_validation(tmp_size, error_match):
27
+ with pytest.raises(ValidationError, match=error_match):
28
+ DockerContainerConfiguration(cpu="4", memory="8Gi", tmp_size=tmp_size)
@@ -95,7 +95,7 @@ def fixture_running_trajectory_id(
95
95
 
96
96
 
97
97
  @pytest.mark.live_agent
98
- @pytest.mark.timeout(300)
98
+ @pytest.mark.timeout(120)
99
99
  @pytest.mark.flaky(reruns=3)
100
100
  def test_futurehouse_dummy_env_crow(admin_client: RestClient, task_req: TaskRequest):
101
101
  task_id = admin_client.create_task(task_req)
@@ -240,7 +240,7 @@ def test_deployment_config_hidden_from_non_fh_admin(pub_client: RestClient):
240
240
 
241
241
 
242
242
  @pytest.mark.live_agent
243
- @pytest.mark.timeout(300)
243
+ @pytest.mark.timeout(240)
244
244
  @pytest.mark.flaky(reruns=3)
245
245
  def test_run_until_done_futurehouse_dummy_env_crow(
246
246
  admin_client: RestClient, task_req: TaskRequest
@@ -254,7 +254,7 @@ def test_run_until_done_futurehouse_dummy_env_crow(
254
254
 
255
255
 
256
256
  @pytest.mark.live_agent
257
- @pytest.mark.timeout(300)
257
+ @pytest.mark.timeout(240)
258
258
  @pytest.mark.flaky(reruns=3)
259
259
  def test_run_until_done_returns_task_response(
260
260
  admin_client: RestClient, task_req: TaskRequest
@@ -274,7 +274,7 @@ def test_run_until_done_returns_task_response(
274
274
 
275
275
 
276
276
  @pytest.mark.live_agent
277
- @pytest.mark.timeout(300)
277
+ @pytest.mark.timeout(240)
278
278
  @pytest.mark.flaky(reruns=3)
279
279
  @pytest.mark.asyncio
280
280
  async def test_arun_until_done_futurehouse_dummy_env_crow(
@@ -289,7 +289,7 @@ async def test_arun_until_done_futurehouse_dummy_env_crow(
289
289
 
290
290
 
291
291
  @pytest.mark.live_agent
292
- @pytest.mark.timeout(300)
292
+ @pytest.mark.timeout(240)
293
293
  @pytest.mark.flaky(reruns=3)
294
294
  @pytest.mark.asyncio
295
295
  async def test_arun_until_done_returns_task_response(
@@ -310,13 +310,13 @@ async def test_arun_until_done_returns_task_response(
310
310
 
311
311
 
312
312
  @pytest.mark.live_agent
313
- @pytest.mark.timeout(500)
313
+ @pytest.mark.timeout(180)
314
314
  def test_continuation_task(admin_client: RestClient):
315
315
  """Test running a continuation task against data-analysis-crow-high."""
316
316
  initial_task = TaskRequest(
317
317
  name=JobNames.ANALYSIS,
318
318
  query="What is 2 + 2? Just give me the number.",
319
- runtime_config=RuntimeConfig(timeout=200),
319
+ runtime_config=RuntimeConfig(timeout=180),
320
320
  )
321
321
  print(initial_task)
322
322
 
@@ -324,7 +324,7 @@ def test_continuation_task(admin_client: RestClient):
324
324
  task_data=initial_task,
325
325
  verbose=True,
326
326
  progress_bar=True,
327
- timeout=220,
327
+ timeout=180,
328
328
  )
329
329
 
330
330
  assert len(initial_results) == 1
@@ -338,7 +338,7 @@ def test_continuation_task(admin_client: RestClient):
338
338
  name=JobNames.ANALYSIS,
339
339
  query="Now multiply that result by 3. Just give me the number.",
340
340
  runtime_config=RuntimeConfig(
341
- timeout=200,
341
+ timeout=180,
342
342
  continued_job_id=initial_result.task_id,
343
343
  ),
344
344
  )
@@ -347,7 +347,7 @@ def test_continuation_task(admin_client: RestClient):
347
347
  task_data=continuation_task,
348
348
  verbose=True,
349
349
  progress_bar=True,
350
- timeout=220,
350
+ timeout=180,
351
351
  )
352
352
 
353
353
  assert len(continuation_results) == 1
@@ -359,7 +359,7 @@ def test_continuation_task(admin_client: RestClient):
359
359
 
360
360
 
361
361
  @pytest.mark.live_agent
362
- @pytest.mark.timeout(300)
362
+ @pytest.mark.timeout(60)
363
363
  @pytest.mark.flaky(reruns=3)
364
364
  @pytest.mark.asyncio
365
365
  async def test_timeout_run_until_done_futurehouse_dummy_env_crow(
@@ -622,7 +622,7 @@ class TestAsyncProjectOperations:
622
622
 
623
623
 
624
624
  @pytest.mark.live_agent
625
- @pytest.mark.timeout(300)
625
+ @pytest.mark.timeout(120)
626
626
  @pytest.mark.flaky(reruns=3)
627
627
  def test_get_tasks_with_project_filter(admin_client: RestClient, task_req: TaskRequest):
628
628
  """Test retrieving trajectories filtered by project_id using real API calls."""
@@ -645,7 +645,7 @@ def test_get_tasks_with_project_filter(admin_client: RestClient, task_req: TaskR
645
645
 
646
646
 
647
647
  @pytest.mark.live_agent
648
- @pytest.mark.timeout(300)
648
+ @pytest.mark.timeout(120)
649
649
  @pytest.mark.flaky(reruns=3)
650
650
  @pytest.mark.asyncio
651
651
  async def test_aget_tasks_with_project_filter(
@@ -1,34 +0,0 @@
1
- # file generated by setuptools-scm
2
- # don't change, don't track in version control
3
-
4
- __all__ = [
5
- "__version__",
6
- "__version_tuple__",
7
- "version",
8
- "version_tuple",
9
- "__commit_id__",
10
- "commit_id",
11
- ]
12
-
13
- TYPE_CHECKING = False
14
- if TYPE_CHECKING:
15
- from typing import Tuple
16
- from typing import Union
17
-
18
- VERSION_TUPLE = Tuple[Union[int, str], ...]
19
- COMMIT_ID = Union[str, None]
20
- else:
21
- VERSION_TUPLE = object
22
- COMMIT_ID = object
23
-
24
- version: str
25
- __version__: str
26
- __version_tuple__: VERSION_TUPLE
27
- version_tuple: VERSION_TUPLE
28
- commit_id: COMMIT_ID
29
- __commit_id__: COMMIT_ID
30
-
31
- __version__ = version = '0.10.1.dev352'
32
- __version_tuple__ = version_tuple = (0, 10, 1, 'dev352')
33
-
34
- __commit_id__ = commit_id = 'gbbd65e731'
@@ -75,14 +75,14 @@ def fetch_environment_function_docstring(
75
75
  directory (Path): The base directory containing the environment files.
76
76
  function_name (str): The name of the function to retrieve the docstring for.
77
77
 
78
+ Returns:
79
+ str | None: The docstring of the specified function, or None if not found.
80
+
78
81
  Raises:
79
82
  ValueError: If multiple classes with the same name are found in different files
80
83
  (making the intended class ambiguous), an error is raised requiring
81
84
  disambiguation.
82
85
 
83
- Returns:
84
- str | None: The docstring of the specified function, or None if not found.
85
-
86
86
  """
87
87
  parts = environment_name.split(".")
88
88
  class_name = parts[-1]