355 lines
14 KiB
Python
355 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Callable, Literal, overload
|
|
|
|
from .client import HarnessClient, HarnessConfig, UserQuestionHandler
|
|
from .errors import SdkProtocolError, TransportClosedError
|
|
from .models import CommandExecution, JsonObject, Notification
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class DeepSeekHarnessConfig:
|
|
"""Configuration for launching the local DeepSeek Harness SDK runtime.
|
|
|
|
The runtime inherits the caller's environment by default, so existing
|
|
DEEPSEEK_API_KEY and DEEPSEEK_BASE_URL settings keep working. Use ``env`` to
|
|
intentionally override or inject variables for a subprocess.
|
|
These values are startup input, not a live Session state mirror; Session
|
|
commands do not rewrite this object.
|
|
"""
|
|
|
|
provider: str = "deepseek-official"
|
|
model: str = "deepseek-v4-flash"
|
|
reasoning_effort: str | None = None
|
|
max_tokens: int | None = None
|
|
cwd: str | None = None
|
|
runtime_cwd: str | None = None
|
|
dsh_bin: str | None = None
|
|
profile: str = "sdk"
|
|
patches: tuple[str, ...] = ()
|
|
dsh_home: str | None = None
|
|
env: dict[str, str] = field(default_factory=dict)
|
|
initialize_timeout_seconds: float = 30.0
|
|
request_timeout_seconds: float | None = None
|
|
shutdown_timeout_seconds: float | None = 1.0
|
|
base_url: str | None = None
|
|
api_key: str | None = None
|
|
on_user_question: UserQuestionHandler | None = None
|
|
approval_requests: bool = False
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class RunResult:
|
|
session_id: str
|
|
final_response: str
|
|
finish_reason: str | None
|
|
events: list[JsonObject]
|
|
notifications: list[Notification]
|
|
|
|
|
|
class _PromptActivity:
|
|
"""Collect an ordinary prompt's notifications from its inbox receipt through idle."""
|
|
|
|
def __init__(self, session_id: str, message_id: str) -> None:
|
|
self.session_id = session_id
|
|
self.message_id = message_id
|
|
self.received = False
|
|
self.events: list[JsonObject] = []
|
|
self.notifications: list[Notification] = []
|
|
|
|
def record(self, notification: Notification, on_notification: Callable[[Notification], None] | None = None) -> bool:
|
|
if not self.received:
|
|
if not _is_inbox_receipt(notification, self.session_id, self.message_id):
|
|
return False
|
|
self.received = True
|
|
self.notifications.append(notification)
|
|
if on_notification is not None:
|
|
on_notification(notification)
|
|
if notification.method == "session.event" and notification.payload.get("sessionId") == self.session_id:
|
|
event = notification.payload.get("event")
|
|
if isinstance(event, dict):
|
|
self.events.append(event)
|
|
return True
|
|
|
|
def result(self) -> RunResult:
|
|
return RunResult(
|
|
session_id=self.session_id, final_response=final_response(self.events),
|
|
finish_reason=finish_reason(self.events), events=self.events, notifications=self.notifications,
|
|
)
|
|
|
|
|
|
class DeepSeekHarness:
|
|
"""Reusable synchronous SDK for running DeepSeek Harness agent turns.
|
|
|
|
The runtime subprocess starts lazily and remains owned by this instance
|
|
across calls to :meth:`run`. Use the instance as a context manager, or call
|
|
:meth:`close` explicitly when finished, so the subprocess is always reaped.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
config: DeepSeekHarnessConfig | None = None,
|
|
*,
|
|
_launch_args: tuple[str, ...] | None = None,
|
|
**kwargs: object,
|
|
) -> None:
|
|
if config is not None and kwargs:
|
|
raise TypeError("pass either DeepSeekHarnessConfig or keyword options, not both")
|
|
self.config = config or DeepSeekHarnessConfig(**kwargs)
|
|
cwd = str(Path(self.config.cwd or Path.cwd()).resolve())
|
|
runtime_cwd = str(Path(self.config.runtime_cwd).resolve()) if self.config.runtime_cwd is not None else cwd
|
|
self._cwd = cwd
|
|
env = dict(self.config.env)
|
|
if self.config.base_url is not None:
|
|
env["DEEPSEEK_BASE_URL"] = self.config.base_url
|
|
if self.config.api_key is not None:
|
|
env["DEEPSEEK_API_KEY"] = self.config.api_key
|
|
|
|
self._client = HarnessClient(
|
|
HarnessConfig(
|
|
dsh_bin=self.config.dsh_bin,
|
|
profile=self.config.profile,
|
|
patches=self.config.patches,
|
|
dsh_home=self.config.dsh_home,
|
|
cwd=runtime_cwd,
|
|
env=env,
|
|
initialize_timeout_seconds=self.config.initialize_timeout_seconds,
|
|
request_timeout_seconds=self.config.request_timeout_seconds,
|
|
shutdown_timeout_seconds=self.config.shutdown_timeout_seconds,
|
|
on_user_question=self.config.on_user_question,
|
|
),
|
|
_launch_args=_launch_args,
|
|
)
|
|
self._initialized = False
|
|
|
|
def __enter__(self) -> "DeepSeekHarness":
|
|
self.start()
|
|
return self
|
|
|
|
def __exit__(self, _exc_type, _exc, _tb) -> None:
|
|
self.close()
|
|
|
|
@property
|
|
def client(self) -> HarnessClient:
|
|
return self._client
|
|
|
|
def start(self) -> None:
|
|
if self._initialized:
|
|
return
|
|
self._client.start()
|
|
self._client.initialize(
|
|
cwd=self._cwd,
|
|
provider=self.config.provider,
|
|
model=self.config.model,
|
|
reasoning_effort=self.config.reasoning_effort,
|
|
max_tokens=self.config.max_tokens,
|
|
approval_requests=self.config.approval_requests,
|
|
)
|
|
self._initialized = True
|
|
|
|
def close(self) -> None:
|
|
self._client.close()
|
|
self._initialized = False
|
|
|
|
def start_session(self, session_id: str | None = None) -> "Session":
|
|
self.start()
|
|
return Session(self, session_id or f"session-{uuid.uuid4().hex}")
|
|
|
|
def resume_session(self, session_id: str) -> "Session":
|
|
"""Restore persisted history from the configured Harness home without sending input.
|
|
|
|
Missing, unreadable, incompatible, or occupied sessions raise instead of creating a new session.
|
|
"""
|
|
self.start()
|
|
self._client.resume_session(session_id)
|
|
return Session(self, session_id)
|
|
|
|
@overload
|
|
def run(
|
|
self, input: list[JsonObject], *, session_id: str | None = None,
|
|
on_notification: Callable[[Notification], None] | None = None,
|
|
) -> RunResult: ...
|
|
|
|
@overload
|
|
def run(
|
|
self, input: str | list[JsonObject], *, session_id: str | None = None,
|
|
on_notification: Callable[[Notification], None] | None = None,
|
|
delivery: Literal["queue", "steer"] = "queue",
|
|
) -> RunResult | CommandExecution: ...
|
|
|
|
def run(
|
|
self,
|
|
input: str | list[JsonObject],
|
|
*,
|
|
session_id: str | None = None,
|
|
on_notification: Callable[[Notification], None] | None = None,
|
|
delivery: Literal["queue", "steer"] = "queue",
|
|
) -> RunResult | CommandExecution:
|
|
"""Run text or a native command on a fresh or named Session."""
|
|
return self.start_session(session_id).run(input, on_notification=on_notification, delivery=delivery)
|
|
|
|
|
|
class Session:
|
|
def __init__(self, harness: DeepSeekHarness, session_id: str) -> None:
|
|
self.harness = harness
|
|
self.id = session_id
|
|
|
|
def get_working_directory(self) -> str:
|
|
"""Read the effective directory, recovering a missing directory to the Session origin."""
|
|
self.harness.start()
|
|
return self.harness.client.get_working_directory(self.id)
|
|
|
|
def set_working_directory(self, path: str) -> str:
|
|
"""Change the effective directory while existing processes retain their own directories."""
|
|
self.harness.start()
|
|
return self.harness.client.set_working_directory(self.id, path)
|
|
|
|
@overload
|
|
def run(
|
|
self, input: list[JsonObject], *,
|
|
on_notification: Callable[[Notification], None] | None = None,
|
|
) -> RunResult: ...
|
|
|
|
@overload
|
|
def run(
|
|
self, input: str | list[JsonObject], *,
|
|
on_notification: Callable[[Notification], None] | None = None,
|
|
delivery: Literal["queue", "steer"] = "queue",
|
|
) -> RunResult | CommandExecution: ...
|
|
|
|
def run(
|
|
self,
|
|
input: str | list[JsonObject],
|
|
*,
|
|
on_notification: Callable[[Notification], None] | None = None,
|
|
delivery: Literal["queue", "steer"] = "queue",
|
|
) -> RunResult | CommandExecution:
|
|
"""Run text through native input routing, or queue literal content blocks.
|
|
|
|
Text returns the prompt's receipt-to-idle RunResult or a settled
|
|
CommandExecution; command completion does not wait for later model work.
|
|
Content blocks stay literal and accept only queue delivery. Notifications
|
|
retain wire order. Command callbacks replay Session-tree notifications
|
|
from subscription through their own command/done after the response,
|
|
including overlapping work. Handler errors remain command results;
|
|
routing, transport and observer failures raise without rolling back work.
|
|
"""
|
|
if delivery not in ("queue", "steer"):
|
|
raise ValueError("run delivery must be queue or steer")
|
|
if not isinstance(input, str) and delivery != "queue":
|
|
raise ValueError("content-block input requires queue delivery")
|
|
self.harness.start()
|
|
with self.harness.client.subscribe_session_notifications(self.id) as subscription:
|
|
if isinstance(input, str):
|
|
response = self.harness.client._request_raw(
|
|
"session/query", {"sessionId": self.id, "text": input, "delivery": delivery},
|
|
)
|
|
if not isinstance(response, dict):
|
|
raise SdkProtocolError("session/query response must be a JSON object")
|
|
if response.get("kind") == "command":
|
|
execution = CommandExecution.model_validate(response.get("execution"))
|
|
while True:
|
|
try:
|
|
notification = subscription._next_queued()
|
|
except TransportClosedError:
|
|
break
|
|
if notification is None:
|
|
break
|
|
if on_notification is not None:
|
|
on_notification(notification)
|
|
if _is_command_done(notification, self.id, execution.command_id):
|
|
break
|
|
return execution
|
|
if response.get("kind") != "prompt" or not isinstance(response.get("messageId"), str) or not response["messageId"]:
|
|
raise SdkProtocolError("session/query requires a prompt messageId or command execution result")
|
|
message_id = response["messageId"]
|
|
else:
|
|
message_id = self.harness.client.session_prompt(
|
|
self.id, input, notification_subscription=subscription,
|
|
)
|
|
activity = _PromptActivity(self.id, message_id)
|
|
while True:
|
|
notification = subscription.next()
|
|
if not activity.record(notification, on_notification):
|
|
continue
|
|
if (
|
|
notification.method == "session.status"
|
|
and notification.payload.get("sessionId") == self.id
|
|
and notification.payload.get("status") == "idle"
|
|
):
|
|
break
|
|
|
|
return activity.result()
|
|
|
|
|
|
def _is_inbox_receipt(notification: Notification, session_id: str, message_id: str) -> bool:
|
|
if notification.method != "session.event" or notification.payload.get("sessionId") != session_id:
|
|
return False
|
|
event = notification.payload.get("event")
|
|
if not isinstance(event, dict) or event.get("type") != "agent/inbox/spliced":
|
|
return False
|
|
data = event.get("data")
|
|
inserted = data.get("inserted") if isinstance(data, dict) else None
|
|
return isinstance(inserted, list) and any(
|
|
isinstance(message, dict) and message.get("id") == message_id for message in inserted
|
|
)
|
|
|
|
|
|
def _is_command_done(notification: Notification, session_id: str, command_id: str) -> bool:
|
|
if notification.method != "session.event" or notification.payload.get("sessionId") != session_id:
|
|
return False
|
|
event = notification.payload.get("event")
|
|
if not isinstance(event, dict) or event.get("type") != "command/done":
|
|
return False
|
|
data = event.get("data")
|
|
return isinstance(data, dict) and data.get("commandId") == command_id
|
|
|
|
|
|
def normalize_input(input: str | list[JsonObject]) -> list[JsonObject]:
|
|
if isinstance(input, str):
|
|
return [{"type": "text", "text": input}]
|
|
return input
|
|
|
|
|
|
def final_response(events: list[JsonObject]) -> str:
|
|
for event in reversed(events):
|
|
if event.get("type") != "assistant/message":
|
|
continue
|
|
data = event.get("data")
|
|
if not isinstance(data, dict):
|
|
continue
|
|
message = data.get("message")
|
|
content_owner = message if isinstance(message, dict) else data
|
|
content = content_owner.get("content")
|
|
if not isinstance(content, list):
|
|
continue
|
|
parts: list[str] = []
|
|
for block in content:
|
|
if isinstance(block, dict) and block.get("type") == "text":
|
|
parts.append(str(block.get("text") or ""))
|
|
return "".join(parts)
|
|
return ""
|
|
|
|
|
|
def finish_reason(events: list[JsonObject]) -> str | None:
|
|
"""Return the last turn-ending kind.
|
|
|
|
The input must contain root-session events from one owned run interval.
|
|
|
|
Raises:
|
|
SdkProtocolError: The last ``turn/end`` has no string reason kind.
|
|
"""
|
|
for event in reversed(events):
|
|
if event.get("type") != "turn/end":
|
|
continue
|
|
data = event.get("data")
|
|
reason = data.get("reason") if isinstance(data, dict) else None
|
|
kind = reason.get("kind") if isinstance(reason, dict) else None
|
|
if not isinstance(kind, str):
|
|
raise SdkProtocolError("turn/end event requires a string data.reason.kind")
|
|
return kind
|
|
return None
|