314 lines
12 KiB
Python
314 lines
12 KiB
Python
"""Command-line entry point for the Open Agent API integration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import platform
|
|
import sys
|
|
import time
|
|
from typing import Any, Dict, Optional, Sequence
|
|
|
|
import httpx
|
|
|
|
from agent_integration.client import (
|
|
OpenAgentAPIClient,
|
|
OpenAgentAPIError,
|
|
OpenAgentEvent,
|
|
OpenAgentProtocolError,
|
|
OpenAgentTransportError,
|
|
)
|
|
from agent_integration.config import AgentConfig, AgentConfigurationError
|
|
from agent_integration.events import AgentStreamNormalizer
|
|
from agent_integration.service import AgentResponseError, OpenAgentService, SessionNotFoundError
|
|
from agent_integration.sessions import SessionStoreError, SQLiteSessionStore
|
|
|
|
|
|
RUN_ACTIVE_STATUSES = {"pending", "running", "cancelling"}
|
|
NETWORK_COMMANDS = {"ensure-session", "chat", "send", "get-run", "cancel-run"}
|
|
|
|
|
|
def parse_metadata(value: str) -> Dict[str, Any]:
|
|
try:
|
|
parsed = json.loads(value)
|
|
except json.JSONDecodeError as exc:
|
|
raise argparse.ArgumentTypeError(f"metadata must be valid JSON: {exc}") from exc
|
|
if not isinstance(parsed, dict):
|
|
raise argparse.ArgumentTypeError("metadata must be a JSON object")
|
|
return parsed
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description="Connect this project to the DeerFlow Open Agent API.")
|
|
parser.add_argument("--base-url", help="Override DEERFLOW_BASE_URL (the API key remains env-only).")
|
|
parser.add_argument("--auth-mode", choices=["bearer", "x-api-key"])
|
|
parser.add_argument("--session-db", help="Override DEERFLOW_SESSION_DB.")
|
|
parser.add_argument("--timeout", type=float, help="Overall request timeout in seconds.")
|
|
parser.add_argument("--connect-timeout", type=float, help="Connection timeout in seconds.")
|
|
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
subparsers.add_parser("doctor", help="Validate local configuration without making a network call.")
|
|
|
|
ensure_parser = subparsers.add_parser("ensure-session", help="Create or reuse an Agent session.")
|
|
_add_conversation_args(ensure_parser, include_message=False)
|
|
|
|
chat_parser = subparsers.add_parser("chat", help="Stream one message as SSE-derived output.")
|
|
_add_conversation_args(chat_parser, include_message=True)
|
|
chat_parser.add_argument(
|
|
"--output",
|
|
choices=["jsonl", "text", "raw-jsonl"],
|
|
default="jsonl",
|
|
help=(
|
|
"jsonl emits filtered public events; text emits visible answer text; "
|
|
"raw-jsonl exposes unfiltered server internals for diagnostics only."
|
|
),
|
|
)
|
|
|
|
send_parser = subparsers.add_parser("send", help="Submit a non-streaming run.")
|
|
_add_conversation_args(send_parser, include_message=True)
|
|
send_parser.add_argument("--poll", action="store_true")
|
|
send_parser.add_argument("--poll-interval", type=float, default=2.0)
|
|
send_parser.add_argument("--poll-timeout", type=float, default=120.0)
|
|
|
|
get_parser = subparsers.add_parser("get-run", help="Fetch a run by local conversation.")
|
|
get_parser.add_argument("--conversation-id", required=True)
|
|
get_parser.add_argument("--run-id", required=True)
|
|
|
|
cancel_parser = subparsers.add_parser("cancel-run", help="Cancel a run by local conversation.")
|
|
cancel_parser.add_argument("--conversation-id", required=True)
|
|
cancel_parser.add_argument("--run-id", required=True)
|
|
|
|
show_parser = subparsers.add_parser("show-session", help="Show the local session mapping only.")
|
|
show_parser.add_argument("--conversation-id", required=True)
|
|
|
|
forget_parser = subparsers.add_parser("forget-session", help="Delete a stale local session mapping.")
|
|
forget_parser.add_argument("--conversation-id", required=True)
|
|
return parser
|
|
|
|
|
|
def _add_conversation_args(parser: argparse.ArgumentParser, *, include_message: bool) -> None:
|
|
parser.add_argument("--conversation-id", required=True, help="Stable opaque upstream conversation ID.")
|
|
parser.add_argument("--external-subject-id", help="Stable opaque user/customer ID.")
|
|
parser.add_argument("--metadata", type=parse_metadata, default={})
|
|
if include_message:
|
|
parser.add_argument("--message", required=True)
|
|
parser.add_argument(
|
|
"--message-id",
|
|
help="Stable upstream message ID; pass the same value when safely retrying.",
|
|
)
|
|
|
|
|
|
def _config_for_args(args: argparse.Namespace, *, require_api_key: bool) -> AgentConfig:
|
|
env = dict(os.environ)
|
|
overrides = {
|
|
"DEERFLOW_BASE_URL": args.base_url,
|
|
"DEERFLOW_AUTH_MODE": args.auth_mode,
|
|
"DEERFLOW_SESSION_DB": args.session_db,
|
|
"DEERFLOW_TIMEOUT_SECONDS": args.timeout,
|
|
"DEERFLOW_CONNECT_TIMEOUT_SECONDS": args.connect_timeout,
|
|
}
|
|
for key, value in overrides.items():
|
|
if value is not None:
|
|
env[key] = str(value)
|
|
return AgentConfig.from_env(env, require_api_key=require_api_key)
|
|
|
|
|
|
def main(argv: Optional[Sequence[str]] = None) -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
try:
|
|
config = _config_for_args(args, require_api_key=args.command in NETWORK_COMMANDS)
|
|
if args.command == "doctor":
|
|
report = config.redacted_summary()
|
|
report.update(
|
|
{
|
|
"status": "ready" if config.api_key_configured else "missing_api_key",
|
|
"python_version": platform.python_version(),
|
|
"httpx_version": httpx.__version__,
|
|
"network_checked": False,
|
|
}
|
|
)
|
|
_print_json(report)
|
|
return 0 if config.api_key_configured else 2
|
|
|
|
with SQLiteSessionStore(config.session_db) as store:
|
|
if args.command == "show-session":
|
|
record = store.get(args.conversation_id)
|
|
if record is None:
|
|
raise SessionNotFoundError("no local Agent session mapping was found")
|
|
_print_json(record.to_dict())
|
|
return 0
|
|
if args.command == "forget-session":
|
|
_print_json(
|
|
{
|
|
"conversation_id": args.conversation_id,
|
|
"deleted": store.delete(args.conversation_id),
|
|
}
|
|
)
|
|
return 0
|
|
|
|
if not config.api_key_configured or not config.api_key:
|
|
raise AgentConfigurationError("DEERFLOW_OPEN_API_KEY is required")
|
|
with OpenAgentAPIClient(
|
|
base_url=config.base_url,
|
|
api_key=config.api_key,
|
|
auth_mode=config.auth_mode, # type: ignore[arg-type]
|
|
timeout=config.timeout_seconds,
|
|
connect_timeout=config.connect_timeout_seconds,
|
|
) as client:
|
|
service = OpenAgentService(client, store)
|
|
return _run_network_command(service, args)
|
|
except AgentConfigurationError as exc:
|
|
_print_error("configuration_error", str(exc))
|
|
return 2
|
|
except (ValueError, SessionNotFoundError, AgentResponseError, SessionStoreError) as exc:
|
|
_print_error("input_or_state_error", str(exc))
|
|
return 2
|
|
except OpenAgentAPIError as exc:
|
|
_print_error(
|
|
"api_error",
|
|
str(exc),
|
|
status_code=exc.status_code,
|
|
request_id=exc.request_id,
|
|
retryable=exc.retryable,
|
|
active_run_conflict=exc.active_run_conflict,
|
|
)
|
|
return 1
|
|
except (OpenAgentTransportError, OpenAgentProtocolError) as exc:
|
|
_print_error("transport_or_protocol_error", str(exc))
|
|
return 1
|
|
except TimeoutError as exc:
|
|
_print_error("poll_timeout", str(exc))
|
|
return 1
|
|
except KeyboardInterrupt:
|
|
_print_error("interrupted", "operation interrupted by user")
|
|
return 130
|
|
|
|
|
|
def _run_network_command(service: OpenAgentService, args: argparse.Namespace) -> int:
|
|
if args.command == "ensure-session":
|
|
record = service.ensure_session(
|
|
args.conversation_id,
|
|
external_subject_id=args.external_subject_id,
|
|
metadata=args.metadata,
|
|
)
|
|
_print_json(record.to_dict())
|
|
return 0
|
|
|
|
if args.command == "chat":
|
|
wrote_text = False
|
|
normalizer = AgentStreamNormalizer()
|
|
for raw_event in service.stream_message(
|
|
args.conversation_id,
|
|
args.message,
|
|
message_id=args.message_id,
|
|
external_subject_id=args.external_subject_id,
|
|
metadata=args.metadata,
|
|
):
|
|
public_events = normalizer.feed(raw_event)
|
|
if args.output == "raw-jsonl":
|
|
print(json.dumps(raw_event.to_dict(), ensure_ascii=False), flush=True)
|
|
continue
|
|
for public_event in public_events:
|
|
wrote_text = _emit_public_event(public_event, args.output, wrote_text)
|
|
for public_event in normalizer.finish():
|
|
if args.output != "raw-jsonl":
|
|
wrote_text = _emit_public_event(public_event, args.output, wrote_text)
|
|
if args.output == "text" and wrote_text:
|
|
print()
|
|
return 1 if normalizer.fatal_error else 0
|
|
|
|
if args.command == "send":
|
|
run = service.send_message(
|
|
args.conversation_id,
|
|
args.message,
|
|
message_id=args.message_id,
|
|
external_subject_id=args.external_subject_id,
|
|
metadata=args.metadata,
|
|
)
|
|
_print_json(run)
|
|
if args.poll:
|
|
run_id = run.get("run_id")
|
|
if not isinstance(run_id, str) or not run_id:
|
|
raise AgentResponseError("send response is missing run_id")
|
|
_poll_run(service, args.conversation_id, run_id, args.poll_interval, args.poll_timeout)
|
|
return 0
|
|
|
|
if args.command == "get-run":
|
|
_print_json(service.get_run(args.conversation_id, args.run_id))
|
|
return 0
|
|
|
|
if args.command == "cancel-run":
|
|
_print_json(service.cancel_run(args.conversation_id, args.run_id))
|
|
return 0
|
|
|
|
raise ValueError(f"unsupported command: {args.command}")
|
|
|
|
|
|
def _poll_run(
|
|
service: OpenAgentService,
|
|
conversation_id: str,
|
|
run_id: str,
|
|
poll_interval: float,
|
|
poll_timeout: float,
|
|
) -> None:
|
|
if poll_interval <= 0 or poll_timeout <= 0:
|
|
raise ValueError("poll interval and timeout must be greater than zero")
|
|
deadline = time.monotonic() + poll_timeout
|
|
while True:
|
|
status = service.get_run(conversation_id, run_id)
|
|
_print_json(status)
|
|
if str(status.get("status")) not in RUN_ACTIVE_STATUSES:
|
|
if status.get("status") == "success" and not status.get("final_content"):
|
|
_print_warning(
|
|
"run_result_unavailable",
|
|
"run succeeded but final_content is empty; use streaming or a configured webhook for the answer",
|
|
run_id=run_id,
|
|
)
|
|
return
|
|
if time.monotonic() >= deadline:
|
|
raise TimeoutError(f"run {run_id} did not finish within {poll_timeout} seconds")
|
|
time.sleep(poll_interval)
|
|
|
|
|
|
def _emit_public_event(event: OpenAgentEvent, output: str, wrote_text: bool) -> bool:
|
|
if output == "jsonl":
|
|
print(json.dumps(event.to_dict(), ensure_ascii=False), flush=True)
|
|
return wrote_text
|
|
|
|
if event.event in {"run.warning", "run.error"}:
|
|
print(json.dumps(event.to_dict(), ensure_ascii=False), file=sys.stderr, flush=True)
|
|
return wrote_text
|
|
if event.event == "message.delta" and isinstance(event.data, dict):
|
|
content = event.data.get("content")
|
|
if isinstance(content, str) and content:
|
|
print(content, end="", flush=True)
|
|
return True
|
|
if event.event == "message.completed" and isinstance(event.data, dict):
|
|
content = event.data.get("content")
|
|
if not event.data.get("streamed") and isinstance(content, str) and content:
|
|
print(content, end="", flush=True)
|
|
return True
|
|
return wrote_text
|
|
|
|
|
|
def _print_json(value: Any) -> None:
|
|
print(json.dumps(value, ensure_ascii=False, indent=2))
|
|
|
|
|
|
def _print_error(kind: str, message: str, **extra: Any) -> None:
|
|
payload: Dict[str, Any] = {"error": kind, "message": message}
|
|
payload.update({key: value for key, value in extra.items() if value is not None})
|
|
print(json.dumps(payload, ensure_ascii=False), file=sys.stderr)
|
|
|
|
|
|
def _print_warning(kind: str, message: str, **extra: Any) -> None:
|
|
payload: Dict[str, Any] = {"warning": kind, "message": message}
|
|
payload.update({key: value for key, value in extra.items() if value is not None})
|
|
print(json.dumps(payload, ensure_ascii=False), file=sys.stderr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|