Refactor tasks / commands / api
This commit is contained in:
Seth Howes
+20
-27
@@ -2,38 +2,30 @@ import asyncio
|
||||
import time
|
||||
from asyncio.queues import Queue
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import List, Optional, Sequence, final
|
||||
from typing import Sequence, final
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
from shared.db.sqlite.connector import AsyncSQLiteEventStorage
|
||||
from shared.types.api import (
|
||||
ChatCompletionMessage,
|
||||
ChatCompletionResponse,
|
||||
StreamingChoiceResponse,
|
||||
)
|
||||
from shared.types.events import ChunkGenerated, Event
|
||||
from shared.types.events.chunks import TokenChunk
|
||||
from shared.types.events.commands import (
|
||||
ChatCompletionCommand,
|
||||
Command,
|
||||
CommandId,
|
||||
CommandTypes,
|
||||
)
|
||||
from shared.types.events.components import EventFromEventLog
|
||||
from shared.types.request import APIRequest, RequestId
|
||||
from shared.types.tasks import ChatCompletionTaskParams
|
||||
|
||||
|
||||
class Message(BaseModel):
|
||||
role: str
|
||||
content: str
|
||||
|
||||
class StreamingChoiceResponse(BaseModel):
|
||||
index: int
|
||||
delta: Message
|
||||
finish_reason: Optional[str] = None
|
||||
|
||||
|
||||
class ChatCompletionResponse(BaseModel):
|
||||
id: str
|
||||
object: str = "chat.completion"
|
||||
created: int
|
||||
model: str
|
||||
choices: List[StreamingChoiceResponse]
|
||||
|
||||
def chunk_to_response(chunk: TokenChunk) -> ChatCompletionResponse:
|
||||
return ChatCompletionResponse(
|
||||
id='abc',
|
||||
@@ -42,7 +34,7 @@ def chunk_to_response(chunk: TokenChunk) -> ChatCompletionResponse:
|
||||
choices=[
|
||||
StreamingChoiceResponse(
|
||||
index=0,
|
||||
delta=Message(
|
||||
delta=ChatCompletionMessage(
|
||||
role='assistant',
|
||||
content=chunk.text
|
||||
),
|
||||
@@ -54,7 +46,7 @@ def chunk_to_response(chunk: TokenChunk) -> ChatCompletionResponse:
|
||||
|
||||
@final
|
||||
class API:
|
||||
def __init__(self, command_queue: Queue[APIRequest], global_events: AsyncSQLiteEventStorage) -> None:
|
||||
def __init__(self, command_queue: Queue[Command], global_events: AsyncSQLiteEventStorage) -> None:
|
||||
self._app = FastAPI()
|
||||
self._setup_routes()
|
||||
|
||||
@@ -106,10 +98,11 @@ class API:
|
||||
|
||||
# At the moment, we just create the task in the API.
|
||||
# In the future, a `Request` will be created here and they will be bundled into `Task` objects by the master.
|
||||
request_id=RequestId()
|
||||
command_id=CommandId()
|
||||
|
||||
request = APIRequest(
|
||||
request_id=request_id,
|
||||
request = ChatCompletionCommand(
|
||||
command_id=command_id,
|
||||
command_type=CommandTypes.CHAT_COMPLETION,
|
||||
request_params=payload,
|
||||
)
|
||||
await self.command_queue.put(request)
|
||||
@@ -124,7 +117,7 @@ class API:
|
||||
|
||||
for wrapped_event in events:
|
||||
event = wrapped_event.event
|
||||
if isinstance(event, ChunkGenerated) and event.request_id == request_id:
|
||||
if isinstance(event, ChunkGenerated) and event.command_id == command_id:
|
||||
assert isinstance(event.chunk, TokenChunk)
|
||||
chunk_response: ChatCompletionResponse = chunk_to_response(event.chunk)
|
||||
print(chunk_response)
|
||||
@@ -146,7 +139,7 @@ class API:
|
||||
|
||||
|
||||
def start_fastapi_server(
|
||||
command_queue: Queue[APIRequest],
|
||||
command_queue: Queue[Command],
|
||||
global_events: AsyncSQLiteEventStorage,
|
||||
host: str = "0.0.0.0",
|
||||
port: int = 8000,
|
||||
|
||||
+8
-8
@@ -10,11 +10,11 @@ from shared.db.sqlite.event_log_manager import EventLogManager
|
||||
from shared.types.common import NodeId
|
||||
from shared.types.events import ChunkGenerated
|
||||
from shared.types.events.chunks import TokenChunk
|
||||
from shared.types.request import APIRequest, RequestId
|
||||
from shared.types.events.commands import Command, CommandId
|
||||
|
||||
|
||||
## TODO: Hook this up properly
|
||||
async def fake_tokens_task(events_log: AsyncSQLiteEventStorage, request_id: RequestId):
|
||||
async def fake_tokens_task(events_log: AsyncSQLiteEventStorage, command_id: CommandId):
|
||||
model_id = "testmodelabc"
|
||||
|
||||
for i in range(10):
|
||||
@@ -22,9 +22,9 @@ async def fake_tokens_task(events_log: AsyncSQLiteEventStorage, request_id: Requ
|
||||
|
||||
# Create the event with proper types and consistent IDs
|
||||
chunk_event = ChunkGenerated(
|
||||
request_id=request_id,
|
||||
command_id=command_id,
|
||||
chunk=TokenChunk(
|
||||
request_id=request_id, # Use the same task_id
|
||||
command_id=command_id, # Use the same task_id
|
||||
idx=i,
|
||||
model=model_id, # Use the same model_id
|
||||
text=f'text{i}',
|
||||
@@ -42,9 +42,9 @@ async def fake_tokens_task(events_log: AsyncSQLiteEventStorage, request_id: Requ
|
||||
|
||||
# Create the event with proper types and consistent IDs
|
||||
chunk_event = ChunkGenerated(
|
||||
request_id=request_id,
|
||||
command_id=command_id,
|
||||
chunk=TokenChunk(
|
||||
request_id=request_id, # Use the same task_id
|
||||
command_id=command_id, # Use the same task_id
|
||||
idx=11,
|
||||
model=model_id, # Use the same model_id
|
||||
text=f'text{11}',
|
||||
@@ -68,7 +68,7 @@ async def main():
|
||||
await event_log_manager.initialize()
|
||||
global_events: AsyncSQLiteEventStorage = event_log_manager.global_events
|
||||
|
||||
command_queue: Queue[APIRequest] = asyncio.Queue()
|
||||
command_queue: Queue[Command] = asyncio.Queue()
|
||||
|
||||
api_thread = threading.Thread(
|
||||
target=start_fastapi_server,
|
||||
@@ -88,7 +88,7 @@ async def main():
|
||||
|
||||
print(command)
|
||||
|
||||
await fake_tokens_task(global_events, request_id=command.request_id)
|
||||
await fake_tokens_task(global_events, command_id=command.command_id)
|
||||
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
|
||||
@@ -0,0 +1,503 @@
|
||||
# This file is automatically @generated by Cargo.
|
||||
# It is not intended for manual editing.
|
||||
version = 4
|
||||
|
||||
[[package]]
|
||||
name = "addr2line"
|
||||
version = "0.22.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6e4503c46a5c0c7844e948c9a4d6acd9f50cccb4de1c48eb9e291ea17470c678"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"gimli 0.29.0",
|
||||
"rustc-std-workspace-alloc",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "adler"
|
||||
version = "1.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "alloc"
|
||||
version = "0.0.0"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"core",
|
||||
"rand",
|
||||
"rand_xorshift",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "allocator-api2"
|
||||
version = "0.2.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5c6cb57a04249c6480766f7f7cef5467412af1490f8d1e243141daddada3264f"
|
||||
|
||||
[[package]]
|
||||
name = "cc"
|
||||
version = "1.1.22"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9540e661f81799159abee814118cc139a2004b3a3aa3ea37724a1b66530b90e0"
|
||||
dependencies = [
|
||||
"shlex",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cfg-if"
|
||||
version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "compiler_builtins"
|
||||
version = "0.1.138"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "53f0ea7fff95b51f84371588f06062557e96bbe363d2b36218ddb806f3ca8611"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "core"
|
||||
version = "0.0.0"
|
||||
dependencies = [
|
||||
"rand",
|
||||
"rand_xorshift",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "dlmalloc"
|
||||
version = "0.2.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d9b5e0d321d61de16390ed273b647ce51605b575916d3c25e6ddf27a1e140035"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"compiler_builtins",
|
||||
"libc",
|
||||
"rustc-std-workspace-core",
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fortanix-sgx-abi"
|
||||
version = "0.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "57cafc2274c10fab234f176b25903ce17e690fca7597090d50880e047a0389c5"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "getopts"
|
||||
version = "0.2.21"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "14dbbfd5c71d70241ecf9e6f13737f7b5ce823821063188d7e46c41d371eebd5"
|
||||
dependencies = [
|
||||
"rustc-std-workspace-core",
|
||||
"rustc-std-workspace-std",
|
||||
"unicode-width",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "gimli"
|
||||
version = "0.29.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "40ecd4077b5ae9fd2e9e169b102c6c330d0605168eb0e8bf79952b256dbefffd"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-alloc",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "gimli"
|
||||
version = "0.31.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-alloc",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hashbrown"
|
||||
version = "0.15.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1e087f84d4f86bf4b218b927129862374b72199ae7d8657835f1e89000eea4fb"
|
||||
dependencies = [
|
||||
"allocator-api2",
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-alloc",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hermit-abi"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-alloc",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.162"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "18d287de67fe55fd7e1581fe933d965a5a9477b38e949cfa9f8574ef01506398"
|
||||
dependencies = [
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "memchr"
|
||||
version = "2.7.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "miniz_oxide"
|
||||
version = "0.7.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b8a240ddb74feaf34a79a7add65a741f3167852fba007066dcac1ca548d89c08"
|
||||
dependencies = [
|
||||
"adler",
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-alloc",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "object"
|
||||
version = "0.36.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "aedf0a2d09c573ed1d8d85b30c119153926a2b36dce0ab28322c09a117a4683e"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"memchr",
|
||||
"rustc-std-workspace-alloc",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "panic_abort"
|
||||
version = "0.0.0"
|
||||
dependencies = [
|
||||
"alloc",
|
||||
"cfg-if",
|
||||
"compiler_builtins",
|
||||
"core",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "panic_unwind"
|
||||
version = "0.0.0"
|
||||
dependencies = [
|
||||
"alloc",
|
||||
"cfg-if",
|
||||
"compiler_builtins",
|
||||
"core",
|
||||
"libc",
|
||||
"unwind",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "proc_macro"
|
||||
version = "0.0.0"
|
||||
dependencies = [
|
||||
"core",
|
||||
"std",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "profiler_builtins"
|
||||
version = "0.0.0"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"compiler_builtins",
|
||||
"core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "r-efi"
|
||||
version = "4.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e9e935efc5854715dfc0a4c9ef18dc69dee0ec3bf9cc3ab740db831c0fdd86a3"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "r-efi-alloc"
|
||||
version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "31d6f09fe2b6ad044bc3d2c34ce4979796581afd2f1ebc185837e02421e02fd7"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"r-efi",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.8.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
|
||||
dependencies = [
|
||||
"rand_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand_core"
|
||||
version = "0.6.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
|
||||
|
||||
[[package]]
|
||||
name = "rand_xorshift"
|
||||
version = "0.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d25bf25ec5ae4a3f1b92f929810509a2f53d7dca2f50b794ff57e3face536c8f"
|
||||
dependencies = [
|
||||
"rand_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustc-demangle"
|
||||
version = "0.1.24"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustc-std-workspace-alloc"
|
||||
version = "1.99.0"
|
||||
dependencies = [
|
||||
"alloc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustc-std-workspace-core"
|
||||
version = "1.99.0"
|
||||
dependencies = [
|
||||
"core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustc-std-workspace-std"
|
||||
version = "1.99.0"
|
||||
dependencies = [
|
||||
"std",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "shlex"
|
||||
version = "1.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
|
||||
|
||||
[[package]]
|
||||
name = "std"
|
||||
version = "0.0.0"
|
||||
dependencies = [
|
||||
"addr2line",
|
||||
"alloc",
|
||||
"cfg-if",
|
||||
"compiler_builtins",
|
||||
"core",
|
||||
"dlmalloc",
|
||||
"fortanix-sgx-abi",
|
||||
"hashbrown",
|
||||
"hermit-abi",
|
||||
"libc",
|
||||
"miniz_oxide",
|
||||
"object",
|
||||
"panic_abort",
|
||||
"panic_unwind",
|
||||
"r-efi",
|
||||
"r-efi-alloc",
|
||||
"rand",
|
||||
"rand_xorshift",
|
||||
"rustc-demangle",
|
||||
"std_detect",
|
||||
"unwind",
|
||||
"wasi",
|
||||
"windows-targets 0.0.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "std_detect"
|
||||
version = "0.1.5"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"compiler_builtins",
|
||||
"libc",
|
||||
"rustc-std-workspace-alloc",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sysroot"
|
||||
version = "0.0.0"
|
||||
dependencies = [
|
||||
"proc_macro",
|
||||
"profiler_builtins",
|
||||
"std",
|
||||
"test",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "test"
|
||||
version = "0.0.0"
|
||||
dependencies = [
|
||||
"core",
|
||||
"getopts",
|
||||
"libc",
|
||||
"std",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unicode-width"
|
||||
version = "0.1.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-core",
|
||||
"rustc-std-workspace-std",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unwind"
|
||||
version = "0.0.0"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"compiler_builtins",
|
||||
"core",
|
||||
"libc",
|
||||
"unwinding",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unwinding"
|
||||
version = "0.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "637d511437df708cee34bdec7ba2f1548d256b7acf3ff20e0a1c559f9bf3a987"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"gimli 0.31.1",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasi"
|
||||
version = "0.11.0+wasi-snapshot-preview1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
|
||||
dependencies = [
|
||||
"compiler_builtins",
|
||||
"rustc-std-workspace-alloc",
|
||||
"rustc-std-workspace-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.59.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
|
||||
dependencies = [
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-targets"
|
||||
version = "0.0.0"
|
||||
|
||||
[[package]]
|
||||
name = "windows-targets"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973"
|
||||
dependencies = [
|
||||
"windows_aarch64_gnullvm",
|
||||
"windows_aarch64_msvc",
|
||||
"windows_i686_gnu",
|
||||
"windows_i686_gnullvm",
|
||||
"windows_i686_msvc",
|
||||
"windows_x86_64_gnu",
|
||||
"windows_x86_64_gnullvm",
|
||||
"windows_x86_64_msvc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows_aarch64_gnullvm"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3"
|
||||
|
||||
[[package]]
|
||||
name = "windows_aarch64_msvc"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_gnu"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_gnullvm"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_msvc"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_gnu"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_gnullvm"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_msvc"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
|
||||
@@ -1,28 +0,0 @@
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import BaseModel, Field, TypeAdapter
|
||||
|
||||
from shared.types.common import NewUUID
|
||||
|
||||
|
||||
class ExternalCommandId(NewUUID):
|
||||
pass
|
||||
|
||||
|
||||
class BaseExternalCommand[T: str](BaseModel):
|
||||
command_id: ExternalCommandId
|
||||
command_type: T
|
||||
|
||||
|
||||
class ChatCompletionNonStreamingCommand(
|
||||
BaseExternalCommand[Literal["chat_completion_non_streaming"]]
|
||||
):
|
||||
command_type: Literal["chat_completion_non_streaming"] = (
|
||||
"chat_completion_non_streaming"
|
||||
)
|
||||
|
||||
|
||||
ExternalCommand = Annotated[
|
||||
ChatCompletionNonStreamingCommand, Field(discriminator="command_type")
|
||||
]
|
||||
ExternalCommandParser: TypeAdapter[ExternalCommand] = TypeAdapter(ExternalCommand)
|
||||
@@ -1,120 +0,0 @@
|
||||
from asyncio import Lock, Task
|
||||
from asyncio import Queue as AsyncQueue
|
||||
from collections.abc import MutableMapping
|
||||
from logging import Logger
|
||||
from typing import Any, Hashable, Mapping, Protocol, Sequence
|
||||
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
from shared.event_loops.commands import ExternalCommand
|
||||
from shared.types.events import Event
|
||||
from shared.types.events.components import Apply, EventFromEventLog
|
||||
from shared.types.state import State
|
||||
|
||||
|
||||
class ExhaustiveMapping[K: Hashable, V](MutableMapping[K, V]):
|
||||
__slots__ = ("_store",)
|
||||
|
||||
required_keys: frozenset[K] = frozenset()
|
||||
|
||||
def __init__(self, data: Mapping[K, V]):
|
||||
missing = self.required_keys - data.keys()
|
||||
extra = data.keys() - self.required_keys
|
||||
if missing or extra:
|
||||
raise ValueError(f"missing={missing!r}, extra={extra!r}")
|
||||
self._store: dict[K, V] = dict(data)
|
||||
|
||||
def __getitem__(self, k: K) -> V:
|
||||
return self._store[k]
|
||||
|
||||
def __setitem__(self, k: K, v: V) -> None:
|
||||
self._store[k] = v
|
||||
|
||||
def __delitem__(self, k: K) -> None:
|
||||
del self._store[k]
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._store)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._store)
|
||||
|
||||
|
||||
def apply_events(
|
||||
state: State, apply_fn: Apply, events: Sequence[EventFromEventLog[Event]]
|
||||
) -> State:
|
||||
sorted_events = sorted(events, key=lambda event: event.idx_in_log)
|
||||
state = state.model_copy()
|
||||
for wrapped_event in sorted_events:
|
||||
if wrapped_event.idx_in_log <= state.last_event_applied_idx:
|
||||
continue
|
||||
state.last_event_applied_idx = wrapped_event.idx_in_log
|
||||
state = apply_fn(state, wrapped_event.event)
|
||||
return state
|
||||
|
||||
|
||||
class NodeCommandLoopProtocol(Protocol):
|
||||
_command_runner: Task[Any] | None = None
|
||||
_command_queue: AsyncQueue[ExternalCommand]
|
||||
_response_queue: AsyncQueue[Response | StreamingResponse]
|
||||
_logger: Logger
|
||||
|
||||
@property
|
||||
def is_command_runner_running(self) -> bool:
|
||||
return self._command_runner is not None and not self._command_runner.done()
|
||||
|
||||
async def start_command_runner(self) -> None: ...
|
||||
async def stop_command_runner(self) -> None: ...
|
||||
async def push_command(self, command: ExternalCommand) -> None: ...
|
||||
async def pop_response(self) -> Response | StreamingResponse: ...
|
||||
async def _handle_command(self, command: ExternalCommand) -> None: ...
|
||||
|
||||
|
||||
class NodeEventGetterProtocol(Protocol):
|
||||
_event_fetcher: Task[Any] | None = None
|
||||
_event_queue: AsyncQueue[EventFromEventLog[Event]]
|
||||
_logger: Logger
|
||||
|
||||
@property
|
||||
async def is_event_fetcher_running(self) -> bool:
|
||||
return self._event_fetcher is not None and not self._event_fetcher.done()
|
||||
|
||||
async def start_event_fetcher(self) -> None: ...
|
||||
async def stop_event_fetcher(self) -> None: ...
|
||||
|
||||
|
||||
class NodeStateStorageProtocol(Protocol):
|
||||
_state: State
|
||||
_state_lock: Lock
|
||||
_logger: Logger
|
||||
|
||||
async def _read_state(
|
||||
self,
|
||||
) -> State: ...
|
||||
|
||||
|
||||
class NodeStateManagerProtocol(
|
||||
NodeEventGetterProtocol, NodeStateStorageProtocol
|
||||
):
|
||||
_state_manager: Task[Any] | None = None
|
||||
_logger: Logger
|
||||
|
||||
@property
|
||||
async def is_state_manager_running(self) -> bool:
|
||||
is_task_running = (
|
||||
self._state_manager is not None and not self._state_manager.done()
|
||||
)
|
||||
return (
|
||||
is_task_running
|
||||
and await self.is_event_fetcher_running
|
||||
and await self.is_state_manager_running
|
||||
)
|
||||
|
||||
async def start_state_manager(self) -> None: ...
|
||||
async def stop_state_manager(self) -> None: ...
|
||||
async def _apply_queued_events(self) -> None: ...
|
||||
|
||||
|
||||
class NodeEventLoopProtocol(
|
||||
NodeCommandLoopProtocol, NodeStateManagerProtocol
|
||||
): ...
|
||||
@@ -16,7 +16,7 @@ from shared.types.events import (
|
||||
_EventType,
|
||||
)
|
||||
from shared.types.events.chunks import ChunkType, TokenChunk
|
||||
from shared.types.request import RequestId
|
||||
from shared.types.events.commands import CommandId
|
||||
|
||||
# Type ignore comment for all protected member access in this test file
|
||||
# pyright: reportPrivateUsage=false
|
||||
@@ -439,19 +439,19 @@ class TestAsyncSQLiteEventStorage:
|
||||
await storage.start()
|
||||
|
||||
# Create a ChunkGenerated event with nested TokenChunk
|
||||
request_id = RequestId(uuid=uuid4())
|
||||
command_id = CommandId(uuid=uuid4())
|
||||
token_chunk = TokenChunk(
|
||||
text="Hello, world!",
|
||||
token_id=42,
|
||||
finish_reason="stop",
|
||||
chunk_type=ChunkType.token,
|
||||
request_id=request_id,
|
||||
command_id=command_id,
|
||||
idx=0,
|
||||
model="test-model"
|
||||
)
|
||||
|
||||
chunk_generated_event = ChunkGenerated(
|
||||
request_id=request_id,
|
||||
command_id=command_id,
|
||||
chunk=token_chunk
|
||||
)
|
||||
|
||||
@@ -473,13 +473,13 @@ class TestAsyncSQLiteEventStorage:
|
||||
retrieved_event = retrieved_event_wrapper.event
|
||||
assert isinstance(retrieved_event, ChunkGenerated)
|
||||
assert retrieved_event.event_type == _EventType.ChunkGenerated
|
||||
assert retrieved_event.request_id == request_id
|
||||
assert retrieved_event.command_id == command_id
|
||||
|
||||
# Verify the nested chunk was deserialized correctly
|
||||
retrieved_chunk = retrieved_event.chunk
|
||||
assert isinstance(retrieved_chunk, TokenChunk)
|
||||
assert retrieved_chunk.chunk_type == ChunkType.token
|
||||
assert retrieved_chunk.request_id == request_id
|
||||
assert retrieved_chunk.command_id == command_id
|
||||
assert retrieved_chunk.idx == 0
|
||||
assert retrieved_chunk.model == "test-model"
|
||||
|
||||
|
||||
+65
-1
@@ -2,6 +2,8 @@ from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from shared.openai_compat import FinishReason
|
||||
|
||||
|
||||
class ChatCompletionMessage(BaseModel):
|
||||
role: Literal["system", "user", "assistant", "developer", "tool", "function"]
|
||||
@@ -12,6 +14,68 @@ class ChatCompletionMessage(BaseModel):
|
||||
function_call: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class TopLogprobItem(BaseModel):
|
||||
token: str
|
||||
logprob: float
|
||||
bytes: list[int] | None = None
|
||||
|
||||
|
||||
class LogprobsContentItem(BaseModel):
|
||||
token: str
|
||||
logprob: float
|
||||
bytes: list[int] | None = None
|
||||
top_logprobs: list[TopLogprobItem]
|
||||
|
||||
|
||||
class Logprobs(BaseModel):
|
||||
content: list[LogprobsContentItem] | None = None
|
||||
|
||||
|
||||
class StreamingChoiceResponse(BaseModel):
|
||||
index: int
|
||||
delta: ChatCompletionMessage
|
||||
logprobs: Logprobs | None = None
|
||||
finish_reason: FinishReason | None = None
|
||||
|
||||
|
||||
class ChatCompletionChoice(BaseModel):
|
||||
index: int
|
||||
message: ChatCompletionMessage
|
||||
logprobs: Logprobs | None = None
|
||||
finish_reason: FinishReason | None = None
|
||||
|
||||
|
||||
class PromptTokensDetails(BaseModel):
|
||||
cached_tokens: int = 0
|
||||
audio_tokens: int = 0
|
||||
|
||||
|
||||
class CompletionTokensDetails(BaseModel):
|
||||
reasoning_tokens: int = 0
|
||||
audio_tokens: int = 0
|
||||
accepted_prediction_tokens: int = 0
|
||||
rejected_prediction_tokens: int = 0
|
||||
|
||||
|
||||
class Usage(BaseModel):
|
||||
prompt_tokens: int
|
||||
completion_tokens: int
|
||||
total_tokens: int
|
||||
prompt_tokens_details: PromptTokensDetails | None = None
|
||||
completion_tokens_details: CompletionTokensDetails | None = None
|
||||
|
||||
|
||||
|
||||
class ChatCompletionResponse(BaseModel):
|
||||
id: str
|
||||
object: Literal["chat.completion"] = "chat.completion"
|
||||
created: int
|
||||
model: str
|
||||
choices: list[ChatCompletionChoice | StreamingChoiceResponse]
|
||||
usage: Usage | None = None
|
||||
service_tier: str | None = None
|
||||
|
||||
|
||||
class ChatCompletionTaskParams(BaseModel):
|
||||
model: str
|
||||
frequency_penalty: float | None = None
|
||||
@@ -31,4 +95,4 @@ class ChatCompletionTaskParams(BaseModel):
|
||||
tools: list[dict[str, Any]] | None = None
|
||||
tool_choice: str | dict[str, Any] | None = None
|
||||
parallel_tool_calls: bool | None = None
|
||||
user: str | None = None
|
||||
user: str | None = None
|
||||
|
||||
@@ -3,7 +3,7 @@ from typing import Literal
|
||||
from shared.topology import Connection, ConnectionProfile, Node, NodePerformanceProfile
|
||||
from shared.types.common import NodeId
|
||||
from shared.types.events.chunks import GenerationChunk
|
||||
from shared.types.request import RequestId
|
||||
from shared.types.events.commands import CommandId
|
||||
from shared.types.tasks import Task, TaskId, TaskStatus
|
||||
from shared.types.worker.common import InstanceId, NodeStatus
|
||||
from shared.types.worker.instances import InstanceParams, TypeOfInstance
|
||||
@@ -101,7 +101,7 @@ class WorkerDisconnected(_BaseEvent[_EventType.WorkerDisconnected]):
|
||||
|
||||
class ChunkGenerated(_BaseEvent[_EventType.ChunkGenerated]):
|
||||
event_type: Literal[_EventType.ChunkGenerated] = _EventType.ChunkGenerated
|
||||
request_id: RequestId
|
||||
command_id: CommandId
|
||||
chunk: GenerationChunk
|
||||
|
||||
|
||||
|
||||
@@ -4,8 +4,8 @@ from typing import Annotated, Literal
|
||||
from pydantic import BaseModel, Field, TypeAdapter
|
||||
|
||||
from shared.openai_compat import FinishReason
|
||||
from shared.types.events.commands import CommandId
|
||||
from shared.types.models import ModelId
|
||||
from shared.types.request import RequestId
|
||||
|
||||
|
||||
class ChunkType(str, Enum):
|
||||
@@ -15,7 +15,7 @@ class ChunkType(str, Enum):
|
||||
|
||||
class BaseChunk[ChunkTypeT: ChunkType](BaseModel):
|
||||
chunk_type: ChunkTypeT
|
||||
request_id: RequestId
|
||||
command_id: CommandId
|
||||
idx: int
|
||||
model: ModelId
|
||||
|
||||
|
||||
@@ -1,19 +1,12 @@
|
||||
from enum import Enum
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Callable,
|
||||
Sequence,
|
||||
)
|
||||
from typing import Annotated, Callable, Sequence
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, Field, TypeAdapter
|
||||
|
||||
from shared.types.api import ChatCompletionTaskParams
|
||||
from shared.types.common import NewUUID
|
||||
from shared.types.state import State
|
||||
|
||||
from . import Event
|
||||
from shared.types.events import Event
|
||||
from shared.types.state import InstanceId, State
|
||||
|
||||
|
||||
class CommandId(NewUUID):
|
||||
@@ -21,19 +14,36 @@ class CommandId(NewUUID):
|
||||
|
||||
|
||||
class CommandTypes(str, Enum):
|
||||
Create = "Create"
|
||||
Update = "Update"
|
||||
Delete = "Delete"
|
||||
CHAT_COMPLETION = "CHAT_COMPLETION"
|
||||
CREATE_INSTANCE = "CREATE_INSTANCE"
|
||||
DELETE_INSTANCE = "DELETE_INSTANCE"
|
||||
|
||||
|
||||
class Command[
|
||||
CommandType: CommandTypes,
|
||||
](BaseModel):
|
||||
command_type: CommandType
|
||||
class _BaseCommand[T: CommandTypes](BaseModel):
|
||||
command_id: CommandId
|
||||
command_type: T
|
||||
|
||||
|
||||
type Decide[CommandTypeT: CommandTypes] = Callable[
|
||||
[State, Command[CommandTypeT]],
|
||||
class ChatCompletionCommand(_BaseCommand[CommandTypes.CHAT_COMPLETION]):
|
||||
request_params: ChatCompletionTaskParams
|
||||
|
||||
|
||||
class CreateInstanceCommand(_BaseCommand[CommandTypes.CREATE_INSTANCE]):
|
||||
model_id: str
|
||||
|
||||
|
||||
class DeleteInstanceCommand(_BaseCommand[CommandTypes.DELETE_INSTANCE]):
|
||||
instance_id: InstanceId
|
||||
|
||||
|
||||
Command = Annotated[
|
||||
ChatCompletionCommand, Field(discriminator="command_type")
|
||||
]
|
||||
|
||||
CommandParser: TypeAdapter[Command] = TypeAdapter(Command)
|
||||
|
||||
|
||||
type Decide = Callable[
|
||||
[State, Command],
|
||||
Sequence[Event],
|
||||
]
|
||||
|
||||
@@ -1,12 +0,0 @@
|
||||
from pydantic import BaseModel
|
||||
|
||||
from shared.types.api import ChatCompletionTaskParams
|
||||
from shared.types.common import NewUUID
|
||||
|
||||
|
||||
class RequestId(NewUUID):
|
||||
pass
|
||||
|
||||
class APIRequest(BaseModel):
|
||||
request_id: RequestId
|
||||
request_params: ChatCompletionTaskParams
|
||||
@@ -1,21 +1,17 @@
|
||||
from collections.abc import Mapping, Sequence
|
||||
from enum import Enum
|
||||
from typing import List
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from shared.topology import Topology
|
||||
from shared.types.common import NodeId
|
||||
from shared.types.profiling import NodePerformanceProfile
|
||||
from shared.types.tasks import Task, TaskId, TaskSagaEntry
|
||||
from shared.types.tasks import Task, TaskId
|
||||
from shared.types.worker.common import InstanceId, NodeStatus
|
||||
from shared.types.worker.instances import BaseInstance
|
||||
from shared.types.worker.runners import RunnerId, RunnerStatus
|
||||
|
||||
|
||||
class ExternalCommand(BaseModel): ...
|
||||
|
||||
|
||||
class CachePolicy(str, Enum):
|
||||
KEEP_ALL = "KEEP_ALL"
|
||||
|
||||
@@ -26,11 +22,8 @@ class State(BaseModel):
|
||||
instances: Mapping[InstanceId, BaseInstance] = {}
|
||||
runners: Mapping[RunnerId, RunnerStatus] = {}
|
||||
tasks: Mapping[TaskId, Task] = {}
|
||||
task_sagas: Mapping[TaskId, Sequence[TaskSagaEntry]] = {}
|
||||
node_profiles: Mapping[NodeId, NodePerformanceProfile] = {}
|
||||
topology: Topology = Topology()
|
||||
history: Sequence[Topology] = []
|
||||
task_inbox: List[Task] = Field(default_factory=list)
|
||||
task_outbox: List[Task] = Field(default_factory=list)
|
||||
cache_policy: CachePolicy = CachePolicy.KEEP_ALL
|
||||
last_event_applied_idx: int = Field(default=0, ge=0)
|
||||
|
||||
+11
-12
@@ -1,6 +1,7 @@
|
||||
from enum import Enum
|
||||
from typing import Annotated
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from shared.types.api import ChatCompletionTaskParams
|
||||
from shared.types.common import NewUUID
|
||||
@@ -10,25 +11,23 @@ from shared.types.worker.common import InstanceId
|
||||
class TaskId(NewUUID):
|
||||
pass
|
||||
|
||||
|
||||
class TaskType(str, Enum):
|
||||
CHAT_COMPLETION = "CHAT_COMPLETION"
|
||||
|
||||
|
||||
class TaskStatus(str, Enum):
|
||||
Pending = "Pending"
|
||||
Running = "Running"
|
||||
Complete = "Complete"
|
||||
Failed = "Failed"
|
||||
PENDING = "PENDING"
|
||||
RUNNING = "RUNNING"
|
||||
COMPLETE = "COMPLETE"
|
||||
FAILED = "FAILED"
|
||||
|
||||
|
||||
class Task(BaseModel):
|
||||
task_id: TaskId
|
||||
class ChatCompletionTask(BaseModel):
|
||||
task_type: TaskType
|
||||
task_id: TaskId
|
||||
instance_id: InstanceId
|
||||
task_status: TaskStatus
|
||||
task_params: ChatCompletionTaskParams
|
||||
|
||||
|
||||
|
||||
class TaskSagaEntry(BaseModel):
|
||||
task_id: TaskId
|
||||
instance_id: InstanceId
|
||||
Task = Annotated[ChatCompletionTask, Field(discriminator="task_type")]
|
||||
|
||||
+2
-5
@@ -243,7 +243,7 @@ class Worker:
|
||||
await queue.put(ChunkGenerated(
|
||||
# todo: at some point we will no longer have a bijection between task_id and row_id.
|
||||
# So we probably want to store a mapping between these two in our Worker object.
|
||||
request_id=chunk.request_id,
|
||||
command_id=chunk.command_id,
|
||||
chunk=chunk
|
||||
))
|
||||
|
||||
@@ -338,12 +338,9 @@ class Worker:
|
||||
async def _loop(self):
|
||||
while True:
|
||||
state_copy = self.state.model_copy(deep=False)
|
||||
state_copy.task_inbox = []
|
||||
state_copy.task_outbox = []
|
||||
|
||||
op: RunnerOp | None = self.plan(state_copy)
|
||||
|
||||
# Run the op, synchronously blocking for now.
|
||||
# run the op, synchronously blocking for now
|
||||
if op is not None:
|
||||
async for event in self._execute_op(op):
|
||||
print(event)
|
||||
|
||||
@@ -6,7 +6,7 @@ from types import CoroutineType
|
||||
from typing import Any, Callable
|
||||
|
||||
from shared.types.events.chunks import GenerationChunk, TokenChunk
|
||||
from shared.types.request import RequestId
|
||||
from shared.types.events.commands import CommandId
|
||||
from shared.types.tasks import ChatCompletionTaskParams, Task
|
||||
from shared.types.worker.commands_runner import (
|
||||
ChatTaskMessage,
|
||||
@@ -181,7 +181,7 @@ class RunnerSupervisor:
|
||||
text=text, token=token, finish_reason=finish_reason
|
||||
):
|
||||
yield TokenChunk(
|
||||
request_id=RequestId(uuid=task.task_id.uuid),
|
||||
command_id=CommandId(uuid=task.task_id.uuid),
|
||||
idx=token,
|
||||
model=self.model_shard_meta.model_meta.model_id,
|
||||
text=text,
|
||||
|
||||
@@ -11,7 +11,7 @@ from shared.types.common import NodeId
|
||||
from shared.types.models import ModelId, ModelMetadata
|
||||
from shared.types.state import State
|
||||
from shared.types.tasks import (
|
||||
Task,
|
||||
ChatCompletionTask,
|
||||
TaskId,
|
||||
TaskStatus,
|
||||
TaskType,
|
||||
@@ -105,21 +105,13 @@ def completion_create_params(user_message: str) -> ChatCompletionTaskParams:
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def chat_completion_task(completion_create_params: ChatCompletionTaskParams) -> Task:
|
||||
"""Creates a ChatCompletionTask directly for serdes testing"""
|
||||
return Task(task_id=TaskId(), instance_id=InstanceId(), task_type=TaskType.CHAT_COMPLETION, task_status=TaskStatus.Pending, task_params=completion_create_params)
|
||||
|
||||
@pytest.fixture
|
||||
def chat_task(
|
||||
completion_create_params: ChatCompletionTaskParams,
|
||||
) -> Task:
|
||||
"""Creates the final Task object"""
|
||||
return Task(
|
||||
def chat_completion_task(completion_create_params: ChatCompletionTaskParams) -> ChatCompletionTask:
|
||||
return ChatCompletionTask(
|
||||
task_id=TaskId(),
|
||||
instance_id=InstanceId(),
|
||||
task_type=TaskType.CHAT_COMPLETION,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=completion_create_params,
|
||||
task_status=TaskStatus.PENDING,
|
||||
task_params=completion_create_params
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
|
||||
Reference in New Issue
Block a user