lm-deluge 0.0.3__py3-none-any.whl

lm_deluge/prompt.py

@@ -0,0 +1,355 @@
+import io
+import json
+import tiktoken
+import xxhash
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal, Sequence
+from lm_deluge.models import APIModel
+from lm_deluge.image import Image
+
+###############################################################################
+# 1. Low-level content blocks – either text or an image                       #
+###############################################################################
+Role = Literal["system", "user", "assistant"]
+
+
+@dataclass(slots=True)
+class Text:
+    text: str
+    type: str = field(init=False, default="text")
+
+    @property
+    def fingerprint(self) -> str:
+        return xxhash.xxh64(self.text.encode()).hexdigest()
+
+    # ── provider-specific emission ────────────────────────────────────────────
+    def oa_chat(self) -> dict | str:  # OpenAI Chat Completions
+        return {"type": "text", "text": self.text}
+
+    def oa_resp(self) -> dict:  # OpenAI *Responses*  (new)
+        return {"type": "input_text", "text": self.text}
+
+    def anthropic(self) -> dict:  # Anthropic Messages
+        return {"type": "text", "text": self.text}
+
+    def gemini(self) -> dict:
+        return {"text": self.text}
+
+
+###############################################################################
+# 2. One conversational turn (role + parts)                                   #
+###############################################################################
+@dataclass(slots=True)
+class Message:
+    role: Role
+    parts: list[Text | Image]
+
+    @property
+    def fingerprint(self) -> str:
+        return self.role + "," + ",".join(part.fingerprint for part in self.parts)
+
+    def add_text(self, content: str) -> "Message":
+        """Append a text block and return self for chaining."""
+        self.parts.append(Text(content))
+        return self
+
+    def add_image(
+        self,
+        data: bytes | str | Path | io.BytesIO,
+        *,
+        media_type: str | None = None,
+        detail: Literal["low", "high", "auto"] = "auto",
+        max_size: int | None = None,
+    ) -> "Message":
+        """
+        Append an image block and return self for chaining.
+
+        If max_size is provided, the image will be resized so that its longer
+        dimension equals max_size, but only if the longer dimension is currently
+        larger than max_size.
+        """
+        img = Image(data, media_type=media_type, detail=detail)
+
+        # Resize if max_size is provided
+        if max_size is not None:
+            img.resize(max_size)
+
+        self.parts.append(img)
+        return self
+
+    # -------- convenient constructors --------
+    @classmethod
+    def user(
+        cls,
+        text: str | None = None,
+        *,
+        image: str | bytes | Path | io.BytesIO | None = None,
+    ) -> "Message":
+        res = cls("user", [])
+        if text is not None:
+            res.add_text(text)
+        if image is not None:
+            res.add_image(image)
+        return res
+
+    @classmethod
+    def system(cls, text: str | None = None) -> "Message":
+        res = cls("system", [])
+        if text is not None:
+            res.add_text(text)
+        return res
+
+    @classmethod
+    def ai(cls, text: str | None = None) -> "Message":
+        res = cls("assistant", [])
+        if text is not None:
+            res.add_text(text)
+        return res
+
+    # ──── provider-specific constructors ───
+    @classmethod
+    def from_oa(cls, msg: dict):
+        role = (
+            "system"
+            if msg["role"] in ["developer", "system"]
+            else ("user" if msg["role"] == "user" else "assistant")
+        )
+        parts: Sequence[Text | Image] = []
+        content = msg["content"]
+        if isinstance(content, str):
+            parts.append(Text(content))
+        else:
+            for item in content:
+                if item["type"] == "text":
+                    parts.append(Text(item["text"]))
+                elif item["type"] == "image_url":
+                    parts.append(Image(data=item["image_url"]["url"]))
+        return cls(role, parts)
+
+    @classmethod
+    def from_oa_resp(cls, msg: dict):
+        raise NotImplementedError("not implemented")
+
+    @classmethod
+    def from_anthropic(cls, msg: dict):
+        pass
+
+    # ───── provider-specific emission ─────
+    def oa_chat(self) -> dict:
+        content = []
+        for p in self.parts:
+            content.append(p.oa_chat())
+        return {"role": self.role, "content": content}
+
+    def oa_resp(self) -> dict:
+        content = [p.oa_resp() for p in self.parts]
+        return {"role": self.role, "content": content}
+
+    def anthropic(self) -> dict:
+        # Anthropic: system message is *not* in the list
+        if self.role == "system":
+            raise ValueError("Anthropic keeps system outside message list")
+        content = [p.anthropic() for p in self.parts]
+        # Shortcut: single text becomes a bare string
+        if len(content) == 1 and content[0]["type"] == "text":
+            content = content[0]["text"]
+        return {"role": self.role, "content": content}
+
+    def gemini(self) -> dict:
+        parts = [p.gemini() for p in self.parts]
+        # Shortcut: single text becomes a bare string
+        role = "user" if self.role == "user" else "model"
+        return {"role": role, "parts": parts}
+
+
+###############################################################################
+# 3. A whole conversation (ordered list of messages)                          #
+###############################################################################
+
+
+@dataclass(slots=True)
+class Conversation:
+    messages: list[Message] = field(default_factory=list)
+
+    # ── convenience shorthands ------------------------------------------------
+    @classmethod
+    def system(cls, text: str) -> "Conversation":
+        return cls([Message.system(text)])
+
+    @classmethod
+    def user(
+        cls, text: str, *, image: bytes | str | Path | None = None
+    ) -> "Conversation":
+        msg = (
+            Message.user(text) if image is None else Message.user(text).add_image(image)
+        )
+        return cls([msg])
+
+    @classmethod
+    def from_openai(cls, messages: list[dict]):
+        """Compatibility with openai-formatted messages"""
+        pass
+
+    @classmethod
+    def from_anthropic(cls, messages: list[dict], system: str | None = None):
+        """Compatibility with anthropic-formatted messages"""
+        pass
+
+    # fluent additions
+    def add(self, msg: Message) -> "Conversation":
+        self.messages.append(msg)
+        return self
+
+    # ── conversions -----------------------------------------------------------
+    def to_openai(self) -> list[dict]:
+        return [m.oa_chat() for m in self.messages]
+
+    def to_openai_responses(self) -> dict:
+        # OpenAI Responses = single “input” array, role must be user/assistant
+        return {"input": [m.oa_resp() for m in self.messages if m.role != "system"]}
+
+    def to_anthropic(self) -> tuple[str | None, list[dict]]:
+        system_msg = next(
+            (
+                m.parts[0].text
+                for m in self.messages
+                if m.role == "system" and isinstance(m.parts[0], Text)
+            ),
+            None,
+        )
+        other = [m.anthropic() for m in self.messages if m.role != "system"]
+        return system_msg, other
+
+    def to_gemini(self) -> tuple[str | None, list[dict]]:
+        system_msg = next(
+            (
+                m.parts[0].text
+                for m in self.messages
+                if m.role == "system" and isinstance(m.parts[0], Text)
+            ),
+            None,
+        )
+        other = [m.gemini() for m in self.messages if m.role != "system"]
+        return system_msg, other
+
+    def to_cohere(self) -> list[dict]:
+        messages = []
+        for m in self.messages:
+            if len(m.parts) > 1:
+                raise ValueError("Cohere does not support multi-part messages")
+            if isinstance(m.parts[0], Image):
+                raise ValueError("Cohere does not support images")
+            messages.append({"role": m.role, "text": m.parts[0].text})
+        return messages
+
+    # ── misc helpers ----------------------------------------------------------
+    _tok = tiktoken.encoding_for_model("gpt-4")
+
+    def count_tokens(self, max_new_tokens: int = 0, img_tokens: int = 85) -> int:
+        n = max_new_tokens
+        for m in self.messages:
+            for p in m.parts:
+                if isinstance(p, Text):
+                    n += len(self._tok.encode(p.text))
+                else:  # Image – crude flat cost per image
+                    n += img_tokens
+
+        # very rough BOS/EOS padding
+        return n + 6 * len(self.messages)
+
+    def dry_run(self, model_name: str, max_new_tokens: int):
+        model_obj = APIModel.from_registry(model_name)
+        if model_obj.api_spec == "openai":
+            image_tokens = 85
+        elif model_obj.api_spec == "anthropic":
+            image_tokens = 1_200
+        else:
+            image_tokens = 0
+        input_tokens = self.count_tokens(0, image_tokens)
+        output_tokens = max_new_tokens
+
+        min_cost, max_cost = None, None
+        if model_obj.input_cost and model_obj.output_cost:
+            min_cost = model_obj.input_cost * input_tokens / 1e6
+            max_cost = min_cost + model_obj.output_cost * output_tokens / 1e6
+
+        return input_tokens, output_tokens, min_cost, max_cost
+
+    @property
+    def fingerprint(self) -> str:
+        hasher = xxhash.xxh64()
+        hasher.update(json.dumps([m.fingerprint for m in self.messages]).encode())
+        return hasher.hexdigest()
+
+    def to_log(self) -> dict:
+        """
+        Return a JSON-serialisable dict that fully captures the conversation.
+        """
+        serialized: list[dict] = []
+
+        for msg in self.messages:
+            content_blocks: list[dict] = []
+            for p in msg.parts:
+                if isinstance(p, Text):
+                    content_blocks.append({"type": "text", "text": p.text})
+                else:  # Image – redact the bytes, keep a hint
+                    w, h = p.size
+                    content_blocks.append(
+                        {"type": "image", "tag": f"<Image ({w}×{h})>"}
+                    )
+            serialized.append({"role": msg.role, "content": content_blocks})
+
+        return {"messages": serialized}
+
+    @classmethod
+    def from_log(cls, payload: dict) -> "Conversation":
+        """Re-hydrate a Conversation previously produced by `to_log()`."""
+        msgs: list[Message] = []
+
+        for m in payload.get("messages", []):
+            role: Role = m["role"]  # 'system' | 'user' | 'assistant'
+            parts: list[Text | Image] = []
+
+            for p in m["content"]:
+                if p["type"] == "text":
+                    parts.append(Text(p["text"]))
+                elif p["type"] == "image":
+                    # We only stored a placeholder tag, so keep that placeholder.
+                    # You could raise instead if real image bytes are required.
+                    parts.append(Image(p["tag"], detail="low"))
+                else:
+                    raise ValueError(f"Unknown part type {p['type']!r}")
+
+            msgs.append(Message(role, parts))
+
+        return cls(msgs)
+
+
+###############################################################################
+# --------------------------------------------------------------------------- #
+# Basic usage examples                                                        #
+# --------------------------------------------------------------------------- #
+
+# 1️⃣  trivial single-turn (text only)  ---------------------------------------
+# conv = Conversation.user("Hi Claude, who won the 2018 World Cup?")
+# client.messages.create(model="claude-3-7-sonnet", **conv.to_anthropic())
+
+# # 2️⃣  system + vision + follow-up for OpenAI Chat Completions  ---------------
+# conv = (
+#     Conversation.system("You are a visual assistant.")
+#     .add(
+#         Message.with_image(
+#             "user",
+#             "What's in this photo?",
+#             Image("boardwalk.jpg", detail="low"),
+#         )
+#     )
+#     .add(Message.text("assistant", "Looks like a lakeside boardwalk."))
+#     .add(Message.text("user", "Great, write a haiku about it."))
+# )
+
+# openai.chat.completions.create(model="gpt-4o-mini", messages=conv.to_openai_chat())
+
+# # 3️⃣  Same conversation sent through new Responses API -----------------------
+# openai.responses.create(model="gpt-4o-mini", **conv.to_openai_responses())

lm_deluge/rerank.py

@@ -0,0 +1,338 @@
+### specific utility for cohere rerank api
+import os
+import aiohttp
+from tqdm.auto import tqdm
+import asyncio
+import time
+from typing import Optional
+from dataclasses import dataclass
+from .tracker import StatusTracker
+
+registry = [
+    "rerank-english-v3.0",
+    "rerank-multilingual-v3.0",
+    "rerank-english-v2.0",
+    "rerank-multilingual-v2.0",
+]
+
+
+class RerankingRequest:
+    def __init__(
+        self,
+        task_id: int,
+        model_name: str,
+        query: str,
+        documents: list[str],
+        top_k: int,
+        attempts_left: int,
+        status_tracker: StatusTracker,
+        retry_queue: asyncio.Queue,
+        request_timeout: int,
+        pbar: Optional[tqdm] = None,
+    ):
+        self.task_id = task_id
+        self.model_name = model_name
+        self.query = query
+        self.documents = documents
+        self.top_k = top_k
+        self.attempts_left = attempts_left
+        self.status_tracker = status_tracker
+        self.retry_queue = retry_queue
+        self.request_timeout = request_timeout
+        self.pbar = pbar
+        self.result = []
+
+    def increment_pbar(self):
+        if self.pbar is not None:
+            self.pbar.update(1)
+
+    def handle_success(self):
+        self.increment_pbar()
+        self.status_tracker.num_tasks_in_progress -= 1
+        self.status_tracker.num_tasks_succeeded += 1
+
+    def handle_error(self):
+        """
+        If create_new_request is True, will create a new API request (so that it
+        has a chance of being sent to a different model). If false, will retry
+        the same request.
+        """
+        last_result: RerankingResponse = self.result[-1]
+        error_to_print = (
+            f"Error on task {self.task_id}, Code: {last_result.status_code}, "
+        )
+        error_to_print += f"Message: {last_result.error_message}."
+        print(error_to_print)
+        if self.attempts_left > 0:
+            self.attempts_left -= 1
+            self.retry_queue.put_nowait(self)
+            return
+        else:
+            print(f"Task {self.task_id} out of tries.")
+            self.status_tracker.num_tasks_in_progress -= 1
+            self.status_tracker.num_tasks_failed += 1
+
+    async def handle_response(self, response: aiohttp.ClientResponse):
+        try:
+            if response.status == 200:
+                result = await response.json()
+                # TODO: add cost calculation
+                return RerankingResponse(
+                    id=self.task_id,
+                    status_code=response.status,
+                    is_error=False,
+                    error_message=None,
+                    query=self.query,
+                    documents=self.documents,
+                    top_k_indices=[doc["index"] for doc in result["results"]],
+                    top_k_scores=[doc["relevance_score"] for doc in result["results"]],
+                )
+            else:
+                error_msg = await response.text()
+                return RerankingResponse(
+                    id=self.task_id,
+                    status_code=response.status,
+                    is_error=True,
+                    error_message=error_msg,
+                    query=self.query,
+                    documents=[],
+                    top_k_indices=[],
+                    top_k_scores=[],
+                )
+        except Exception as e:
+            return RerankingResponse(
+                id=self.task_id,
+                status_code=response.status,
+                is_error=True,
+                error_message=str(e),
+                query=self.query,
+                documents=[],
+                top_k_indices=[],
+                top_k_scores=[],
+            )
+
+    async def call_api(self):
+        url = "https://api.cohere.com/v1/rerank"
+        headers = {
+            "accept": "application/json",
+            "content-type": "application/json",
+            "Authorization": f"Bearer {os.environ.get('COHERE_API_KEY')}",
+        }
+        data = {
+            "model": self.model_name,
+            "query": self.query,
+            "top_n": self.top_k,
+            "documents": self.documents,
+        }
+        try:
+            self.status_tracker.total_requests += 1
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                    url, headers=headers, json=data, timeout=self.request_timeout
+                ) as response:
+                    # print("got response!!")
+                    response_obj: RerankingResponse = await self.handle_response(
+                        response
+                    )
+            self.result.append(response_obj)
+            if response_obj.is_error:
+                self.handle_error()
+            else:
+                self.handle_success()
+
+        except asyncio.TimeoutError:
+            self.result.append(
+                RerankingResponse(
+                    id=self.task_id,
+                    status_code=None,
+                    is_error=True,
+                    error_message="Timeout",
+                    query=self.query,
+                    documents=[],
+                    top_k_indices=[],
+                    top_k_scores=[],
+                )
+            )
+            self.handle_error()
+
+        except Exception as e:
+            self.result.append(
+                RerankingResponse(
+                    id=self.task_id,
+                    status_code=None,
+                    is_error=True,
+                    error_message=f"Unexpected {type(e).__name__}: {str(e) or 'No message.'}",
+                    query=self.query,
+                    documents=[],
+                    top_k_indices=[],
+                    top_k_scores=[],
+                )
+            )
+            self.handle_error()
+
+
+@dataclass
+class RerankingResponse:
+    id: int
+    status_code: int | None
+    is_error: bool
+    error_message: Optional[str]
+    query: str
+    documents: list[str]
+    top_k_indices: list[int]
+    top_k_scores: list[float]
+
+    @property
+    def ranked_documents(self):
+        return [self.documents[i] for i in self.top_k_indices]
+
+
+async def rerank_parallel_async(
+    queries: list[str],
+    docs: list[list[str]],  # one list per query
+    top_k: int = 3,
+    model: str = "rerank-english-v3.0",
+    max_attempts: int = 5,
+    max_requests_per_minute: int = 4_000,
+    max_concurrent_requests: int = 500,
+    request_timeout: int = 10,
+    progress_bar: Optional[tqdm] = None,
+):
+    """Processes rerank requests in parallel, throttling to stay under rate limits."""
+    ids = range(len(queries))
+    # constants
+    seconds_to_pause_after_rate_limit_error = 5
+    seconds_to_sleep_each_loop = 0.003  # so concurrent tasks can run
+
+    # initialize trackers
+    retry_queue = asyncio.Queue()
+    status_tracker = StatusTracker()
+    next_request = None  # variable to hold the next request to call
+
+    # initialize available capacity counts
+    # throttle over a 1 second window rather than minute,
+    # since some models limit RPS rather than RPM
+    available_request_capacity = max_requests_per_minute
+    last_update_time = time.time()
+    last_pbar_update_time = time.time()
+
+    # initialize flags
+    prompts_not_finished = True
+    prompts_iter = iter(zip(ids, queries, docs))
+    results: list[RerankingRequest] = []
+
+    while True:
+        # get next request (if one is not already waiting for capacity)
+        if next_request is None:
+            if not retry_queue.empty():
+                next_request = retry_queue.get_nowait()
+                print(f"Retrying request {next_request.task_id}.")
+            elif prompts_not_finished:
+                try:
+                    # get new request
+                    req_id, req_query, req_docs = next(prompts_iter)
+                    next_request = RerankingRequest(
+                        task_id=req_id,
+                        model_name=model,
+                        query=req_query,
+                        documents=req_docs,
+                        top_k=top_k,
+                        attempts_left=max_attempts,
+                        status_tracker=status_tracker,
+                        retry_queue=retry_queue,
+                        request_timeout=request_timeout,
+                        pbar=progress_bar,
+                    )
+                    status_tracker.num_tasks_started += 1
+                    status_tracker.num_tasks_in_progress += 1
+                    results.append(next_request)
+
+                except StopIteration:
+                    prompts_not_finished = False
+                    print("API requests finished, only retries remain.")
+
+        # update available capacity
+        current_time = time.time()
+        seconds_since_update = current_time - last_update_time
+        available_request_capacity = min(
+            available_request_capacity
+            + max_requests_per_minute * seconds_since_update / 60.0,
+            max_requests_per_minute,
+        )
+        last_update_time = current_time
+
+        # update pbar status
+        if progress_bar:
+            if current_time - last_pbar_update_time > 1:
+                last_pbar_update_time = current_time
+                progress_bar.set_postfix(
+                    {
+                        "Request Capacity": f"{available_request_capacity:.1f}",
+                        "Requests in Progress": status_tracker.num_tasks_in_progress,
+                    }
+                )
+
+        # if enough capacity available, call API
+        if next_request:
+            if (
+                available_request_capacity >= 1
+                and status_tracker.num_tasks_in_progress < max_concurrent_requests
+            ):
+                # update counters
+                available_request_capacity -= 1
+                next_request.attempts_left -= 1
+
+                # call API
+                asyncio.create_task(next_request.call_api())
+                next_request = None  # reset next_request to empty
+
+        # if all tasks are finished, break
+        if status_tracker.num_tasks_in_progress == 0:
+            break
+
+        # main loop sleeps briefly so concurrent tasks can run
+        await asyncio.sleep(seconds_to_sleep_each_loop)
+
+        # if a rate limit error was hit recently, pause to cool down
+        seconds_since_rate_limit_error = (
+            time.time() - status_tracker.time_of_last_rate_limit_error
+        )
+        if seconds_since_rate_limit_error < seconds_to_pause_after_rate_limit_error:
+            remaining_seconds_to_pause = (
+                seconds_to_pause_after_rate_limit_error - seconds_since_rate_limit_error
+            )
+            await asyncio.sleep(remaining_seconds_to_pause)
+            # ^e.g., if pause is 15 seconds and final limit was hit 5 seconds ago
+            print(
+                f"Pausing to cool down until {time.ctime(status_tracker.time_of_last_rate_limit_error + seconds_to_pause_after_rate_limit_error)}"
+            )
+
+    # after finishing, log final status
+    if status_tracker.num_tasks_failed > 0:
+        print(
+            f"{status_tracker.num_tasks_failed} / {status_tracker.num_tasks_started} requests failed."
+        )
+    if status_tracker.num_rate_limit_errors > 0:
+        print(
+            f"{status_tracker.num_rate_limit_errors} rate limit errors received. Consider running at a lower rate."
+        )
+
+    print(
+        f"After processing, got {len(results)} results for {len(ids)} inputs. Removing duplicates."
+    )
+
+    # deduplicate results by id
+    deduplicated = {}
+    for request in results:
+        if request.task_id not in deduplicated:
+            deduplicated[request.task_id] = request.result[-1]
+        else:
+            current_response: RerankingResponse = deduplicated[request.task_id]
+            # only replace if the current request has no top_k_indexes and the new one does
+            if request.result[-1].top_k_indices and not current_response.top_k_indices:
+                deduplicated[request.task_id] = request.result[-1]
+
+    output = list(deduplicated.values())
+    print(f"Returning {len(output)} unique results.")
+
+    return output