cortexgrid-infer 0.3.6__tar.gz → 0.3.7__tar.gz

{cortexgrid_infer-0.3.6 → cortexgrid_infer-0.3.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid-infer
-Version: 0.3.6
+Version: 0.3.7
 Summary: Serve apps per task (text-to-text, text-to-image, image-to-mesh) and importers per source (HuggingFace) for running models on cortexgrid
 Project-URL: Homepage, https://github.com/robodatalab/cortexgrid-infer
 Project-URL: Repository, https://github.com/robodatalab/cortexgrid-infer

{cortexgrid_infer-0.3.6 → cortexgrid_infer-0.3.7}/cortexgrid_infer/core.py

@@ -37,6 +37,27 @@ class CompletionChunk:
     def has_tool_calls(self) -> bool:
         return bool(self.tool_calls)
 
+    @classmethod
+    def from_wire(
+        cls, answered: dict[str, Any], tools_by_name: dict[str, Callable[..., Any]]
+    ) -> CompletionChunk:
+        tool_calls = []
+        for called in answered["tool_calls"]:
+            tool = tools_by_name.get(called["name"])
+            if tool is None:
+                continue
+            bound = partial(tool, **called["arguments"])
+            tool_call = ToolCall(
+                id=called["id"], name=called["name"], arguments=called["arguments"], _func=bound
+            )
+            tool_calls.append(tool_call)
+        return cls(
+            content=answered["content"],
+            thinking=answered["thinking"],
+            tool_calls=tool_calls,
+            finish_reason=answered["finish_reason"],
+        )
+
 
 Message = dict[str, Any]
 ToolSpec = dict[str, Any]

cortexgrid_infer-0.3.7/cortexgrid_infer/protocols/completion.py

@@ -0,0 +1,92 @@
+"""The wire protocol between a cortexgrid-served completion app and its client.
+
+A completion serve app streams plain text from `POST /complete`, with tool calls
+inlined in the text as `<tool_call>{"name": ..., "arguments": {...}}</tool_call>`.
+That is the format an instruct-tuned causal LM emits on its own, so a serve app
+wrapping one forwards its tokens untouched; a serve app wrapping a provider with
+structured tool calls encodes them into the same form with `encode_tool_call`.
+
+Either way the client is the same `ServedCompletingModel`, which streams text
+back as `CompletionChunk`s and reassembles the tool calls.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+import json
+from typing import Any, Sequence
+
+import httpx
+
+from cortexgrid_infer.core import CompletingModel, CompletionChunk, Message, Tool
+from cortexgrid_infer.utils import build_tool_map, normalize_tools
+
+@dataclass
+class ServedCompletingModel(CompletingModel):
+    """Client for a completion app cortexgrid has deployed at `url`.
+
+    Provider-agnostic: every completion serve app speaks the protocol above, so
+    which model is behind it only shows up in `model_id`."""
+
+    url: str
+    model_id: str
+
+    @property
+    def name(self) -> str:
+        return self.model_id
+
+    async def loglikelihoods(
+        self, messages: list[Message], continuations: list[str]
+    ) -> list[float]:
+        body = {"messages": messages, "continuations": continuations}
+        loglikelihoods_url = f"{self.url}/loglikelihoods"
+        async with httpx.AsyncClient(timeout=None) as client:
+            response = await client.post(loglikelihoods_url, json=body)
+        response.raise_for_status()
+        answered = response.json()
+        return answered["loglikelihoods"]
+
+    async def complete(
+        self,
+        messages: list[Message],
+        tools: Sequence[Tool] | None = None,
+        max_new_tokens: int = 16 * 1024,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        top_k: int = 50,
+        repetition_penalty: float = 1.0,
+        **kwargs: Any,
+    ) -> AsyncIterator[CompletionChunk]:
+        tool_specs = normalize_tools(tools)
+        tools_by_name = build_tool_map(tools)
+        body: dict[str, Any] = {
+            "messages": messages,
+            "tools": tool_specs,
+            "max_new_tokens": max_new_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": top_k,
+            "repetition_penalty": repetition_penalty,
+            **kwargs,
+        }
+        complete_url = f"{self.url}/complete"
+        async with httpx.AsyncClient(timeout=None) as client:
+            async with client.stream("POST", complete_url, json=body) as response:
+                response.raise_for_status()
+                async for line in response.aiter_lines():
+                    answered = json.loads(line)
+                    chunk = CompletionChunk.from_wire(answered, tools_by_name)
+                    yield chunk
+
+    async def last_hidden_states(
+        self, parts: list[str], continuations_of_each_part: list[list[str]]
+    ) -> AsyncIterator[list[list[float]]]:
+        body = {"parts": parts, "continuations_of_each_part": continuations_of_each_part}
+        last_hidden_states_url = f"{self.url}/last_hidden_states"
+        async with httpx.AsyncClient(timeout=None) as client:
+            async with client.stream("POST", last_hidden_states_url, json=body) as response:
+                response.raise_for_status()
+                async for line in response.aiter_lines():
+                    answered = json.loads(line)
+                    yield answered["last_hidden_states"]

{cortexgrid_infer-0.3.6 → cortexgrid_infer-0.3.7}/cortexgrid_infer/serve_apps/anthropic.py

@@ -19,12 +19,13 @@ from typing import Any
 import cortexgrid
 from anthropic import AsyncAnthropic
 from cortexgrid import serve
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException, status
 from fastapi.responses import StreamingResponse
 
 from cortexgrid_infer.core import Message, ToolSpec
-from cortexgrid_infer.protocols.completion import ServedCompletingModel, encode_tool_call
+from cortexgrid_infer.protocols.completion import ServedCompletingModel
 from cortexgrid_infer.serve_apps.base import HostedModel
+from cortexgrid_infer.serve_apps.streamed_completion import StreamedCompletionWriter
 
 _app = FastAPI()
 
@@ -162,19 +163,30 @@ class AnthropicText2Text(HostedModel):
             create_kwargs["tools"] = tools
 
         async def stream() -> AsyncIterator[bytes]:
+            writer = StreamedCompletionWriter()
             async with self._client.messages.stream(**create_kwargs) as events:
                 async for event in events:
                     if (
                         event.type == "content_block_delta"
                         and event.delta.type == "text_delta"
                     ):
-                        yield event.delta.text.encode("utf-8")
+                        line = writer.read_content(event.delta.text)
+                        yield line
 
                     elif event.type == "content_block_stop":
                         block = events.current_message_snapshot.content[event.index]
                         if block.type == "tool_use":
-                            yield encode_tool_call(
-                                block.name, dict(block.input)
-                            ).encode("utf-8")
-
-        return StreamingResponse(stream(), media_type="text/plain")
+                            arguments = dict(block.input)
+                            line = writer.read_tool_call(block.name, arguments, block.id)
+                            yield line
+            for line in writer.finish():
+                yield line
+
+        return StreamingResponse(stream(), media_type="application/x-ndjson")
+
+    @_app.post("/last_hidden_states")
+    async def last_hidden_states(self, body: dict[str, Any]) -> None:
+        raise HTTPException(
+            status_code=status.HTTP_501_NOT_IMPLEMENTED,
+            detail="Anthropic does not expose its models' hidden states",
+        )

{cortexgrid_infer-0.3.6 → cortexgrid_infer-0.3.7}/cortexgrid_infer/serve_apps/gemini/text2text.py

@@ -12,17 +12,14 @@ from collections.abc import AsyncIterator
 from typing import Any
 
 from cortexgrid import serve
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException, status
 from fastapi.responses import StreamingResponse
 
 from cortexgrid_infer.core import Message, ToolSpec
-from cortexgrid_infer.protocols.completion import (
-    ServedCompletingModel,
-    encode_thinking,
-    encode_tool_call,
-)
+from cortexgrid_infer.protocols.completion import ServedCompletingModel
 from cortexgrid_infer.serve_apps.base import ENABLE_THINKING_PARAM
 from cortexgrid_infer.serve_apps.gemini.base import DEFAULT_API_KEY_SECRET, GeminiModel
+from cortexgrid_infer.serve_apps.streamed_completion import StreamedCompletionWriter
 
 _app = FastAPI()
 
@@ -150,6 +147,7 @@ class GeminiText2Text(GeminiModel):
             generate_config["tools"] = tools
 
         async def stream() -> AsyncIterator[bytes]:
+            writer = StreamedCompletionWriter()
             chunks = await self._client.aio.models.generate_content_stream(
                 model=self._model, contents=contents, config=generate_config
             )
@@ -158,14 +156,27 @@ class GeminiText2Text(GeminiModel):
                     continue
                 for part in chunk.candidates[0].content.parts or []:
                     if part.function_call:
-                        yield encode_tool_call(
-                            part.function_call.name, dict(part.function_call.args or {})
-                        ).encode("utf-8")
+                        arguments = dict(part.function_call.args or {})
+                        line = writer.read_tool_call(
+                            part.function_call.name, arguments, part.function_call.id
+                        )
+                        yield line
 
                     elif part.text and part.thought:
-                        yield encode_thinking(part.text).encode("utf-8")
+                        line = writer.read_thinking(part.text)
+                        yield line
 
                     elif part.text:
-                        yield part.text.encode("utf-8")
-
-        return StreamingResponse(stream(), media_type="text/plain")
+                        line = writer.read_content(part.text)
+                        yield line
+            for line in writer.finish():
+                yield line
+
+        return StreamingResponse(stream(), media_type="application/x-ndjson")
+
+    @_app.post("/last_hidden_states")
+    async def last_hidden_states(self, body: dict[str, Any]) -> None:
+        raise HTTPException(
+            status_code=status.HTTP_501_NOT_IMPLEMENTED,
+            detail="Gemini does not expose its models' hidden states",
+        )

cortexgrid_infer-0.3.7/cortexgrid_infer/serve_apps/streamed_completion.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import itertools
+import json
+import re
+from typing import Any, Sequence
+
+_tool_call_id_counter = itertools.count()
+
+TOOL_CALL_OPENERS = ["<tool_call>", "<|tool_call|>", "```tool_call"]
+
+THINKING_OPENER = "<think>"
+THINKING_CLOSER = "</think>"
+
+TOOL_CALL_PATTERNS = [
+    re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL),
+    re.compile(r"<\|tool_call\|>\s*(\{.*?\})\s*<\|/tool_call\|>", re.DOTALL),
+    re.compile(r"```tool_call\s*(\{.*?\})\s*```", re.DOTALL),
+]
+
+
+def new_tool_call_id() -> str:
+    number = next(_tool_call_id_counter)
+    return f"call_{number}"
+
+
+def completion_chunk(
+    content: str = "",
+    thinking: str = "",
+    tool_calls: Sequence[dict[str, Any]] = (),
+    finish_reason: str | None = None,
+) -> bytes:
+    chunk = {
+        "content": content,
+        "thinking": thinking,
+        "tool_calls": list(tool_calls),
+        "finish_reason": finish_reason,
+    }
+    line = json.dumps(chunk)
+    return f"{line}\n".encode("utf-8")
+
+
+def _find_opener(text: str, openers: Sequence[str]) -> int | None:
+    """Find the earliest position of any of `openers` in text, or None."""
+    earliest = None
+    for opener in openers:
+        idx = text.find(opener)
+        if idx != -1 and (earliest is None or idx < earliest):
+            earliest = idx
+    return earliest
+
+
+def _split_at_potential_prefix(text: str, openers: Sequence[str]) -> tuple[str, str]:
+    """Split text into (safe_to_stream, potential_opener_prefix).
+
+    The second part is a suffix that could be the beginning of one of
+    `openers`, so it must be held back until more text arrives.
+    """
+    for opener in openers:
+        for length in range(1, len(opener)):
+            if text.endswith(opener[:length]):
+                return text[:-length], text[-length:]
+    return text, ""
+
+
+def tool_calls_in(text: str) -> tuple[str, list[dict[str, Any]]]:
+    tool_calls = []
+    clean_text = text
+    for pattern in TOOL_CALL_PATTERNS:
+        matches = pattern.findall(text)
+        for match in matches:
+            try:
+                data = json.loads(match)
+                name = data["name"]
+            except (json.JSONDecodeError, KeyError):
+                continue
+            arguments = data.get("arguments", {})
+            call_id = new_tool_call_id()
+            tool_calls.append({"id": call_id, "name": name, "arguments": arguments})
+        clean_text = pattern.sub("", clean_text)
+    stripped = clean_text.strip()
+    return stripped, tool_calls
+
+
+class StreamedCompletionWriter:
+    def __init__(self) -> None:
+        self._pending = ""
+        self._in_thinking = False
+        self._in_tool_call = False
+        self._tool_call_text = ""
+        self._made_a_tool_call = False
+
+    def read_content(self, text: str) -> bytes:
+        return completion_chunk(content=text)
+
+    def read_thinking(self, text: str) -> bytes:
+        return completion_chunk(thinking=text)
+
+    def read_tool_call(
+        self, name: str, arguments: dict[str, Any], call_id: str | None = None
+    ) -> bytes:
+        self._made_a_tool_call = True
+        tool_call_id = call_id or new_tool_call_id()
+        tool_call = {"id": tool_call_id, "name": name, "arguments": arguments}
+        return completion_chunk(tool_calls=[tool_call])
+
+    def read(self, text: str) -> list[bytes]:
+        if not text:
+            return []
+        if self._in_tool_call:
+            self._tool_call_text += text
+            return []
+        self._pending += text
+        lines: list[bytes] = []
+        openers = [THINKING_OPENER, *TOOL_CALL_OPENERS]
+        while self._pending and not self._in_tool_call:
+            if self._in_thinking:
+                closer_pos = self._pending.find(THINKING_CLOSER)
+                if closer_pos == -1:
+                    thought, self._pending = _split_at_potential_prefix(
+                        self._pending, [THINKING_CLOSER]
+                    )
+                    if thought:
+                        lines.append(completion_chunk(thinking=thought))
+                    break
+                if closer_pos:
+                    thought = self._pending[:closer_pos]
+                    lines.append(completion_chunk(thinking=thought))
+                self._pending = self._pending[closer_pos + len(THINKING_CLOSER) :]
+                self._in_thinking = False
+                continue
+
+            opener_pos = _find_opener(self._pending, openers)
+            if opener_pos is None:
+                safe, self._pending = _split_at_potential_prefix(self._pending, openers)
+                if safe:
+                    lines.append(completion_chunk(content=safe))
+                break
+
+            before = self._pending[:opener_pos]
+            if before:
+                lines.append(completion_chunk(content=before))
+            if self._pending.startswith(THINKING_OPENER, opener_pos):
+                self._pending = self._pending[opener_pos + len(THINKING_OPENER) :]
+                self._in_thinking = True
+            else:
+                self._tool_call_text = self._pending[opener_pos:]
+                self._in_tool_call = True
+                self._pending = ""
+        return lines
+
+    def finish(self) -> list[bytes]:
+        lines: list[bytes] = []
+        if self._pending and self._in_thinking:
+            lines.append(completion_chunk(thinking=self._pending))
+        elif self._pending:
+            lines.append(completion_chunk(content=self._pending))
+        clean_text, tool_calls = tool_calls_in(self._tool_call_text)
+        if clean_text:
+            lines.append(completion_chunk(content=clean_text))
+        if tool_calls:
+            lines.append(completion_chunk(tool_calls=tool_calls))
+        elif not self._made_a_tool_call:
+            lines.append(completion_chunk(finish_reason="stop"))
+        return lines

{cortexgrid_infer-0.3.6 → cortexgrid_infer-0.3.7}/cortexgrid_infer/serve_apps/text2text.py

@@ -10,8 +10,10 @@ its tokens are forwarded untouched.
 from __future__ import annotations
 
 import asyncio
+import copy
+import json
 import logging
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Callable
 from threading import Thread
 from typing import Any
 
@@ -22,7 +24,9 @@ from fastapi.responses import StreamingResponse
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
+    Cache,
     CompileConfig,
+    PreTrainedModel,
     TextIteratorStreamer,
 )
 import torch
@@ -35,6 +39,7 @@ from cortexgrid_infer.serve_apps.base import (
     ENABLE_THINKING_PARAM,
     LocalModel,
 )
+from cortexgrid_infer.serve_apps.streamed_completion import StreamedCompletionWriter
 
 
 log = logging.getLogger(__name__)
@@ -51,6 +56,9 @@ _app = FastAPI()
 MAX_TOTAL_TOKENS_PARAM = "max_total_tokens"
 DEFAULT_MAX_TOTAL_TOKENS = 4096
 
+DTYPE_PARAM = "dtype"
+DEFAULT_DTYPE = "float16"
+
 _RESERVED_BODY_KEYS = {
     "messages",
     "tools",
@@ -59,9 +67,55 @@ _RESERVED_BODY_KEYS = {
     "top_p",
     "top_k",
     "repetition_penalty",
+    "choices",
 }
 
 
+def allowed_next_tokens(
+    choices: list[list[int]], generated: list[int], end_of_text: int
+) -> list[int]:
+    matching = [choice for choice in choices if choice[: len(generated)] == generated]
+    continuing = {choice[len(generated)] for choice in matching if len(choice) > len(generated)}
+    finished = any(choice == generated for choice in matching)
+    allowed = sorted(continuing | {end_of_text}) if finished else sorted(continuing)
+    return allowed or [end_of_text]
+
+
+def continuation_loglikelihood(
+    model: PreTrainedModel, prompt_ids: list[int], continuation_ids: list[int]
+) -> float:
+    read = torch.tensor([prompt_ids + continuation_ids], device=model.device)
+    with torch.inference_mode():
+        outputs = model(input_ids=read, logits_to_keep=len(continuation_ids) + 1)
+    predicting_the_continuation = outputs.logits[0, :-1].float()
+    log_probabilities = torch.log_softmax(predicting_the_continuation, dim=-1)
+    positions = torch.arange(len(continuation_ids), device=log_probabilities.device)
+    continuation = torch.tensor(continuation_ids, device=log_probabilities.device)
+    return float(log_probabilities[positions, continuation].sum())
+
+
+def read_further(model: PreTrainedModel, read: Cache | None, ids: list[int]) -> Cache:
+    unread = torch.tensor([ids], device=model.device)
+    with torch.inference_mode():
+        outputs = model(input_ids=unread, past_key_values=read, use_cache=True, logits_to_keep=1)
+    return outputs.past_key_values
+
+
+def last_hidden_state_after(model: PreTrainedModel, read: Cache, ids: list[int]) -> torch.Tensor:
+    read_for_the_continuation = copy.deepcopy(read)
+    continuation = torch.tensor([ids], device=model.device)
+    with torch.inference_mode():
+        outputs = model(
+            input_ids=continuation,
+            past_key_values=read_for_the_continuation,
+            use_cache=True,
+            output_hidden_states=True,
+            logits_to_keep=1,
+        )
+    last_layer = outputs.hidden_states[-1]
+    return last_layer[0, -1].float()
+
+
 @serve.ingress(_app)
 class Text2Text(LocalModel):
     @classmethod
@@ -76,15 +130,15 @@ class Text2Text(LocalModel):
         path = cortexgrid.load_model(
             deployment.family, deployment.suffix, deployment.run_name
         )
+        settings = cortexgrid.model_config(deployment)
         self._device = detect_device()
         self._tokenizer = AutoTokenizer.from_pretrained(path)
         if self._tokenizer.pad_token is None:
             self._tokenizer.pad_token = self._tokenizer.eos_token
         self._model = AutoModelForCausalLM.from_pretrained(
-            str(path), torch_dtype=torch.float16
+            str(path), torch_dtype=getattr(torch, settings.get(DTYPE_PARAM, DEFAULT_DTYPE))
         )
         self._model.to(self._device)  # type: ignore
-        settings = cortexgrid.model_config(deployment)
         context = getattr(
             self._model.config, "max_position_embeddings", DEFAULT_MAX_TOTAL_TOKENS
         )
@@ -127,6 +181,82 @@ class Text2Text(LocalModel):
             inputs[key] = value[..., -limit:]
         return inputs
 
+    def _prompt_ids(self, messages: list[dict[str, Any]]) -> list[int]:
+        prompt = self._tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=self._enable_thinking,
+        )
+        return self._tokenizer(prompt)["input_ids"]
+
+    def _continuation_ids(self, text: str) -> list[int]:
+        return self._tokenizer(text, add_special_tokens=False)["input_ids"]
+
+    def _only_the_choices(
+        self, choices: list[str], prompt_length: int
+    ) -> Callable[[int, torch.Tensor], list[int]]:
+        choice_ids = [self._continuation_ids(choice) for choice in choices]
+        end_of_text = self._tokenizer.eos_token_id
+
+        def allowed(_batch: int, sequence: torch.Tensor) -> list[int]:
+            generated = sequence[prompt_length:].tolist()
+            return allowed_next_tokens(choice_ids, generated, end_of_text)
+
+        return allowed
+
+    def _loglikelihoods(
+        self, messages: list[dict[str, Any]], continuations: list[str]
+    ) -> list[float]:
+        prompt_ids = self._prompt_ids(messages)
+        continuation_ids = [self._continuation_ids(continuation) for continuation in continuations]
+        return [
+            continuation_loglikelihood(self._model, prompt_ids, ids) for ids in continuation_ids
+        ]
+
+    @_app.post("/loglikelihoods")
+    async def loglikelihoods(self, body: dict[str, Any]) -> dict[str, list[float]]:
+        loop = asyncio.get_running_loop()
+        loglikelihoods = await loop.run_in_executor(
+            None, self._loglikelihoods, body["messages"], body["continuations"]
+        )
+        return {"loglikelihoods": loglikelihoods}
+
+    def _read_part(
+        self, read: Cache | None, part: str, continuations: list[str]
+    ) -> tuple[Cache, list[list[float]]]:
+        tokenized = self._tokenizer(part, add_special_tokens=read is None)
+        part_ids = tokenized["input_ids"]
+        read = read_further(self._model, read, part_ids)
+        continuation_ids = [self._continuation_ids(continuation) for continuation in continuations]
+        last_hidden_states = [
+            last_hidden_state_after(self._model, read, ids) for ids in continuation_ids
+        ]
+        answered = [hidden_state.tolist() for hidden_state in last_hidden_states]
+        return read, answered
+
+    async def _last_hidden_states_of_each_part(
+        self, parts: list[str], continuations_of_each_part: list[list[str]]
+    ) -> AsyncIterator[bytes]:
+        loop = asyncio.get_running_loop()
+        read: Cache | None = None
+        for part, continuations in zip(parts, continuations_of_each_part):
+            read, last_hidden_states = await loop.run_in_executor(
+                None, self._read_part, read, part, continuations
+            )
+            answered = json.dumps({"last_hidden_states": last_hidden_states})
+            line = f"{answered}\n".encode("utf-8")
+            yield line
+
+    @_app.post("/last_hidden_states")
+    async def last_hidden_states(self, body: dict[str, Any]) -> StreamingResponse:
+        last_hidden_states_of_each_part = self._last_hidden_states_of_each_part(
+            body["parts"], body["continuations_of_each_part"]
+        )
+        return StreamingResponse(
+            last_hidden_states_of_each_part, media_type="application/x-ndjson"
+        )
+
     @_app.post("/complete")
     async def complete(self, body: dict[str, Any]) -> StreamingResponse:
         messages = body["messages"]
@@ -138,6 +268,7 @@ class Text2Text(LocalModel):
         top_p = body.get("top_p", 0.9)
         top_k = body.get("top_k", 50)
         repetition_penalty = body.get("repetition_penalty", 1.0)
+        choices = body.get("choices")
         extra = {k: v for k, v in body.items() if k not in _RESERVED_BODY_KEYS}
 
         prompt = self._tokenizer.apply_chat_template(
@@ -151,6 +282,9 @@ class Text2Text(LocalModel):
             self._tokenizer(prompt, return_tensors="pt"),
             self._max_total_tokens - max_new_tokens,
         ).to(self._model.device)
+        if choices:
+            prompt_length = inputs["input_ids"].shape[-1]
+            extra["prefix_allowed_tokens_fn"] = self._only_the_choices(choices, prompt_length)
 
         queue: asyncio.Queue[str | None] = asyncio.Queue()
         error: list[BaseException] = []
@@ -192,13 +326,16 @@ class Text2Text(LocalModel):
         loop.run_in_executor(None, generate)
 
         async def stream() -> AsyncIterator[bytes]:
+            writer = StreamedCompletionWriter()
             while True:
                 text = await queue.get()
                 if text is None:
                     break
-                if text:
-                    yield text.encode("utf-8")
+                for line in writer.read(text):
+                    yield line
             if error:
                 raise error[0]
+            for line in writer.finish():
+                yield line
 
-        return StreamingResponse(stream(), media_type="text/plain")
+        return StreamingResponse(stream(), media_type="application/x-ndjson")

{cortexgrid_infer-0.3.6 → cortexgrid_infer-0.3.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cortexgrid-infer"
-version = "0.3.6"
+version = "0.3.7"
 description = "Serve apps per task (text-to-text, text-to-image, image-to-mesh) and importers per source (HuggingFace) for running models on cortexgrid"
 readme = "README.md"
 license = "Apache-2.0"

cortexgrid_infer-0.3.6/cortexgrid_infer/protocols/completion.py

@@ -1,222 +0,0 @@
-"""The wire protocol between a cortexgrid-served completion app and its client.
-
-A completion serve app streams plain text from `POST /complete`, with tool calls
-inlined in the text as `<tool_call>{"name": ..., "arguments": {...}}</tool_call>`.
-That is the format an instruct-tuned causal LM emits on its own, so a serve app
-wrapping one forwards its tokens untouched; a serve app wrapping a provider with
-structured tool calls encodes them into the same form with `encode_tool_call`.
-
-Either way the client is the same `ServedCompletingModel`, which streams text
-back as `CompletionChunk`s and reassembles the tool calls.
-"""
-
-from __future__ import annotations
-
-from collections.abc import AsyncIterator
-from dataclasses import dataclass
-from functools import partial
-import itertools
-import json
-import re
-from typing import Any, Callable, Sequence
-
-import httpx
-
-from cortexgrid_infer.core import (
-    CompletingModel,
-    CompletionChunk,
-    Message,
-    Tool,
-    ToolCall,
-)
-from cortexgrid_infer.utils import build_tool_map, normalize_tools
-
-_tool_call_id_counter = itertools.count()
-
-TOOL_CALL_OPENERS = ["<tool_call>", "<|tool_call|>", "```tool_call"]
-
-THINKING_OPENER = "<think>"
-THINKING_CLOSER = "</think>"
-
-TOOL_CALL_PATTERNS = [
-    re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL),
-    re.compile(r"<\|tool_call\|>\s*(\{.*?\})\s*<\|/tool_call\|>", re.DOTALL),
-    re.compile(r"```tool_call\s*(\{.*?\})\s*```", re.DOTALL),
-]
-
-
-def encode_tool_call(name: str, arguments: dict[str, Any]) -> str:
-    """Render a tool call in the form the client parses back out of the stream.
-
-    For serve apps whose upstream reports tool calls as structured data rather
-    than as text the model generated - they are re-encoded here so one client
-    handles every completion app."""
-    return f'<tool_call>{json.dumps({"name": name, "arguments": arguments})}</tool_call>'
-
-
-def encode_thinking(text: str) -> str:
-    return f"{THINKING_OPENER}{text}{THINKING_CLOSER}"
-
-
-def _find_opener(text: str, openers: Sequence[str]) -> int | None:
-    """Find the earliest position of any of `openers` in text, or None."""
-    earliest = None
-    for opener in openers:
-        idx = text.find(opener)
-        if idx != -1 and (earliest is None or idx < earliest):
-            earliest = idx
-    return earliest
-
-
-def _split_at_potential_prefix(text: str, openers: Sequence[str]) -> tuple[str, str]:
-    """Split text into (safe_to_stream, potential_opener_prefix).
-
-    The second part is a suffix that could be the beginning of one of
-    `openers`, so it must be held back until more text arrives.
-    """
-    for opener in openers:
-        for length in range(1, len(opener)):
-            if text.endswith(opener[:length]):
-                return text[:-length], text[-length:]
-    return text, ""
-
-
-def parse_tool_calls(
-    text: str, tool_map: dict[str, Callable[..., Any]]
-) -> tuple[str, list[ToolCall]]:
-    """Parse tool calls from generated text and bind them to their functions."""
-    tool_calls: list[ToolCall] = []
-    clean_text = text
-
-    for pattern in TOOL_CALL_PATTERNS:
-        matches = pattern.findall(text)
-        for match in matches:
-            try:
-                data = json.loads(match)
-                name = data["name"]
-                arguments = data.get("arguments", {})
-                if name in tool_map:
-                    tool_calls.append(
-                        ToolCall(
-                            id=f"call_{next(_tool_call_id_counter)}",
-                            name=name,
-                            arguments=arguments,
-                            _func=partial(tool_map[name], **arguments),
-                        )
-                    )
-            except (json.JSONDecodeError, KeyError):
-                continue
-        clean_text = pattern.sub("", clean_text)
-
-    return clean_text.strip(), tool_calls
-
-
-@dataclass
-class ServedCompletingModel(CompletingModel):
-    """Client for a completion app cortexgrid has deployed at `url`.
-
-    Provider-agnostic: every completion serve app speaks the protocol above, so
-    which model is behind it only shows up in `model_id`."""
-
-    url: str
-    model_id: str
-
-    @property
-    def name(self) -> str:
-        return self.model_id
-
-    async def complete(
-        self,
-        messages: list[Message],
-        tools: Sequence[Tool] | None = None,
-        max_new_tokens: int = 16 * 1024,
-        temperature: float = 0.7,
-        top_p: float = 0.9,
-        top_k: int = 50,
-        repetition_penalty: float = 1.0,
-        **kwargs: Any,
-    ) -> AsyncIterator[CompletionChunk]:
-        tool_specs = normalize_tools(tools)
-        tool_map = build_tool_map(tools)
-
-        body: dict[str, Any] = {
-            "messages": messages,
-            "tools": tool_specs,
-            "max_new_tokens": max_new_tokens,
-            "temperature": temperature,
-            "top_p": top_p,
-            "top_k": top_k,
-            "repetition_penalty": repetition_penalty,
-            **kwargs,
-        }
-
-        pending = ""
-        in_thinking = False
-        in_tool_call = False
-        tool_call_text = ""
-        openers = [THINKING_OPENER, *TOOL_CALL_OPENERS]
-
-        async with httpx.AsyncClient(timeout=None) as client:
-            async with client.stream(
-                "POST", f"{self.url}/complete", json=body
-            ) as response:
-                response.raise_for_status()
-                async for text in response.aiter_text():
-                    if not text:
-                        continue
-
-                    if in_tool_call:
-                        tool_call_text += text
-                        continue
-
-                    pending += text
-
-                    while pending and not in_tool_call:
-                        if in_thinking:
-                            closer_pos = pending.find(THINKING_CLOSER)
-                            if closer_pos == -1:
-                                thought, pending = _split_at_potential_prefix(
-                                    pending, [THINKING_CLOSER]
-                                )
-                                if thought:
-                                    yield CompletionChunk(thinking=thought)
-                                break
-                            if closer_pos:
-                                yield CompletionChunk(thinking=pending[:closer_pos])
-                            pending = pending[closer_pos + len(THINKING_CLOSER) :]
-                            in_thinking = False
-                            continue
-
-                        opener_pos = _find_opener(pending, openers)
-                        if opener_pos is None:
-                            safe, pending = _split_at_potential_prefix(pending, openers)
-                            if safe:
-                                yield CompletionChunk(content=safe)
-                            break
-
-                        before = pending[:opener_pos]
-                        if before:
-                            yield CompletionChunk(content=before)
-                        if pending.startswith(THINKING_OPENER, opener_pos):
-                            pending = pending[opener_pos + len(THINKING_OPENER) :]
-                            in_thinking = True
-                        else:
-                            tool_call_text = pending[opener_pos:]
-                            in_tool_call = True
-                            pending = ""
-
-        if pending and in_thinking:
-            yield CompletionChunk(thinking=pending)
-        elif pending:
-            yield CompletionChunk(content=pending)
-
-        if tool_call_text:
-            clean_text, tool_calls = parse_tool_calls(tool_call_text, tool_map)
-            if clean_text:
-                yield CompletionChunk(content=clean_text)
-            if tool_calls:
-                yield CompletionChunk(tool_calls=tool_calls)
-            else:
-                yield CompletionChunk(finish_reason="stop")
-        else:
-            yield CompletionChunk(finish_reason="stop")