From 7d9cff3b1d72192935663215c77c7fa47646ec9e Mon Sep 17 00:00:00 2001
From: Ryan Marten <ryanmarten2000@gmail.com>
Date: Tue, 17 Dec 2024 19:43:30 -0800
Subject: [PATCH] fix types

---
 src/bespokelabs/curator/dataset.py                           | 2 +-
 .../request_processor/base_batch_request_processor.py        | 4 ++--
 .../request_processor/base_online_request_processor.py       | 2 +-
 .../request_processor/litellm_online_request_processor.py    | 2 +-
 .../request_processor/openai_online_request_processor.py     | 2 +-
 .../curator/status_tracker/online_status_tracker.py          | 5 ++++-
 6 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/bespokelabs/curator/dataset.py b/src/bespokelabs/curator/dataset.py
index 180be4c6..8f0d6920 100644
--- a/src/bespokelabs/curator/dataset.py
+++ b/src/bespokelabs/curator/dataset.py
@@ -8,7 +8,7 @@
 from pydantic import BaseModel
 
 from bespokelabs.curator.llm.prompt_formatter import PromptFormatter
-from bespokelabs.curator.request_processor.generic_response import GenericResponse
+from bespokelabs.curator.types.generic_response import GenericResponse
 
 T = TypeVar("T")
 
diff --git a/src/bespokelabs/curator/request_processor/base_batch_request_processor.py b/src/bespokelabs/curator/request_processor/base_batch_request_processor.py
index 0b97162a..d2f2db27 100644
--- a/src/bespokelabs/curator/request_processor/base_batch_request_processor.py
+++ b/src/bespokelabs/curator/request_processor/base_batch_request_processor.py
@@ -5,7 +5,7 @@
 import litellm
 
 from bespokelabs.curator.dataset import Dataset
-from bespokelabs.curator.prompter.prompt_formatter import PromptFormatter
+from bespokelabs.curator.llm.prompt_formatter import PromptFormatter
 from bespokelabs.curator.request_processor.base_request_processor import (
     parse_response_message,
     BaseRequestProcessor,
@@ -121,7 +121,7 @@ def requests_from_generic_request_file(self, request_file: str) -> list[dict]:
         return api_specific_requests
 
     def generic_response_file_from_responses(
-        self, responses: str, batch: BetaMessageBatch, response_file: str
+        self, responses: str, batch: GenericBatchObject, response_file: str
     ) -> str | None:
         """Processes API-specific responses and creates a generic response file.
 
diff --git a/src/bespokelabs/curator/request_processor/base_online_request_processor.py b/src/bespokelabs/curator/request_processor/base_online_request_processor.py
index bdcff800..c48bb734 100644
--- a/src/bespokelabs/curator/request_processor/base_online_request_processor.py
+++ b/src/bespokelabs/curator/request_processor/base_online_request_processor.py
@@ -13,7 +13,7 @@
 
 from bespokelabs.curator.dataset import Dataset
 from bespokelabs.curator.request_processor.base_request_processor import BaseRequestProcessor
-from bespokelabs.curator.prompter.prompt_formatter import PromptFormatter
+from bespokelabs.curator.llm.prompt_formatter import PromptFormatter
 from bespokelabs.curator.status_tracker.online_status_tracker import OnlineStatusTracker
 from bespokelabs.curator.types.generic_request import GenericRequest
 from bespokelabs.curator.request_processor.event_loop import run_in_event_loop
diff --git a/src/bespokelabs/curator/request_processor/litellm_online_request_processor.py b/src/bespokelabs/curator/request_processor/litellm_online_request_processor.py
index 44ab4443..7ded6dd5 100644
--- a/src/bespokelabs/curator/request_processor/litellm_online_request_processor.py
+++ b/src/bespokelabs/curator/request_processor/litellm_online_request_processor.py
@@ -11,7 +11,7 @@
 from bespokelabs.curator.request_processor.base_online_request_processor import APIRequest
 from bespokelabs.curator.status_tracker.online_status_tracker import OnlineStatusTracker
 from bespokelabs.curator.types.generic_request import GenericRequest
-from bespokelabs.curator.request_processor.generic_response import TokenUsage, GenericResponse
+from bespokelabs.curator.types.generic_response import TokenUsage, GenericResponse
 from pydantic import BaseModel
 import time
 
diff --git a/src/bespokelabs/curator/request_processor/openai_online_request_processor.py b/src/bespokelabs/curator/request_processor/openai_online_request_processor.py
index 0794685a..77bf3204 100644
--- a/src/bespokelabs/curator/request_processor/openai_online_request_processor.py
+++ b/src/bespokelabs/curator/request_processor/openai_online_request_processor.py
@@ -16,7 +16,7 @@
 from bespokelabs.curator.request_processor.base_online_request_processor import APIRequest
 from bespokelabs.curator.status_tracker.online_status_tracker import OnlineStatusTracker
 from bespokelabs.curator.types.generic_request import GenericRequest
-from bespokelabs.curator.request_processor.generic_response import TokenUsage, GenericResponse
+from bespokelabs.curator.types.generic_response import TokenUsage, GenericResponse
 
 T = TypeVar("T")
 logger = logger = logging.getLogger(__name__)
diff --git a/src/bespokelabs/curator/status_tracker/online_status_tracker.py b/src/bespokelabs/curator/status_tracker/online_status_tracker.py
index 2227181e..812dd2eb 100644
--- a/src/bespokelabs/curator/status_tracker/online_status_tracker.py
+++ b/src/bespokelabs/curator/status_tracker/online_status_tracker.py
@@ -3,10 +3,13 @@
 import tqdm
 
 from dataclasses import dataclass, field
-from bespokelabs.curator.request_processor.base_online_request_processor import SECONDS_TO_PAUSE_ON_RATE_LIMIT
+from bespokelabs.curator.request_processor.base_online_request_processor import (
+    SECONDS_TO_PAUSE_ON_RATE_LIMIT,
+)
 
 logger = logging.getLogger(__name__)
 
+
 @dataclass
 class OnlineStatusTracker:
     """Tracks the status of all requests."""