diff --git a/danyapi/accounts.py b/danyapi/accounts.py index fd348f5..e90d259 100644 --- a/danyapi/accounts.py +++ b/danyapi/accounts.py @@ -10,7 +10,7 @@ from .deepseek.client import DeepSeekClient from .pow import PowManager from .sessions import SessionRegistry -from .store import JsonStore +from .store import _MAX_AFFINITY, JsonStore log = logging.getLogger("danyapi.accounts") @@ -161,7 +161,7 @@ def _touch(self, session_id: str, now: float) -> None: class DeepSeekAccount: - __slots__ = ("broken", "client", "index", "pow", "pow_upload", "sem", "sessions", "stable_id") + __slots__ = ("broken", "broken_at", "client", "index", "pow", "pow_upload", "sem", "sessions", "stable_id") def __init__( self, @@ -180,10 +180,12 @@ def __init__( self.sessions = SessionRegistry(client, session_cache_size, ttl, store=store, key_prefix=f"{index}:") self.stable_id = stable_id self.broken = False + self.broken_at: float | None = None def mark_broken(self) -> None: if not self.broken: self.broken = True + self.broken_at = time.monotonic() log.warning("account #%d marked broken (invalid/expired token)", self.index) @property @@ -193,13 +195,17 @@ def label(self) -> str: class _PoolAccount(Protocol): broken: bool + broken_at: float | None sem: asyncio.Semaphore + label: str AccountT = TypeVar("AccountT", bound=_PoolAccount) class AccountPool(Generic[AccountT]): + _REVIVE_COOLDOWN = 300.0 + def __init__( self, accounts: list[AccountT], @@ -220,6 +226,7 @@ def __init__( self._rr = 0 self._ttl = max(0.0, ttl) self._affinity_store = affinity_store + self._affinity_lock = threading.Lock() self._contexts = ContextIndex(session_cache_size, ttl, store=context_store) self._restore_affinities() @@ -259,18 +266,27 @@ def healthy(self) -> list[AccountT]: def register(self, account_index: int, session_id: str) -> None: now = time.monotonic() - self._by_session[session_id] = (account_index, now) - if self._affinity_store is not None: - self._affinity_store.set(session_id, self._affinity_record(account_index)) - if self._ttl > 0 and len(self._by_session) > max(4096, len(self.accounts) * 1024): - stale = [sid for sid, (_, ts) in self._by_session.items() if now - ts > self._ttl] - for sid in stale: - self._by_session.pop(sid, None) + record = self._affinity_record(account_index) + with self._affinity_lock: + self._by_session[session_id] = (account_index, now) + if self._affinity_store is not None: + if self._affinity_store.get(session_id) != record: + self._affinity_store.set(session_id, record) + while len(self._by_session) > _MAX_AFFINITY: + oldest = next(iter(self._by_session)) + self._by_session.pop(oldest, None) if self._affinity_store is not None: - self._affinity_store.discard(sid) + self._affinity_store.discard(oldest) + if self._ttl > 0 and len(self._by_session) > max(4096, len(self.accounts) * 1024): + stale = [sid for sid, (_, ts) in self._by_session.items() if now - ts > self._ttl] + for sid in stale: + self._by_session.pop(sid, None) + if self._affinity_store is not None: + self._affinity_store.discard(sid) def forget(self, session_id: str) -> None: - self._by_session.pop(session_id, None) + with self._affinity_lock: + self._by_session.pop(session_id, None) if self._affinity_store is not None: self._affinity_store.discard(session_id) @@ -284,35 +300,38 @@ def forget_context(self, session_id: str) -> None: self._contexts.forget(session_id) def account_for_session(self, session_id: str) -> AccountT | None: - entry = self._by_session.get(session_id) - if entry is None: - return None - idx, ts = entry - now = time.monotonic() - if self._ttl > 0 and now - ts > self._ttl: - self._by_session.pop(session_id, None) - self._contexts.forget(session_id) - if self._affinity_store is not None: - self._affinity_store.discard(session_id) - return None - acct = self.accounts[idx] - if acct is None or acct.broken: - self._by_session.pop(session_id, None) - self._contexts.forget(session_id) - if self._affinity_store is not None: - self._affinity_store.discard(session_id) - return None - if self._ttl > 0 and now != ts: - self._by_session[session_id] = (idx, now) - return acct + with self._affinity_lock: + entry = self._by_session.get(session_id) + if entry is None: + return None + idx, ts = entry + now = time.monotonic() + if self._ttl > 0 and now - ts > self._ttl: + self._by_session.pop(session_id, None) + self._contexts.forget(session_id) + if self._affinity_store is not None: + self._affinity_store.discard(session_id) + return None + acct = self.accounts[idx] + if acct is None or acct.broken: + self._by_session.pop(session_id, None) + self._contexts.forget(session_id) + if self._affinity_store is not None: + self._affinity_store.discard(session_id) + return None + if self._ttl > 0 and now != ts: + self._by_session[session_id] = (idx, now) + return acct def stats(self) -> dict[str, Any]: + with self._affinity_lock: + affinities = len(self._by_session) return { "label": self.label, "accounts": len(self.accounts), "healthy": len(self.healthy), "broken": len(self.accounts) - len(self.healthy), - "session_affinities": len(self._by_session), + "session_affinities": affinities, "context_entries": self._contexts.size, "context_hits": self._contexts.hits, "context_misses": self._contexts.misses, @@ -323,6 +342,11 @@ def stats(self) -> dict[str, Any]: async def acquire(self, session_id: str | None, max_wait: float | None = None) -> tuple[AccountT, str | None]: healthy = self.healthy if not healthy: + revived = await self.revive_broken() + if revived is not None: + healthy = [revived] + elif any(getattr(acct, "broken_at", None) is not None for acct in self.accounts): + raise AccountPoolBusy() raise RuntimeError(f"all {self.label} accounts are unavailable") if session_id: acct = self.account_for_session(session_id) @@ -367,6 +391,30 @@ async def _wait_free( raise AccountPoolBusy() await asyncio.sleep(0.05) + async def revive_broken(self) -> AccountT | None: + now = time.monotonic() + candidates: list[AccountT] = [] + for acct in self.accounts: + if not acct.broken: + continue + broken_at = getattr(acct, "broken_at", None) + if broken_at is not None and now - broken_at >= self._REVIVE_COOLDOWN: + candidates.append(acct) + for acct in candidates: + client = getattr(acct, "client", None) + if client is None: + continue + try: + ok = await client.check_auth() + except Exception: + ok = False + if ok: + acct.broken = False + acct.broken_at = None + log.info("%s revived after auth recheck", acct.label) + return acct + return None + def add_account(self, account: AccountT) -> None: idx = len(self.accounts) self.accounts.append(account) diff --git a/danyapi/api/openai.py b/danyapi/api/openai.py index e208507..e832e46 100644 --- a/danyapi/api/openai.py +++ b/danyapi/api/openai.py @@ -68,9 +68,9 @@ "FINISHED": "stop", "CONTEXT_LENGTH_EXCEEDED": "length", "CONTENT_FILTER": "content_filter", - "INCOMPLETE": "stop", - "WIP": "stop", - "TIMEOUT": "stop", + "INCOMPLETE": "length", + "WIP": "length", + "TIMEOUT": "length", } CONTEXT_LENGTH_STATUS = "CONTEXT_LENGTH_EXCEEDED" @@ -369,7 +369,7 @@ async def _fetch_qwen_models(client: QwenClient) -> list[dict]: app.add_middleware( CORSMiddleware, allow_origins=["*"], - allow_credentials=False, + allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) @@ -465,6 +465,15 @@ def _responses_store() -> JsonStore: return _shared_store("responses_store", "responses") +def _pool_account_by_stable(pool: AccountPool | None, stable_id: str) -> Any | None: + if pool is None: + return None + for acct in pool.accounts: + if getattr(acct, "stable_id", None) == stable_id: + return acct + return None + + @app.post("/v1/tokens") async def add_tokens(tokens: dict) -> dict: async with _TOKENS_LOCK: @@ -477,7 +486,55 @@ async def add_tokens(tokens: dict) -> dict: ds_candidates = [t for t in dict.fromkeys(new_ds) if t not in existing_ds] qw_candidates = [t for t in dict.fromkeys(new_qw) if t not in existing_qw] + pool: AccountPool | None = getattr(app.state, "pool", None) + qwen_pool: AccountPool | None = getattr(app.state, "qwen_pool", None) + + activated_ds = 0 + activated_qw = 0 + + for token in dict.fromkeys(new_ds): + if token in existing_ds: + acct = _pool_account_by_stable(pool, _token_stable_id(token)) + if acct is None or not acct.broken: + continue + try: + valid_auth = await acct.client.check_auth() + except Exception as exc: + log.warning("reactivation check failed for existing deepseek token: %s", exc) + continue + if not valid_auth: + continue + acct.broken = False + acct.broken_at = None + activated_ds += 1 + log.info("reactivated deepseek token (total accounts: %d)", len(pool.accounts) if pool else 0) + + for token in dict.fromkeys(new_qw): + if token in existing_qw: + acct = _pool_account_by_stable(qwen_pool, _token_stable_id(token)) + if acct is None or not acct.broken: + continue + try: + valid_auth = await acct.client.check_auth() + except Exception as exc: + log.warning("reactivation check failed for existing qwen token: %s", exc) + continue + if not valid_auth: + continue + acct.broken = False + acct.broken_at = None + activated_qw += 1 + log.info("reactivated qwen token (total accounts: %d)", len(qwen_pool.accounts) if qwen_pool else 0) + if not ds_candidates and not qw_candidates: + if activated_ds or activated_qw: + return { + "success": True, + "message": "Tokens reactivated.", + "added": {"deepseek": 0, "qwen": 0}, + "skipped": {"deepseek": 0, "qwen": 0}, + "reactivated": {"deepseek": activated_ds, "qwen": activated_qw}, + } raise HTTPException(400, "all provided tokens already exist") added_ds = 0 @@ -492,8 +549,6 @@ async def add_tokens(tokens: dict) -> dict: ds_affinity_store = _shared_store("deepseek_affinity_store", "deepseek-affinities") qw_affinity_store = _shared_store("qwen_affinity_store", "qwen-affinities") - pool: AccountPool | None = getattr(app.state, "pool", None) - qwen_pool: AccountPool | None = getattr(app.state, "qwen_pool", None) accepted_ds: list[str] = [] accepted_qw: list[str] = [] @@ -585,25 +640,71 @@ async def add_tokens(tokens: dict) -> dict: "message": "Tokens added and activated." if (added_ds or added_qw) else "No valid tokens to add.", "added": {"deepseek": added_ds, "qwen": added_qw}, "skipped": {"deepseek": skipped_ds, "qwen": skipped_qw}, + "reactivated": {"deepseek": activated_ds, "qwen": activated_qw}, } MAX_LOGGED_BODY = 256 * 1024 +MAX_REQUEST_BODY = 100 * 1024 * 1024 + + +async def _read_request_body(request: Request, limit: int) -> bytes: + raw_length = -1 + content_length = request.headers.get("content-length") + if content_length: + try: + raw_length = int(content_length) + except ValueError: + raw_length = -1 + if raw_length > 0: + if raw_length > limit: + raise HTTPException(413, "request body too large") + body = await request.body() + if len(body) > limit: + raise HTTPException(413, "request body too large") + request._body = body + return body + cached = getattr(request, "_body", None) + if cached: + if len(cached) > limit: + raise HTTPException(413, "request body too large") + return cached + chunks: list[bytes] = [] + total = 0 + async for chunk in request.stream(): + total += len(chunk) + if total > limit: + raise HTTPException(413, "request body too large") + chunks.append(chunk) + body = b"".join(chunks) + request._body = body + return body async def _extract_request_body(request: Request) -> dict[str, Any]: content_length = request.headers.get("content-length") if content_length: try: - if int(content_length) > MAX_LOGGED_BODY: + raw_length = int(content_length) + if raw_length <= 0: + return {} + if raw_length > MAX_REQUEST_BODY: + raise HTTPException(413, "request body too large") + if raw_length > MAX_LOGGED_BODY: return {} except ValueError: return {} + if getattr(request, "method", None) in ("GET", "DELETE", "HEAD", "OPTIONS"): + return {} try: - body = await request.body() + body = await _read_request_body(request, MAX_REQUEST_BODY) + except HTTPException: + raise except Exception: return {} - if not body or len(body) > MAX_LOGGED_BODY: + if not body: + return {} + if len(body) > MAX_LOGGED_BODY: return {} try: payload = json.loads(body) @@ -716,6 +817,23 @@ async def _log_requests(request: Request, call_next): MAX_FILES_PER_REQUEST = 50 MAX_FILE_SIZE = 100 * 1024 * 1024 +MAX_ATTACHMENT_TOTAL_SIZE = 10 * 1024 * 1024 + + +def _raw_data_uri_length(uri: str) -> int: + if not uri.startswith("data:"): + raise HTTPException(400, "image_url must be a data URI (data:;base64,...)") + _, _, payload = uri[5:].partition(",") + if not payload: + raise HTTPException(400, "invalid data URI: missing base64 payload") + compact = "".join(payload.split()).rstrip("=") + units, remainder = divmod(len(compact), 4) + decoded = units * 3 + if remainder == 2: + decoded += 1 + elif remainder == 3: + decoded += 2 + return decoded @dataclass @@ -743,6 +861,7 @@ def _split_data_uri(uri: str) -> tuple[str, bytes]: def _collect_attachments(req: ChatCompletionRequest) -> list[Attachment]: attachments: list[Attachment] = [] + raw_total = 0 for msg in req.messages: if not isinstance(msg.content, list): continue @@ -757,6 +876,9 @@ def _collect_attachments(req: ChatCompletionRequest) -> list[Attachment]: uri = image_url["url"] else: raise HTTPException(400, "invalid image_url value") + raw_total += _raw_data_uri_length(uri) + if raw_total > MAX_ATTACHMENT_TOTAL_SIZE: + raise HTTPException(413, "attachments too large") content_type, data = _split_data_uri(uri) name = f"image_{len(attachments)}.{content_type.split('/')[-1] or 'bin'}" attachments.append(Attachment(data, name, content_type, True)) @@ -791,8 +913,11 @@ async def _fresh_pow_upload_headers(account) -> dict: async def _upload_attachments(account, attachments: list[Attachment], model_type: str, thinking: bool) -> list[str]: file_ids: list[str] = [] - for att in attachments: - pow_headers = await _fresh_pow_upload_headers(account) + if attachments: + pow_headers_list = await asyncio.gather(*(_fresh_pow_upload_headers(account) for _ in attachments)) + else: + pow_headers_list = [] + for att, pow_headers in zip(attachments, pow_headers_list, strict=True): try: info = await account.client.upload_file( att.data, @@ -990,7 +1115,9 @@ async def _extract_request_api_key(request: Request) -> str | None: if not content_type.startswith("application/json"): return None try: - body = await request.body() + body = await _read_request_body(request, MAX_REQUEST_BODY) + except HTTPException: + raise except Exception: return None if not body: @@ -1012,6 +1139,10 @@ async def _close_pool(pool: Any) -> None: acct.sessions.close_all() except Exception as exc: log.info("session cleanup failed for byok account %r: %s", getattr(acct, "label", acct), exc) + sem = getattr(acct, "sem", None) + if sem is not None and sem.locked(): + log.info("skip client close for busy byok account %r", getattr(acct, "label", acct)) + continue try: await acct.client.aclose() except Exception as exc: @@ -1273,50 +1404,59 @@ async def _image_generations(req: ImageGenerationRequest, pool: AccountPool | No raise HTTPException(503, "qwen provider is not configured (required for image generation)") dims = _parse_image_size(req.size) + count = max(1, int(getattr(req, "n", 1) or 1)) account, existing_sid = await _acquire_account(pool, req.session_id) + want_b64 = req.response_format == "b64_json" + data: list[dict] = [] + usage = None + result_sid = existing_sid + revised_prompt = "" try: - result = await qwen_api.collect_image( - account=account, - pool=pool, - existing_sid=existing_sid, - lock=account.sem, - prompt=req.prompt, - model=req.model, - model_id=req.model, - user=req.user, - ) + for _ in range(count): + result = await qwen_api.collect_image( + account=account, + pool=pool, + existing_sid=result_sid, + lock=account.sem, + prompt=req.prompt, + model=req.model, + model_id=req.model, + user=req.user, + ) + result_sid = result.get("session_id") or result_sid + if result.get("usage"): + usage = result.get("usage") + if result.get("revised_prompt"): + revised_prompt = result["revised_prompt"] + for url in result["image_urls"]: + if not (want_b64 or dims): + data.append({"url": url}) + continue + try: + async with httpx.AsyncClient(follow_redirects=True, timeout=30) as hc: + img_resp = await hc.get(url) + if img_resp.status_code != 200: + log.warning("image download failed (%s) for %s, returning url", img_resp.status_code, url) + data.append({"url": url}) + continue + payload_bytes = _resize_image_bytes(img_resp.content, dims) + data.append({"b64_json": base64.b64encode(payload_bytes).decode()}) + except Exception as exc: + log.warning("image fetch failed for %s, returning url: %s", url, exc) + data.append({"url": url}) except AccountPoolBusy: raise HTTPException(429, "all accounts are busy, try again later") from None - want_b64 = req.response_format == "b64_json" - data: list[dict] = [] - for url in result["image_urls"]: - if not (want_b64 or dims): - data.append({"url": url}) - continue - try: - async with httpx.AsyncClient(follow_redirects=True, timeout=30) as hc: - img_resp = await hc.get(url) - if img_resp.status_code != 200: - log.warning("image download failed (%s) for %s, returning url", img_resp.status_code, url) - data.append({"url": url}) - continue - payload_bytes = _resize_image_bytes(img_resp.content, dims) - data.append({"b64_json": base64.b64encode(payload_bytes).decode()}) - except Exception as exc: - log.warning("image fetch failed for %s, returning url: %s", url, exc) - data.append({"url": url}) - if not data: - data.append({"url": "", "revised_prompt": result.get("revised_prompt", "")}) + data.append({"url": "", "revised_prompt": revised_prompt}) return { "created": int(time.time()), "data": data, - "usage": result.get("usage"), - "session_id": result.get("session_id"), + "usage": usage, + "session_id": result_sid, } @@ -1367,9 +1507,15 @@ def _include_usage(req: ChatCompletionRequest) -> bool: return bool(opts.get("include_usage")) -def _deepseek_usage(total: int, prompt: str = "") -> dict: +def _deepseek_usage(total: int, prompt: str = "", provider_usage: dict | None = None) -> dict: value = max(0, int(total or 0)) - prompt_tokens = estimate_tokens(prompt) + prompt_tokens = 0 + if isinstance(provider_usage, dict): + p_tokens = provider_usage.get("prompt_tokens") + if isinstance(p_tokens, int) and p_tokens > 0: + prompt_tokens = p_tokens + if not prompt_tokens: + prompt_tokens = estimate_tokens(prompt) return {"prompt_tokens": prompt_tokens, "completion_tokens": value, "total_tokens": prompt_tokens + value} @@ -1392,7 +1538,7 @@ def _stream_error_sse( session_key: str | None = None, error_finish: str | None = None, choice_finish: str | None = None, -) -> tuple[str, str, str]: +) -> tuple[str, str]: error: dict = {"message": message} if error_finish is not None: error["finish_reason"] = error_finish @@ -1402,19 +1548,12 @@ def _stream_error_sse( "object": "chat.completion.chunk", "created": created, "model": model, + "session_id": session_key, "error": error, "choices": [{"index": 0, "delta": {}, "finish_reason": choice_finish or error_finish or "error"}], } ) - tail_chunk = _sse( - { - "id": chunk_id, - "session_id": session_key, - "object": "chat.completion.chunk", - "choices": [], - } - ) - return error_chunk, tail_chunk, "data: [DONE]\n\n" + return error_chunk, "data: [DONE]\n\n" async def _stream_guard(gen, model: str): @@ -1496,6 +1635,14 @@ async def _chat_completions_qwen(req: ChatCompletionRequest, pool: AccountPool | account, existing_sid, context_seq, prompt, tool_mode = await _acquire_and_build(pool, req, {"model": req.model}) + attachments = _collect_attachments(req) + if attachments: + _validate_attachments(attachments) + for att in attachments: + if not att.is_image: + raise HTTPException(400, "qwen only supports image attachments, use deepseek for files") + prompt = f"{prompt}\n![image](data:{att.content_type};base64,{base64.b64encode(att.data).decode('ascii')})" + common = { "account": account, "pool": pool, @@ -1576,6 +1723,9 @@ async def _send_completion( except httpx.HTTPError as exc: raise HTTPException(502, f"DeepSeek request failed: {exc}") from exc + if resp is None or not hasattr(resp, "status_code"): + raise HTTPException(502, "unexpected provider response") + if resp.status_code != 200: body = await resp.aread() await resp.aclose() @@ -2135,7 +2285,7 @@ async def _collect_non_stream( raise HTTPException(502, _fake_context_error_body()) raise HTTPException(429, _busy_error_body(rec)) request_tokens = _advance_session_usage(session, rec.accumulated_tokens) - usage = _deepseek_usage(request_tokens, prompt) + usage = _deepseek_usage(request_tokens, prompt, rec.usage) account.sessions.touch_last_message(session_key, rec.id or response_message_id) record_usage( "deepseek", @@ -2204,6 +2354,8 @@ async def _stream_openai( response_message_id = None stop_message_id: str | None = None content_buf = "" + content_shown_len = 0 + tool_hidden = False role_sent = False started = time.monotonic() had_cached_session = bool(existing_sid) and account.sessions.get(existing_sid) is not None @@ -2301,6 +2453,19 @@ async def _stream_openai( if c_diff: if tool_mode: content_buf += c_diff + if not tool_hidden: + search_from = content_shown_len + marker_pos = content_buf.find('{"tool_calls"', search_from) + if marker_pos < 0: + marker_pos = content_buf.find(" search_from: + delta["content"] = content_buf[search_from:marker_pos] + content_shown_len = marker_pos + tool_hidden = True else: delta["content"] = c_diff if r_diff: @@ -2352,6 +2517,19 @@ async def _stream_openai( if c_diff: if tool_mode: content_buf += c_diff + if not tool_hidden: + search_from = content_shown_len + marker_pos = content_buf.find('{"tool_calls"', search_from) + if marker_pos < 0: + marker_pos = content_buf.find(" search_from: + delta2["content"] = content_buf[search_from:marker_pos] + content_shown_len = marker_pos + tool_hidden = True else: delta2["content"] = c_diff if r_diff: @@ -2448,6 +2626,28 @@ async def _stream_openai( ) if tool_mode: content_buf += cont_rec.content + if not tool_hidden: + search_from = content_shown_len + marker_pos = content_buf.find('{"tool_calls"', search_from) + if marker_pos < 0: + marker_pos = content_buf.find(" dict | None: +def _reasoning_text_from_output(output: Any) -> str: + parts: list[str] = [] + if not isinstance(output, list): + return "" + for item in output: + if not isinstance(item, dict): + continue + if item.get("type") not in ("reasoning", "reasoning_summary", "summary_text"): + continue + text = item.get("text") + if isinstance(text, str) and text: + parts.append(text) + continue + for key in ("summary", "content"): + content = item.get(key) + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict): + continue + part_text = part.get("text") + if isinstance(part_text, str) and part_text: + parts.append(part_text) + return "".join(parts) + + +def _usage_to_responses(usage: Any, output: Any = None) -> dict | None: if not isinstance(usage, dict): return None input_tokens = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0) output_tokens = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0) total_tokens = int(usage.get("total_tokens") or (input_tokens + output_tokens)) + reasoning_tokens = int(usage.get("reasoning_tokens") or 0) + if not reasoning_tokens: + reasoning_tokens = estimate_tokens(_reasoning_text_from_output(output)) return { "input_tokens": input_tokens, "input_tokens_details": {"cached_tokens": int(usage.get("cached_tokens") or 0)}, "output_tokens": output_tokens, - "output_tokens_details": {"reasoning_tokens": int(usage.get("reasoning_tokens") or 0)}, + "output_tokens_details": {"reasoning_tokens": reasoning_tokens}, "total_tokens": total_tokens, } @@ -286,7 +317,7 @@ def build_response_object( "tools": info.tools or [], "top_p": info.top_p, "truncation": info.truncation, - "usage": _usage_to_responses(usage), + "usage": _usage_to_responses(usage, output), "user": info.user, "metadata": info.metadata if isinstance(info.metadata, dict) else {}, } @@ -568,13 +599,29 @@ def tool_delta(self, tool_calls: Any) -> Iterator[str]: entry = { "id": f"fc_{uuid.uuid4().hex}", "call_id": call.get("id") or f"call_{uuid.uuid4().hex[:12]}", - "name": function.get("name") or "", + "name": "", "arguments": "", "output_index": self.output_index, "open": True, + "added": False, + "buffered_args": [], } self.output_index += 1 self.tool_items[index] = entry + name = function.get("name") + if isinstance(name, str) and name and not entry["name"]: + entry["name"] = name + arguments = function.get("arguments") + if isinstance(arguments, str) and arguments: + entry["arguments"] += arguments + if entry["added"]: + yield self.emit( + "response.function_call_arguments.delta", + {"item_id": entry["id"], "output_index": entry["output_index"], "delta": arguments}, + ) + else: + entry["buffered_args"].append(arguments) + if entry["name"] and not entry["added"]: item = { "id": entry["id"], "type": "function_call", @@ -584,15 +631,14 @@ def tool_delta(self, tool_calls: Any) -> Iterator[str]: "status": "in_progress", } yield self.emit("response.output_item.added", {"output_index": entry["output_index"], "item": item}) - if function.get("name") and not entry["name"]: - entry["name"] = function["name"] - arguments = function.get("arguments") - if isinstance(arguments, str) and arguments: - entry["arguments"] += arguments - yield self.emit( - "response.function_call_arguments.delta", - {"item_id": entry["id"], "output_index": entry["output_index"], "delta": arguments}, - ) + entry["added"] = True + if entry["buffered_args"]: + for buffered in entry["buffered_args"]: + yield self.emit( + "response.function_call_arguments.delta", + {"item_id": entry["id"], "output_index": entry["output_index"], "delta": buffered}, + ) + entry["buffered_args"] = [] def close_tools(self) -> Iterator[str]: for index in sorted(self.tool_items): @@ -600,6 +646,24 @@ def close_tools(self) -> Iterator[str]: if not entry.get("open"): continue entry["open"] = False + if not entry["added"]: + item = { + "id": entry["id"], + "type": "function_call", + "call_id": entry["call_id"], + "name": entry["name"], + "arguments": "", + "status": "in_progress", + } + yield self.emit("response.output_item.added", {"output_index": entry["output_index"], "item": item}) + entry["added"] = True + if entry["buffered_args"]: + for buffered in entry["buffered_args"]: + yield self.emit( + "response.function_call_arguments.delta", + {"item_id": entry["id"], "output_index": entry["output_index"], "delta": buffered}, + ) + entry["buffered_args"] = [] yield self.emit( "response.function_call_arguments.done", {"item_id": entry["id"], "output_index": entry["output_index"], "arguments": entry["arguments"]}, diff --git a/danyapi/deepseek/client.py b/danyapi/deepseek/client.py index 87c60e6..71bda91 100644 --- a/danyapi/deepseek/client.py +++ b/danyapi/deepseek/client.py @@ -63,7 +63,7 @@ def __init__( self.http = httpx.AsyncClient( base_url=BASE_URL, headers=headers, - timeout=httpx.Timeout(timeout), + timeout=httpx.Timeout(timeout, read=max(float(timeout) * 5, 300.0)), follow_redirects=True, ) @@ -105,9 +105,15 @@ async def check_auth(self) -> bool: "/api/v0/client/settings", params={"did": self.device_id, "scope": "main"}, ) - return resp.json().get("code") == 0 - except (httpx.HTTPError, ValueError): + except httpx.HTTPError: return False + if resp.status_code != 200: + return False + try: + payload = resp.json() + except ValueError: + return False + return payload.get("code") == 0 async def get_user(self) -> dict: resp = await self._post("/api/v0/users", None) diff --git a/danyapi/deepseek/pow_solver.c b/danyapi/deepseek/pow_solver.c index a3efdf5..f1ccaa3 100644 --- a/danyapi/deepseek/pow_solver.c +++ b/danyapi/deepseek/pow_solver.c @@ -14,6 +14,14 @@ #define ROUNDS 23 #define MAX_DIGITS 32 +#if defined(_MSC_VER) +#define POW_MEMORY_BARRIER() MemoryBarrier() +#else +#define POW_MEMORY_BARRIER() __sync_synchronize() +#endif + +static volatile int g_found = 0; + static const uint64_t RC[24] = { 0x0000000000000001ULL, 0x0000000000008082ULL, @@ -215,7 +223,13 @@ static int check_counter(const uint64_t base[25], size_t off0, } st[16] ^= (uint64_t)0x80 << 56; keccak_f(st); - return memcmp(st, target, 32) == 0; + for (int i = 0; i < 32; i++) + { + uint64_t lane = st[i >> 3]; + if ((uint8_t)((lane >> (8 * (i & 7))) & 0xffu) != target[i]) + return 0; + } + return 1; } typedef struct @@ -237,9 +251,13 @@ static void run_worker(WorkerArgs *a) int dlen = to_digits(a->start, digits); for (uint64_t c = a->start; c < a->end; c++) { + if (g_found) + return; if (check_counter(a->base, a->off0, digits, dlen, a->target)) { a->result = c; + POW_MEMORY_BARRIER(); + g_found = 1; return; } inc_digits(digits, &dlen); @@ -276,7 +294,7 @@ static int detect_threads(void) static int hex_to_bytes(const char *hex, uint8_t *out) { size_t n = strlen(hex); - if (n % 2) + if (n % 2 || n > 64) return -1; for (size_t i = 0; i < n; i += 2) { @@ -330,7 +348,12 @@ static long long find_json_ll(const char *json, const char *key) p = strchr(p + strlen(pat), ':'); if (!p) return -1; - return strtoll(p + 1, NULL, 10); + p++; + while (*p == ' ' || *p == '\t') + p++; + if (*p == '"') + p++; + return strtoll(p, NULL, 10); } int main(void) @@ -415,6 +438,7 @@ int main(void) } uint64_t chunk = (limit + (uint64_t)nthreads - 1) / (uint64_t)nthreads; + g_found = 0; for (int i = 0; i < nthreads; i++) { args[i].base = base; diff --git a/danyapi/deepseek/pow_solver.js b/danyapi/deepseek/pow_solver.js index c368f48..9251796 100644 --- a/danyapi/deepseek/pow_solver.js +++ b/danyapi/deepseek/pow_solver.js @@ -1,23 +1,156 @@ const fs = require("fs"); const path = require("path"); const wasmPath = path.join(__dirname, "sha3_wasm_bg.wasm"); - -if (!fs.existsSync(wasmPath)) { - process.stdout.write( - JSON.stringify({ error: "sha3_wasm_bg.wasm not found" }), - ); - process.exit(1); +const useWasm = fs.existsSync(wasmPath); +const RATE = 136; +const ROUNDS = 23; +const M64 = 0xffffffffffffffffn; +const RC = [ + 0x0000000000000001n, + 0x0000000000008082n, + 0x800000000000808an, + 0x8000000080008000n, + 0x000000000000808bn, + 0x0000000080000001n, + 0x8000000080008081n, + 0x8000000000008009n, + 0x000000000000008an, + 0x0000000000000088n, + 0x0000000080008009n, + 0x000000008000000an, + 0x000000008000808bn, + 0x800000000000008bn, + 0x8000000000008089n, + 0x8000000000008003n, + 0x8000000000008002n, + 0x8000000000000080n, + 0x000000000000800an, + 0x800000008000000an, + 0x8000000080008081n, + 0x8000000000008080n, + 0x0000000080000001n, + 0x8000000080008008n, +]; +const ROUND_RC = RC.slice(1, 1 + ROUNDS); +const ROT = [ + [0, 36, 3, 41, 18], + [1, 44, 10, 45, 2], + [62, 6, 43, 15, 61], + [28, 55, 25, 21, 56], + [27, 20, 39, 8, 14], +]; +function rol64(x, n) { + return ((x << BigInt(n)) | (x >> BigInt(64 - n))) & M64; +} +function keccakF(st) { + const c = new Array(5); + const d = new Array(5); + for (let r = 0; r < ROUNDS; r++) { + for (let x = 0; x < 5; x++) + c[x] = st[x] ^ st[x + 5] ^ st[x + 10] ^ st[x + 15] ^ st[x + 20]; + for (let x = 0; x < 5; x++) + d[x] = c[(x + 4) % 5] ^ rol64(c[(x + 1) % 5], 1); + for (let x = 0; x < 5; x++) + for (let y = 0; y < 5; y++) st[x + 5 * y] ^= d[x]; + const b = new Array(25); + for (let x = 0; x < 5; x++) + for (let y = 0; y < 5; y++) + b[y + 5 * ((2 * x + 3 * y) % 5)] = rol64(st[x + 5 * y], ROT[x][y]); + for (let x = 0; x < 5; x++) + for (let y = 0; y < 5; y++) + st[x + 5 * y] = + b[x + 5 * y] ^ (~b[(x + 1) % 5 + 5 * y] & b[(x + 2) % 5 + 5 * y]); + st[0] ^= ROUND_RC[r]; + } +} +function laneAt(bytes, off) { + let v = 0n; + for (let b = 0; b < 8; b++) v |= BigInt(bytes[off + b]) << BigInt(8 * b); + return v; +} +function hexToBytes(hex) { + if (typeof hex !== "string") return null; + if (hex.length % 2 || hex.length > 64) return null; + const out = new Uint8Array(hex.length / 2); + for (let i = 0; i < hex.length; i += 2) { + const hi = parseInt(hex[i], 16); + const lo = parseInt(hex[i + 1], 16); + if (isNaN(hi) || isNaN(lo)) return null; + out[i / 2] = (hi << 4) | lo; + } + return out; +} +function absorbPrefix(bytes) { + const st = new Array(25).fill(0n); + let pos = 0; + while (bytes.length - pos >= RATE) { + for (let i = 0; i < RATE; i += 8) st[i >> 3] ^= laneAt(bytes, pos + i); + keccakF(st); + pos += RATE; + } + const rem = bytes.length - pos; + for (let i = 0; i < rem; i++) + st[i >> 3] ^= BigInt(bytes[pos + i]) << BigInt(8 * (i & 7)); + return { st, off0: rem }; +} +function counterMatches(base, off0, digits, target) { + const st = base.slice(); + let off = off0; + for (let i = 0; i < digits.length; i++) { + st[off >> 3] ^= BigInt(digits.charCodeAt(i)) << BigInt(8 * (off & 7)); + off++; + if (off === RATE) { + keccakF(st); + off = 0; + } + } + st[off >> 3] ^= 6n << BigInt(8 * (off & 7)); + off++; + if (off === RATE) { + keccakF(st); + off = 0; + } + st[16] ^= 0x8000000000000000n; + keccakF(st); + for (let i = 0; i < target.length; i++) { + if (Number((st[i >> 3] >> BigInt(8 * (i & 7))) & 0xffn) !== target[i]) + return false; + } + return true; +} +function nextDigits(s) { + const a = s.split(""); + let i = a.length - 1; + while (i >= 0 && a[i] === "9") { + a[i] = "0"; + i--; + } + if (i < 0) a.unshift("1"); + else a[i] = String.fromCharCode(a[i].charCodeAt(0) + 1); + return a.join(""); +} +function solveJs(challenge, prefix, difficulty) { + const target = hexToBytes(challenge); + if (!target) return null; + const pre = Buffer.from(prefix, "utf8"); + const { st, off0 } = absorbPrefix(pre); + let limit = Number(difficulty); + if (!Number.isFinite(limit) || limit < 0) limit = 0; + limit = Math.min(Math.floor(limit), 2000000000); + let digits = "0"; + for (let c = 0; c < limit; c++) { + if (counterMatches(st, off0, digits, target)) return c; + digits = nextDigits(digits); + } + return null; } - -const wasmBuf = fs.readFileSync(wasmPath); - let instancePromise = null; function getInstance() { - if (!instancePromise) instancePromise = WebAssembly.instantiate(wasmBuf, {}); + if (!instancePromise) + instancePromise = WebAssembly.instantiate(fs.readFileSync(wasmPath), {}); return instancePromise; } - -function solve(challenge, prefix, difficulty) { +function solveWasm(challenge, prefix, difficulty) { return getInstance().then(({ instance }) => { const { memory, @@ -59,7 +192,6 @@ function solve(challenge, prefix, difficulty) { } }); } - let input = ""; process.stdin.setEncoding("utf8"); process.stdin.on("data", (chunk) => { @@ -75,7 +207,10 @@ process.stdin.on("end", () => { return; } const prefix = `${req.salt}_${req.expire_at}_`; - solve(req.challenge, prefix, req.difficulty) + const attempt = useWasm + ? solveWasm(req.challenge, prefix, req.difficulty) + : Promise.resolve(solveJs(req.challenge, prefix, req.difficulty)); + attempt .then((answer) => { if (answer === null) { process.stdout.write( diff --git a/danyapi/pow.py b/danyapi/pow.py index 05c9e3a..eb1d7f5 100644 --- a/danyapi/pow.py +++ b/danyapi/pow.py @@ -4,6 +4,7 @@ import base64 import json import logging +import math import struct import subprocess from pathlib import Path @@ -54,6 +55,29 @@ _PYTHON_SOLVE_LIMIT = 2_000_000 +def _parse_number(value): + if isinstance(value, bool): + return None + if isinstance(value, int): + return value + if isinstance(value, float): + return value if math.isfinite(value) else None + if isinstance(value, str): + text = value.strip() + if not text: + return None + try: + return int(text) + except ValueError: + pass + try: + number = float(text) + except ValueError: + return None + return number if math.isfinite(number) else None + return None + + def _rol(x: int, n: int) -> int: return ((x << n) | (x >> (64 - n))) & _MASK64 @@ -187,11 +211,11 @@ async def _build(self, fetch) -> dict: missing = [k for k in ("challenge", "salt", "algorithm", "signature", "target_path") if not challenge.get(k)] if missing: raise RuntimeError(f"pow challenge missing fields: {', '.join(missing)}") - expire_at = challenge.get("expire_at") - difficulty = challenge.get("difficulty") - if isinstance(expire_at, bool) or not isinstance(expire_at, (int, float)): + expire_at = _parse_number(challenge.get("expire_at")) + difficulty = _parse_number(challenge.get("difficulty")) + if expire_at is None or expire_at < 0: raise RuntimeError("pow challenge has invalid expire_at") - if isinstance(difficulty, bool) or not isinstance(difficulty, (int, float)) or difficulty <= 0: + if difficulty is None or difficulty <= 0: raise RuntimeError("pow challenge has invalid difficulty") answer = await solve_challenge( challenge["challenge"], diff --git a/danyapi/qwen/accounts.py b/danyapi/qwen/accounts.py index eb6a435..57e5618 100644 --- a/danyapi/qwen/accounts.py +++ b/danyapi/qwen/accounts.py @@ -2,9 +2,10 @@ import asyncio import logging +import time from typing import Any -from ..sessions import SessionRegistry +from ..sessions import SessionRegistry, _as_int from ..store import JsonStore from .client import QwenClient, QwenSession @@ -42,8 +43,8 @@ def _deserialize(self, record: Any) -> QwenSession: title=record.get("title") or "", last_response_id=record.get("last_response_id"), model=record.get("model"), - accumulated_input_tokens=int(input_tokens) if isinstance(input_tokens, (int, float)) else 0, - accumulated_output_tokens=int(output_tokens) if isinstance(output_tokens, (int, float)) else 0, + accumulated_input_tokens=_as_int(input_tokens), + accumulated_output_tokens=_as_int(output_tokens), ) async def _create(self, **kwargs) -> QwenSession: @@ -66,7 +67,7 @@ async def obtain(self, session_id: str | None = None, model: str | None = None, class QwenAccount: - __slots__ = ("broken", "client", "index", "sem", "sessions", "stable_id") + __slots__ = ("broken", "broken_at", "client", "index", "sem", "sessions", "stable_id") def __init__( self, @@ -83,10 +84,12 @@ def __init__( self.sessions = QwenSessionRegistry(client, session_cache_size, ttl, store=store, key_prefix=f"{index}:") self.stable_id = stable_id self.broken = False + self.broken_at: float | None = None def mark_broken(self) -> None: if not self.broken: self.broken = True + self.broken_at = time.monotonic() log.warning("qwen account #%d marked broken (invalid/expired token)", self.index) @property diff --git a/danyapi/qwen/api.py b/danyapi/qwen/api.py index ba80793..40622ea 100644 --- a/danyapi/qwen/api.py +++ b/danyapi/qwen/api.py @@ -71,6 +71,43 @@ def _is_retryable_http(exc: HTTPException) -> bool: "tokenlimit", ) +_TOOL_MARKERS = ('{"tool_calls"', "") + + +def _tool_marker_pos(text: str) -> int: + found = -1 + for marker in _TOOL_MARKERS: + pos = text.find(marker) + if pos != -1 and (found == -1 or pos < found): + found = pos + return found + + +def _append_image_markdown(prompt: str, messages: list[Any] | None) -> str: + if not messages: + return prompt + appended: list[str] = [] + for message in messages: + content = getattr(message, "content", None) + if not isinstance(content, list): + continue + for item in content: + if not isinstance(item, dict) or item.get("type") != "image_url": + continue + image_url = item.get("image_url") + if isinstance(image_url, str): + uri = image_url + elif isinstance(image_url, dict) and isinstance(image_url.get("url"), str): + uri = image_url["url"] + else: + continue + if uri.startswith("http") or uri.startswith("data:"): + appended.append(f"![image]({uri})") + if not appended: + return prompt + extra = "\n".join(appended) + return f"{prompt}\n\n{extra}" if prompt else extra + class ContextLimitError(Exception): pass @@ -286,17 +323,19 @@ async def _human_delay() -> None: def _accumulate_usage(session, rec: QwenStreamReconstructor) -> dict: current = rec.usage_tokens - prompt_tokens = current["prompt_tokens"] - completion_tokens = current["completion_tokens"] - total_tokens = current["total_tokens"] or prompt_tokens + completion_tokens + current_input = current["prompt_tokens"] + current_output = current["completion_tokens"] + current_total = current["total_tokens"] or current_input + current_output prev_input = int(getattr(session, "accumulated_input_tokens", 0) or 0) prev_output = int(getattr(session, "accumulated_output_tokens", 0) or 0) - session.accumulated_input_tokens = prev_input + prompt_tokens - session.accumulated_output_tokens = prev_output + completion_tokens + prompt_tokens = max(0, current_input - prev_input) + completion_tokens = max(0, current_output - prev_output) + session.accumulated_input_tokens = current_input + session.accumulated_output_tokens = current_output return { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, - "total_tokens": total_tokens, + "total_tokens": current_total or prompt_tokens + completion_tokens, } @@ -330,7 +369,7 @@ async def _collect_response( resp = await _send_completion( account.client, session, - prompt, + _append_image_markdown(prompt, messages), model_id, thinking, search, @@ -562,13 +601,16 @@ async def stream_openai( rec: QwenStreamReconstructor | None = None content_buf = "" + content_shown_len = 0 + tool_marker_pos = -1 + role_sent = False stop_response_id: str | None = None had_cached_session = bool(existing_sid) and account.sessions.get(existing_sid) is not None stale_rebuilt = False attempt = 0 while True: try: - resp = await _send_completion(account.client, session, prompt, model_id, thinking, search) + resp = await _send_completion(account.client, session, _append_image_markdown(prompt, messages), model_id, thinking, search) except ContextLimitError: _drop_session(pool, account, session_key) for line in _stream_context_limit_lines(chunk_id, created, model, session_key): @@ -621,74 +663,54 @@ async def stream_openai( rec = QwenStreamReconstructor() incremental = IncrementalSSE() got_content = False - pending: list[str] = [] role_sent = False + content_shown_len = 0 + tool_marker_pos = -1 stopped = False try: async for chunk in resp.aiter_bytes(): for event in incremental.feed(chunk): rec.handle(event) c_diff, r_diff = rec.take_diffs() - if c_diff or r_diff: - got_content = True + if not (c_diff or r_diff): + continue + got_content = True if not role_sent: role_sent = True - pending.append( - _sse( - { - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [ - { - "index": 0, - "delta": {"role": "assistant"}, - "finish_reason": None, - } - ], - } - ) + yield _sse( + { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [ + { + "index": 0, + "delta": {"role": "assistant"}, + "finish_reason": None, + } + ], + } ) delta: dict = {} if c_diff: if tool_mode: content_buf += c_diff + if tool_marker_pos == -1: + tool_marker_pos = _tool_marker_pos(content_buf) + if tool_marker_pos == -1: + shown = content_buf[content_shown_len:] + else: + shown = content_buf[content_shown_len:tool_marker_pos] + if shown: + delta["content"] = shown + content_shown_len += len(shown) else: delta["content"] = c_diff if r_diff: delta["reasoning_content"] = r_diff if delta: - pending.append( - _sse( - { - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [ - { - "index": 0, - "delta": delta, - "finish_reason": None, - } - ], - } - ) - ) - if got_content: - for line in pending: - yield line - pending.clear() - for event in incremental.finish(): - rec.handle(event) - c_diff, r_diff = rec.take_diffs() - if c_diff or r_diff: - got_content = True - if not role_sent: - role_sent = True - pending.append( - _sse( + yield _sse( { "id": chunk_id, "object": "chat.completion.chunk", @@ -697,43 +719,68 @@ async def stream_openai( "choices": [ { "index": 0, - "delta": {"role": "assistant"}, + "delta": delta, "finish_reason": None, } ], } ) + for event in incremental.finish(): + rec.handle(event) + c_diff, r_diff = rec.take_diffs() + if not (c_diff or r_diff): + continue + got_content = True + if not role_sent: + role_sent = True + yield _sse( + { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [ + { + "index": 0, + "delta": {"role": "assistant"}, + "finish_reason": None, + } + ], + } ) delta2: dict = {} if c_diff: if tool_mode: content_buf += c_diff + if tool_marker_pos == -1: + tool_marker_pos = _tool_marker_pos(content_buf) + if tool_marker_pos == -1: + shown = content_buf[content_shown_len:] + else: + shown = content_buf[content_shown_len:tool_marker_pos] + if shown: + delta2["content"] = shown + content_shown_len += len(shown) else: delta2["content"] = c_diff if r_diff: delta2["reasoning_content"] = r_diff if delta2: - pending.append( - _sse( - { - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [ - { - "index": 0, - "delta": delta2, - "finish_reason": None, - } - ], - } - ) + yield _sse( + { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [ + { + "index": 0, + "delta": delta2, + "finish_reason": None, + } + ], + } ) - if got_content: - for line in pending: - yield line - pending.clear() except BaseException: stopped = True if rec.response_id: @@ -802,9 +849,9 @@ async def stream_openai( if tool_mode: parsed = toolemu.parse_tool_calls(content_buf, tool_schemas) if parsed is not None: - tool_calls, tool_text = parsed + tool_calls, _ = parsed if tool_calls: - for delta in toolemu.tool_call_deltas(tool_calls, tool_text): + for delta in toolemu.tool_call_deltas(tool_calls): yield _sse( { "id": chunk_id, @@ -816,24 +863,27 @@ async def stream_openai( ) finish = "tool_calls" else: - yield _sse( - { - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [ - { - "index": 0, - "delta": {"content": content_buf}, - "finish_reason": None, - } - ], - } - ) + remainder = content_buf[content_shown_len:] + if remainder: + yield _sse( + { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [ + { + "index": 0, + "delta": {"content": remainder}, + "finish_reason": None, + } + ], + } + ) finish = "stop" else: - if content_buf: + remainder = content_buf[content_shown_len:] + if remainder: yield _sse( { "id": chunk_id, @@ -843,7 +893,7 @@ async def stream_openai( "choices": [ { "index": 0, - "delta": {"content": content_buf}, + "delta": {"content": remainder}, "finish_reason": None, } ], @@ -853,6 +903,24 @@ async def stream_openai( else: finish = "stop" + if not role_sent: + role_sent = True + yield _sse( + { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [ + { + "index": 0, + "delta": {"role": "assistant"}, + "finish_reason": None, + } + ], + } + ) + yield _sse( { "id": chunk_id, diff --git a/danyapi/qwen/stream.py b/danyapi/qwen/stream.py index bd6b0e0..a2d6d60 100644 --- a/danyapi/qwen/stream.py +++ b/danyapi/qwen/stream.py @@ -90,8 +90,8 @@ def handle(self, event: SSEEvent) -> None: if phase in IMAGE_PHASES: text = _delta_text(delta, "content") if text: - self._collect_image_urls(text) self.content += text + self._collect_image_urls(self.content) image_field = delta.get("image_url") or delta.get("image") if isinstance(image_field, str) and image_field.startswith("http"): if image_field not in self._seen_image_urls: @@ -126,7 +126,7 @@ def handle(self, event: SSEEvent) -> None: text = _delta_text(delta, "content") if text: self.content += text - self._collect_image_urls(text) + self._collect_image_urls(self.content) elif phase in THINK_PHASES: text = _delta_text(delta, "content") if text: @@ -156,7 +156,22 @@ def has_content(self) -> bool: @property def usage_tokens(self) -> dict: def _int(value: Any) -> int: - return int(value) if isinstance(value, (int, float)) else 0 + if isinstance(value, bool): + return 0 + if isinstance(value, int): + return value + if isinstance(value, float): + return int(value) + if not isinstance(value, str): + return 0 + try: + return int(value.strip()) + except (TypeError, ValueError): + pass + try: + return int(float(value.strip())) + except (TypeError, ValueError): + return 0 return { "prompt_tokens": _int(self.usage.get("input_tokens")), diff --git a/danyapi/sessions.py b/danyapi/sessions.py index 57dbb7f..b78dce8 100644 --- a/danyapi/sessions.py +++ b/danyapi/sessions.py @@ -9,6 +9,22 @@ from .store import JsonStore +def _as_int(value: Any, default: int = 0) -> int: + if isinstance(value, bool): + return default + if isinstance(value, (int, float)): + return int(value) + if isinstance(value, str): + try: + return int(value) + except ValueError: + try: + return int(float(value)) + except ValueError: + return default + return default + + class SessionRegistry: def __init__( self, @@ -25,7 +41,8 @@ def __init__( self._ttl = max(0.0, ttl) self._store = store self._key_prefix = key_prefix - self._creation_lock = asyncio.Lock() + self._session_locks: dict[str, asyncio.Lock] = {} + self._session_locks_guard = asyncio.Lock() self._restore() def _now(self) -> float: @@ -51,7 +68,7 @@ def _deserialize(self, record: Any) -> Any: from .deepseek.client import DeepSeekSession accumulated = record.get("accumulated_tokens") - accumulated_tokens = int(accumulated) if isinstance(accumulated, (int, float)) else 0 + accumulated_tokens = _as_int(accumulated) return DeepSeekSession( id=record["id"], title=record.get("title") or "", @@ -122,37 +139,51 @@ def get(self, session_id: str | None) -> Any | None: self._sessions[session_id] = (session, now) return session + async def _session_lock(self, session_key: str) -> asyncio.Lock: + async with self._session_locks_guard: + lock = self._session_locks.get(session_key) + if lock is None: + lock = asyncio.Lock() + self._session_locks[session_key] = lock + return lock + async def obtain(self, session_id: str | None, **kwargs: Any) -> tuple[Any, str]: - async with self._creation_lock: - if session_id: - existing = self.get(session_id) - if existing is not None and self._reuse(existing, session_id, **kwargs): - return existing, session_id - session = await self._create(**kwargs) - new_id = session.id - bind_key = session_id or new_id - now = self._now() - with self._lock: - self._sessions[new_id] = (session, now) - if bind_key != new_id: - self._sessions[bind_key] = (session, now) - self._sessions.move_to_end(bind_key) - if bind_key != new_id: - self._sessions.move_to_end(new_id) - while len(self._sessions) > self._maxsize: - protect = {new_id, bind_key} - evictable = [k for k in self._sessions if k not in protect] - if not evictable: - break - oldest = min(evictable, key=lambda k: self._sessions[k][1]) - self._sessions.pop(oldest, None) - if self._store is not None: - self._store.discard(self._session_key(oldest)) + if session_id: + lock = await self._session_lock(session_id) + async with lock: + return await self._obtain(session_id, **kwargs) + return await self._obtain(session_id, **kwargs) + + async def _obtain(self, session_id: str | None, **kwargs: Any) -> tuple[Any, str]: + if session_id: + existing = self.get(session_id) + if existing is not None and self._reuse(existing, session_id, **kwargs): + return existing, session_id + session = await self._create(**kwargs) + new_id = session.id + bind_key = session_id or new_id + now = self._now() + with self._lock: + self._sessions[new_id] = (session, now) + if bind_key != new_id: + self._sessions[bind_key] = (session, now) + self._sessions.move_to_end(bind_key) + if bind_key != new_id: + self._sessions.move_to_end(new_id) + while len(self._sessions) > self._maxsize: + protect = {new_id, bind_key} + evictable = [k for k in self._sessions if k not in protect] + if not evictable: + break + oldest = min(evictable, key=lambda k: self._sessions[k][1]) + self._sessions.pop(oldest, None) if self._store is not None: - self._store.set(self._session_key(new_id), self._serialize(session)) - if bind_key != new_id: - self._store.set(self._session_key(bind_key), self._serialize(session)) - return session, bind_key + self._store.discard(self._session_key(oldest)) + if self._store is not None: + self._store.set(self._session_key(new_id), self._serialize(session)) + if bind_key != new_id: + self._store.set(self._session_key(bind_key), self._serialize(session)) + return session, bind_key def touch_last_message(self, session_id: str, message_id: str | None) -> None: session = self.get(session_id) @@ -168,10 +199,12 @@ def forget(self, session_id: str) -> None: self._sessions.pop(session_id, None) if self._store is not None: self._store.discard(self._session_key(session_id)) + self._session_locks.pop(session_id, None) def close_all(self) -> None: with self._lock: self._sessions.clear() + self._session_locks.clear() if self._store is not None: prefix = self._key_prefix for key, _ in self._store.items(): diff --git a/danyapi/sseutil.py b/danyapi/sseutil.py index 44be453..59e8be1 100644 --- a/danyapi/sseutil.py +++ b/danyapi/sseutil.py @@ -166,7 +166,7 @@ def _apply_delta(message: dict, op: str, path: str, value: Any) -> None: idx = int(key) except ValueError: return - if idx >= len(node) or idx < 0: + if idx < -len(node) or idx >= len(node): return cur = node[idx] if isinstance(cur, str) and isinstance(value, str): diff --git a/danyapi/store.py b/danyapi/store.py index 0cf5cdf..8895c39 100644 --- a/danyapi/store.py +++ b/danyapi/store.py @@ -14,6 +14,8 @@ DEFAULT_CACHE_SUBDIR = "danyapi" +_MAX_AFFINITY = 8192 + def cache_root() -> Path: override = settings.cache_dir diff --git a/danyapi/tools.py b/danyapi/tools.py index 75a26f6..c4acbd1 100644 --- a/danyapi/tools.py +++ b/danyapi/tools.py @@ -232,7 +232,7 @@ _XML_WRAPPER_CLOSE_RE = re.compile(r"", re.IGNORECASE) _XML_TOOL_NAMES = r"invoke|toolinvoke|tool_invoke|use_tool|tool_use|call|function|tool" _XML_TOOL_ELEMENT_RE = re.compile( - rf"<(?:{_XML_TOOL_NAMES})\b([^>]*)>(.*?)", + rf"<((?:{_XML_TOOL_NAMES}))\b([^>]*)>(.*?)", re.DOTALL | re.IGNORECASE, ) _XML_TOOL_SELFCLOSE_RE = re.compile( @@ -414,14 +414,20 @@ def _content_text(content: Any, *, with_images: bool = False, separator: str = " return "" +def _msg_field(msg: Any, key: str, default: Any = None) -> Any: + if isinstance(msg, dict): + return msg.get(key, default) + return getattr(msg, key, default) + + def context_sequence(messages: list[Any], user: str | None = None) -> tuple[str, ...]: sequence: list[str] = [] scope = f"\0{user or ''}" for msg in messages: - role = getattr(msg, "role", "user") + role = _msg_field(msg, "role", "user") if role not in ("system", "user"): continue - content = _content_text(getattr(msg, "content", ""), with_images=True, separator="\n") + content = _content_text(_msg_field(msg, "content", ""), with_images=True, separator="\n") if not content.strip(): continue digest = hashlib.sha256(f"{role}\0{content}{scope}".encode()).hexdigest() @@ -445,30 +451,30 @@ def _render_tool_call_mention(call: Any) -> str: def render_message(msg: Any) -> str: - role = getattr(msg, "role", "user") - text = _strip_dsml(_content_text(getattr(msg, "content", ""))) + role = _msg_field(msg, "role", "user") + text = _strip_dsml(_content_text(_msg_field(msg, "content", ""))) if role in ("user", "system"): return text if role == "assistant": parts = [] if text: parts.append(text) - for call in getattr(msg, "tool_calls", None) or []: + for call in _msg_field(msg, "tool_calls", None) or []: mention = _render_tool_call_mention(call) if mention: parts.append(mention) - content = getattr(msg, "content", None) + content = _msg_field(msg, "content", None) if isinstance(content, list): for item in content: if isinstance(item, dict) and item.get("type") == "tool_call": parts.append(_render_tool_call_mention(item)) return "; ".join(parts) if role == "tool": - tool_call_id = getattr(msg, "tool_call_id", None) or "" + tool_call_id = _msg_field(msg, "tool_call_id", None) or "" prefix = f"Tool result ({tool_call_id})" if tool_call_id else "Tool result" return f"{prefix}: {text}" if role == "function": - name = getattr(msg, "name", None) or "" + name = _msg_field(msg, "name", None) or "" return f"Function {name} returned: {text}" return text @@ -478,7 +484,7 @@ def _render_history(messages: list[Any]) -> str: for msg in messages: text = render_message(msg) if text: - role = getattr(msg, "role", "user") + role = _msg_field(msg, "role", "user") parts.append(f"{role.capitalize()}: {text}") return "\n".join(parts) @@ -486,7 +492,7 @@ def _render_history(messages: list[Any]) -> str: def _render_tool_tail(messages: list[Any]) -> str: parts = [] for msg in messages: - role = getattr(msg, "role", None) + role = _msg_field(msg, "role", None) if role in ("tool", "function"): parts.append(render_message(msg)) parts.append(TOOL_TAIL_REMINDER) @@ -497,37 +503,41 @@ def extract_last_user(messages: list[Any]) -> str: if not messages: raise ValueError("messages is required") for msg in reversed(messages): - if getattr(msg, "role", None) in ("user", "system"): - content = getattr(msg, "content", "") - if isinstance(content, str): - return _strip_dsml(content) - if isinstance(content, list): - parts = [] - for item in content: - if isinstance(item, str): - parts.append(item) - elif isinstance(item, dict): - if item.get("type") == "text" and isinstance(item.get("text"), str): - parts.append(item["text"]) - elif item.get("type") == "image_url": - continue - text = _strip_dsml("".join(parts)).strip() - if text: - return text - continue - raise ValueError("unsupported message content") + if _msg_field(msg, "role", None) != "user": + continue + content = _msg_field(msg, "content", None) + if content is None: + continue + if isinstance(content, str): + return _strip_dsml(content) + if isinstance(content, list): + parts = [] + for item in content: + if isinstance(item, str): + parts.append(item) + elif isinstance(item, dict): + if item.get("type") == "text" and isinstance(item.get("text"), str): + parts.append(item["text"]) + elif item.get("type") == "image_url": + continue + text = _strip_dsml("".join(parts)).strip() + if text: + return text + continue + raise ValueError("unsupported message content") raise ValueError("no user message found") def is_tool_round(messages: list[Any]) -> bool: for msg in messages: - role = getattr(msg, "role", None) + role = _msg_field(msg, "role", None) if role in ("tool", "function"): return True - if role == "assistant" and getattr(msg, "tool_calls", None): + if role == "assistant" and _msg_field(msg, "tool_calls", None): return True - if role == "assistant" and isinstance(getattr(msg, "content", None), list): - for item in msg.content: + content = _msg_field(msg, "content", None) + if role == "assistant" and isinstance(content, list): + for item in content: if isinstance(item, dict) and item.get("type") == "tool_call": return True return False @@ -536,12 +546,12 @@ def is_tool_round(messages: list[Any]) -> bool: def _has_history(messages: list[Any]) -> bool: user_count = 0 for msg in messages: - role = getattr(msg, "role", None) + role = _msg_field(msg, "role", None) if role == "user": user_count += 1 continue - text = _content_text(getattr(msg, "content", "")) - if role == "assistant" and (text or getattr(msg, "tool_calls", None)): + text = _content_text(_msg_field(msg, "content", "")) + if role == "assistant" and (text or _msg_field(msg, "tool_calls", None)): return True if role in ("tool", "function") and text: return True @@ -551,7 +561,7 @@ def _has_history(messages: list[Any]) -> bool: def _tail_after_last_user(messages: list[Any]) -> list[Any]: index = -1 for i, msg in enumerate(messages): - if getattr(msg, "role", None) in ("user", "system"): + if _msg_field(msg, "role", None) in ("user", "system"): index = i if index < 0: return list(messages) @@ -561,8 +571,8 @@ def _tail_after_last_user(messages: list[Any]) -> list[Any]: def extract_system(messages: list[Any]) -> str: parts = [] for msg in messages: - if getattr(msg, "role", None) == "system": - text = _strip_dsml(_content_text(getattr(msg, "content", ""))).strip() + if _msg_field(msg, "role", None) == "system": + text = _strip_dsml(_content_text(_msg_field(msg, "content", ""))).strip() if text: parts.append(text) return "\n".join(parts) @@ -609,6 +619,9 @@ def build_prompt( blocks = [] if json_block: blocks.append(json_block) + choice = _choice_name(tool_choice) + if schema and choice is not None and choice not in ("auto", "none"): + blocks.append(schema) blocks.append(base) return "\n\n".join(blocks), tools_present @@ -736,6 +749,32 @@ def _is_bare_literal(token: str) -> bool: return _NUMBER_RE.fullmatch(token) is not None +_URL_SCHEME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9+.\-]*://") + + +def _url_like_after(text: str, i: int) -> bool: + fragment = text[i : i + 32] + if not fragment: + return False + if fragment.startswith("://"): + return True + if _URL_SCHEME_RE.match(fragment): + return True + if fragment[0] == ":": + rest = fragment[1:] + for k, ch in enumerate(rest): + if k >= 16: + break + if ch in "{}[],\"'\\ \t\r\n": + if ch in "[],": + return False + break + else: + return True + return False + return False + + def _normalize_bare_json(text: str) -> str | None: out: list[str] = [] i = 0 @@ -802,7 +841,7 @@ def _normalize_bare_json(text: str) -> str | None: i = j changed = True continue - if prev == ":" and not _is_bare_literal(token): + if prev == ":" and not _is_bare_literal(token) and not _url_like_after(text, i): out.append('"') out.append(token) out.append('"') @@ -1030,8 +1069,12 @@ def _extract_calls(obj: dict) -> list[ToolCall] | None: return None +_XML_ENTITY_RE = re.compile(r"&(amp|lt|gt|quot|apos);") +_XML_ENTITY_MAP = {"lt": "<", "gt": ">", "quot": '"', "apos": "'", "amp": "&"} + + def _unescape_xml(text: str) -> str: - return text.replace("<", "<").replace(">", ">").replace(""", '"').replace("'", "'").replace("&", "&") + return _XML_ENTITY_RE.sub(lambda m: _XML_ENTITY_MAP.get(m.group(1), m.group(0)), text) def _coerce_scalar(value: str, json_type: Any) -> Any: @@ -1219,7 +1262,7 @@ def _xml_invoke_arguments(body: str, param_types: dict[str, Any] | None = None, if params: if len(params) == 1: for key in _ARGS_ALIASES: - if key in params and isinstance(params[key], dict): + if key in params and isinstance(params[key], dict) and (param_types is None or key not in param_types): return params[key] return params if not allow_content: @@ -1275,8 +1318,8 @@ def blank(start: int, end: int) -> None: for match in _XML_TOOL_ELEMENT_RE.finditer(text): start, end = match.span() - attrs_text = match.group(1) - body = match.group(2) + attrs_text = match.group(2) + body = match.group(3) name_match = _XML_NAME_ATTR_RE.search(attrs_text) tool_name = name_match.group(2) if name_match else None if not tool_name: @@ -1501,7 +1544,6 @@ def _iter_json_objects(text: str) -> Iterator[tuple[dict, int, int]]: in_string = False escaped = False end = start - parsed = False closed = False while end < length: ch = text[end] @@ -1520,18 +1562,19 @@ def _iter_json_objects(text: str) -> Iterator[tuple[dict, int, int]]: depth -= 1 if depth == 0: closed = True - candidate = text[start : end + 1] - try: - obj = _loads_lenient(candidate) - yield obj, start, end - parsed = True - except ValueError: - pass break end += 1 scanned += end - start + 1 attempts += 1 - i = (end + 1) if (parsed or not closed) else (start + 1) + if not closed: + return + candidate = text[start : end + 1] + try: + obj = _loads_lenient(candidate) + yield obj, start, end + except ValueError: + pass + i = end + 1 def _split_top_level(text: str, delimiter: str = ",") -> list[str]: diff --git a/danyapi/usage.py b/danyapi/usage.py index 478b0eb..9f9b199 100644 --- a/danyapi/usage.py +++ b/danyapi/usage.py @@ -79,6 +79,7 @@ def _restore(self) -> None: if field in row and isinstance(value, (int, float)): row[field] = int(value) restored[name] = row + self._evict_overflow(restored) setattr(self, attr, restored) def _serialize(self) -> dict[str, Any]: @@ -91,11 +92,18 @@ def _serialize(self) -> dict[str, Any]: @staticmethod def _add(bucket: dict[str, dict[str, int]], key: str, prompt_tokens: int, completion_tokens: int, total_tokens: int) -> None: - entry = bucket.setdefault(key, {"requests": 0, "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}) + entry = bucket.pop(key, None) + if entry is None: + entry = {"requests": 0, "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} entry["requests"] += 1 entry["prompt_tokens"] += prompt_tokens entry["completion_tokens"] += completion_tokens entry["total_tokens"] += total_tokens + bucket[key] = entry + + def _evict_overflow(self, bucket: dict[str, dict[str, int]]) -> None: + while len(bucket) > self._max_records: + bucket.pop(next(iter(bucket))) def record( self, @@ -118,9 +126,12 @@ def record( self._totals["completion_tokens"] += completion_tokens self._totals["total_tokens"] += total_tokens self._add(self._by_model, model or "unknown", prompt_tokens, completion_tokens, total_tokens) + self._evict_overflow(self._by_model) self._add(self._by_provider, provider or "unknown", prompt_tokens, completion_tokens, total_tokens) + self._evict_overflow(self._by_provider) if user: self._add(self._by_user, user, prompt_tokens, completion_tokens, total_tokens) + self._evict_overflow(self._by_user) self._recent.append( { "ts": time.time(), diff --git a/tests/test_api_helpers.py b/tests/test_api_helpers.py index d594788..8cdc864 100644 --- a/tests/test_api_helpers.py +++ b/tests/test_api_helpers.py @@ -152,9 +152,9 @@ def test_finish_reason(): assert openai_mod._finish_reason("FINISHED") == "stop" assert openai_mod._finish_reason("CONTEXT_LENGTH_EXCEEDED") == "length" assert openai_mod._finish_reason("CONTENT_FILTER") == "content_filter" - assert openai_mod._finish_reason("INCOMPLETE") == "stop" - assert openai_mod._finish_reason("WIP") == "stop" - assert openai_mod._finish_reason("TIMEOUT") == "stop" + assert openai_mod._finish_reason("INCOMPLETE") == "length" + assert openai_mod._finish_reason("WIP") == "length" + assert openai_mod._finish_reason("TIMEOUT") == "length" assert openai_mod._finish_reason("WEIRD") == "stop" assert openai_mod._finish_reason(None) == "stop" assert openai_mod._finish_reason(42) == "stop" @@ -2229,3 +2229,56 @@ async def test_image_generations_requires_qwen_pool(): with pytest.raises(openai_mod.HTTPException) as excinfo: await openai_mod._image_generations(req) assert excinfo.value.status_code == 503 + + +def test_stream_error_sse_shape(): + first, done = openai_mod._stream_error_sse("c1", 123, "m1", "boom", session_key="s1", error_finish="length") + assert done == "data: [DONE]\n\n" + payload = json.loads(first[6:]) + assert payload["id"] == "c1" + assert payload["session_id"] == "s1" + assert payload["error"]["message"] == "boom" + assert payload["error"]["finish_reason"] == "length" + assert payload["choices"][0]["delta"] == {} + assert payload["choices"][0]["finish_reason"] == "length" + + +def test_collect_attachments_image_total_cap_413(): + big = b64.b64encode(b"x" * (openai_mod.MAX_ATTACHMENT_TOTAL_SIZE + 1)).decode() + req = SimpleNamespace( + messages=[ + openai_mod.ChatMessage( + role="user", + content=[{"type": "image_url", "image_url": f"data:image/png;base64,{big}"}], + ) + ], + files=[], + ) + with pytest.raises(openai_mod.HTTPException) as excinfo: + openai_mod._collect_attachments(req) + assert excinfo.value.status_code == 413 + + +async def test_add_tokens_reactivates_broken_account(monkeypatch, tmp_path): + from danyapi import store as store_mod + from danyapi.accounts import AccountPool, DeepSeekAccount + + token = "tok1" + monkeypatch.setattr(store_mod.settings, "cache_dir", str(tmp_path)) + env_file = tmp_path / ".env" + env_file.write_text(f"DEEPSEEK_TOKENS={token}\nQWEN_TOKENS=\n", encoding="utf-8") + monkeypatch.setattr(openai_mod, "_env_path", lambda: env_file) + client = MagicMock() + client.check_auth = AsyncMock(return_value=True) + client.aclose = AsyncMock() + acct = DeepSeekAccount(0, client, stable_id=openai_mod._token_stable_id(token)) + acct.mark_broken() + app.state.pool = AccountPool([acct]) + result = await openai_mod.add_tokens({"deepseek_tokens": [token]}) + assert result["success"] is True + assert result["message"] == "Tokens reactivated." + assert result["reactivated"]["deepseek"] == 1 + assert result["added"]["deepseek"] == 0 + assert acct.broken is False + assert acct.broken_at is None + assert env_file.read_text(encoding="utf-8").count(token) == 1 diff --git a/tests/test_deepseek_client.py b/tests/test_deepseek_client.py index 82aaf3a..64054b1 100644 --- a/tests/test_deepseek_client.py +++ b/tests/test_deepseek_client.py @@ -132,6 +132,14 @@ async def test_check_auth_bad_code(): assert not await client.check_auth() +async def test_check_auth_non_200(): + client = make_client() + resp = make_resp({"code": 0}, status=503) + client.http.get = AsyncMock(return_value=resp) + + assert not await client.check_auth() + + async def test_check_auth_exception(): client = make_client() client.http.get = AsyncMock(side_effect=httpx.ConnectError("boom")) @@ -139,6 +147,14 @@ async def test_check_auth_exception(): assert not await client.check_auth() +def test_client_timeout_read_extended(): + client = DeepSeekClient(token="tok", timeout=10.0) + assert client.http.timeout.read == 300.0 + assert client.http.timeout.connect == 10.0 + assert client.http.timeout.write == 10.0 + assert client.http.timeout.pool == 10.0 + + async def test_get_user(): client = make_client() resp = make_resp({"code": 0, "data": {"biz_data": {"user": {"id": 7}}}}) diff --git a/tests/test_qwen_api.py b/tests/test_qwen_api.py index b14a94c..f043a4a 100644 --- a/tests/test_qwen_api.py +++ b/tests/test_qwen_api.py @@ -455,15 +455,15 @@ def test_accumulate_usage(): usage = qwen_api._accumulate_usage(session, rec) assert usage == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} usage2 = qwen_api._accumulate_usage(session, rec) - assert usage2 == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} - assert session.accumulated_input_tokens == 20 - assert session.accumulated_output_tokens == 10 + assert usage2 == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 15} + assert session.accumulated_input_tokens == 10 + assert session.accumulated_output_tokens == 5 rec2 = QwenStreamReconstructor() - rec2.usage = {"input_tokens": 3, "output_tokens": 4} + rec2.usage = {"input_tokens": 13, "output_tokens": 8} usage3 = qwen_api._accumulate_usage(session, rec2) - assert usage3 == {"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7} - assert session.accumulated_input_tokens == 23 - assert session.accumulated_output_tokens == 14 + assert usage3 == {"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 21} + assert session.accumulated_input_tokens == 13 + assert session.accumulated_output_tokens == 8 def test_drop_session(): @@ -636,7 +636,65 @@ async def test_stream_tool_mode_falls_back_to_content(): lines = await _collect(gen) joined = "".join(lines) assert '"tool_calls"' not in joined - assert '"content": "Hello world"' in joined + assert '"content": "Hello"' in joined + assert '"content": " world"' in joined + assert joined.rstrip().endswith("data: [DONE]") + + +async def test_stream_tool_mode_streams_prefix_then_tool_deltas(): + sse = ( + 'data: {"response.created":{"chat_id":"c1","parent_id":"p0","response_id":"r1","response_index":"0"}} \n' + "\n" + 'data: {"choices": [{"delta": {"content": "Sure, here: ", "phase": "answer", "status": "typing"}}], "response_id": "r1"}\n' + "\n" + f'data: {{"choices": [{{"delta": {{"content": {json.dumps(TOOL_JSON)}, "phase": "answer", "status": "typing"}}}}], "response_id": "r1"}}\n' + "\n" + 'data: {"choices": [{"delta": {"content": "", "role": "assistant", "status": "finished", "phase": "answer"}}], "response_id": "r1"}\n' + "\n" + ) + acct = FakeAccount([sse]) + args = _args(acct, tool_mode=True) + gen = qwen_api.stream_openai(**args) + joined = "".join(await _collect(gen)) + assert '"content": "Sure, here: "' in joined + assert '"tool_calls"' in joined + assert json.dumps(TOOL_JSON) not in joined + assert '"finish_reason": "tool_calls"' in joined + assert joined.rstrip().endswith("data: [DONE]") + + +class _ImgMsg: + def __init__(self, content): + self.content = content + + +def test_append_image_markdown(): + messages = [ + _ImgMsg("plain"), + _ImgMsg([{"type": "image_url", "image_url": {"url": "https://x/y.png"}}]), + _ImgMsg([{"type": "image_url", "image_url": "data:image/png;base64,AAAA"}]), + _ImgMsg([{"type": "image_url", "image_url": 42}]), + ] + prompt = qwen_api._append_image_markdown("hello", messages) + assert prompt.startswith("hello") + assert "![image](https://x/y.png)" in prompt + assert "![image](data:image/png;base64,AAAA)" in prompt + assert qwen_api._append_image_markdown("hello", None) == "hello" + assert qwen_api._append_image_markdown("hello", [_ImgMsg("nope")]) == "hello" + + +async def test_stream_empty_response_sends_role_delta(): + sse = ( + 'data: {"response.created":{"chat_id":"c1","parent_id":"p0","response_id":"r1"}} \n' + "\n" + 'data: {"choices": [{"delta": {"content": "", "role": "assistant", "status": "finished", "phase": "answer"}}], "response_id": "r1"}\n' + "\n" + ) + acct = FakeAccount([sse]) + gen = qwen_api.stream_openai(**_args(acct)) + joined = "".join(await _collect(gen)) + assert '"role": "assistant"' in joined + assert '"finish_reason": "stop"' in joined assert joined.rstrip().endswith("data: [DONE]") diff --git a/tests/test_qwen_stream.py b/tests/test_qwen_stream.py index ec475ff..2583e24 100644 --- a/tests/test_qwen_stream.py +++ b/tests/test_qwen_stream.py @@ -283,14 +283,43 @@ def test_usage_tokens_defaults(): assert rec.usage_tokens == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} +def test_usage_tokens_parses_string_numbers(): + rec = QwenStreamReconstructor() + rec.usage = {"input_tokens": "1234", "output_tokens": "56.0", "total_tokens": "1290.5"} + assert rec.usage_tokens == {"prompt_tokens": 1234, "completion_tokens": 56, "total_tokens": 1290} + + +def test_usage_tokens_booleans_zero(): + rec = QwenStreamReconstructor() + rec.usage = {"input_tokens": True, "output_tokens": False, "total_tokens": True} + assert rec.usage_tokens == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + def test_usage_tokens_ignores_invalid_values(): rec = QwenStreamReconstructor() rec.usage = {"input_tokens": None, "output_tokens": "5", "total_tokens": [15]} - assert rec.usage_tokens == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + assert rec.usage_tokens == {"prompt_tokens": 0, "completion_tokens": 5, "total_tokens": 0} rec.usage = {"input_tokens": 10.0, "output_tokens": 5.0, "total_tokens": 15.0} assert rec.usage_tokens == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} +def test_image_url_split_across_answer_chunks(): + rec = QwenStreamReconstructor() + rec.handle(SSEEvent(None, {"choices": [{"delta": {"content": "see ![cat](https://cdn.qwenlm.ai/", "phase": "answer"}}]})) + assert rec.image_urls == [] + rec.handle(SSEEvent(None, {"choices": [{"delta": {"content": "cat.png) here", "phase": "answer"}}]})) + assert rec.image_urls == ["https://cdn.qwenlm.ai/cat.png"] + assert rec.content == "see ![cat](https://cdn.qwenlm.ai/cat.png) here" + + +def test_image_phase_url_split_across_chunks(): + rec = QwenStreamReconstructor() + rec.handle(SSEEvent(None, {"choices": [{"delta": {"content": "![img](https://cdn.qwenlm.ai/", "phase": "image"}}]})) + rec.handle(SSEEvent(None, {"choices": [{"delta": {"content": "img.webp)", "phase": "image"}}]})) + assert rec.image_urls == ["https://cdn.qwenlm.ai/img.webp"] + assert rec.content == "![img](https://cdn.qwenlm.ai/img.webp)" + + def test_has_content_false_initially(): rec = QwenStreamReconstructor() assert not rec.has_content diff --git a/tests/test_responses.py b/tests/test_responses.py index 696d59c..1948910 100644 --- a/tests/test_responses.py +++ b/tests/test_responses.py @@ -255,6 +255,27 @@ def test_response_from_chat_incomplete_length(): assert obj["incomplete_details"] == {"reason": "max_output_tokens"} +def test_response_from_chat_reasoning_tokens(): + info = resp.RequestInfo(model="m") + chat = { + "choices": [ + { + "message": { + "role": "assistant", + "content": "Hi", + "reasoning_content": "Let me think about this carefully first", + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + } + obj = resp.response_from_chat(chat, info, "resp_1", 1) + details = obj["usage"]["output_tokens_details"] + assert details["reasoning_tokens"] > 0 + assert details["reasoning_tokens"] == 9 + + async def test_translate_stream_text(): info = resp.RequestInfo(model="m") chat_stream = _agen( @@ -292,6 +313,42 @@ async def test_translate_stream_tools(): assert '"name": "f"' in joined +async def test_translate_stream_tool_name_after_args(): + info = resp.RequestInfo(model="m") + chat_stream = _agen( + [ + 'data: {"id":"x","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function",' + '"function":{"arguments":"{\\"a\\":1}"}}]},"finish_reason":null}]}\n\n', + 'data: {"id":"x","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"name":"f","arguments":""}}]},"finish_reason":null}]}\n\n', + 'data: {"id":"x","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}\n\n', + ] + ) + joined = "".join(await _collect(resp.translate_stream(chat_stream, info, "resp_1", 1))) + added_at = joined.index("response.output_item.added") + delta_at = joined.index("response.function_call_arguments.delta") + done_at = joined.index("response.function_call_arguments.done") + assert added_at < delta_at < done_at + assert '"name": "f"' in joined + assert '"delta": "{\\"a\\":1}"' in joined + + +async def test_translate_stream_interrupted_tool_call_closed(): + info = resp.RequestInfo(model="m") + chat_stream = _agen( + [ + 'data: {"id":"x","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function",' + '"function":{"arguments":"{\\"a\\":"}}]},"finish_reason":null}]}\n\n', + ] + ) + joined = "".join(await _collect(resp.translate_stream(chat_stream, info, "resp_1", 1))) + added_at = joined.index("response.output_item.added") + delta_at = joined.index("response.function_call_arguments.delta") + done_at = joined.index("response.function_call_arguments.done") + item_done_at = joined.index("event: response.output_item.done") + assert added_at < delta_at < done_at < item_done_at + assert '"arguments": "{\\"a\\":"}' in joined + + async def test_translate_stream_error(): info = resp.RequestInfo(model="m") chat_stream = _agen(['data: {"id":"x","error":{"message":"boom","finish_reason":"server_busy"},"choices":[]}\n\n']) diff --git a/tests/test_stream.py b/tests/test_stream.py index 623aa8d..f8684ec 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -340,7 +340,7 @@ def test_append_list_negative_index(): rec = MessageReconstructor() rec.message = {"items": ["a", "b"]} rec.handle(SSEEvent(None, {"p": "response/items/-1", "o": "APPEND", "v": "X"})) - assert rec.message["items"] == ["a", "b"] + assert rec.message["items"] == ["a", "bX"] def test_init_message_non_dict(): diff --git a/tests/test_tools.py b/tests/test_tools.py index fab7330..5979efb 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -1799,8 +1799,7 @@ def test_parse_xml_wrapper_element_overlap_mismatched_close(): assert parsed is not None calls, _ = parsed assert calls is not None - assert calls[0].name == "a" - assert json.loads(calls[0].arguments) == {"x": "1"} + assert not any(call.name == "a" for call in calls) def test_parse_xml_wrapper_element_basic(): diff --git a/tests/test_tools_api.py b/tests/test_tools_api.py index c6ae605..a8db8f5 100644 --- a/tests/test_tools_api.py +++ b/tests/test_tools_api.py @@ -1,6 +1,7 @@ import asyncio import json -from unittest.mock import AsyncMock, MagicMock +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -229,3 +230,75 @@ async def test_qwen_stream_emits_tool_call_deltas(): assert '"tool_calls"' in joined assert '"finish_reason": "tool_calls"' in joined assert joined.rstrip().endswith("data: [DONE]") + + +async def test_stream_tool_mode_does_not_leak_json_as_content(): + acct = FakeAccount([DS_TOOL_SSE]) + gen = openai_mod._stream_openai(**_deepseek_args(acct)) + lines = await collect_stream(gen) + for line in lines: + if not line.startswith("data: ") or line.startswith("data: [DONE]"): + continue + payload = json.loads(line[6:]) + for chunk in payload.get("choices") or []: + content = (chunk.get("delta") or {}).get("content") + if content: + assert '{"tool_calls"' not in content + assert "