Upload deepseek.py

llm_services/providers/deepseek.py

@@ -1,14 +1,20 @@
-"""DeepSeek OpenAI-compatible Chat Completions adapter for research-only calls.
+"""DeepSeek research-only adapter for WaterScope-AI.
 
-Secrets live in HF Space environment, never in code/client-facing responses.
-Non-thinking mode is explicit by default, and can be configured with the
-DEEPSEEK_THINKING HF Space variable. Record it for experimental provenance.
+Preserves the previous DeepSeek integration, HF/Azure routing, user prompts,
+full paper evidence, and scientific answer word ceiling. In thinking mode,
+DeepSeek `max_tokens` counts *both* reasoning and final answer tokens; the
+research endpoint's 2-tokens-per-word allowance is only for final text.
+This adapter adds an independent reasoning headroom so thinking responses
+are not prematurely cut off. The setting is a ceiling, not a token target.
 """
 from __future__ import annotations
+
 import os
 from openai import OpenAI
 
 _ALLOWED_THINKING = {"enabled", "disabled"}
+_ALLOWED_EFFORT = {"low", "high", "max"}
+_DEFAULT_REASONING_RESERVE = 32768
 
 
 def _config():
@@ -18,6 +24,21 @@ def _config():
     return mode
 
 
+def _reasoning_config():
+    effort = os.getenv("DEEPSEEK_REASONING_EFFORT", "high").strip().lower()
+    if effort not in _ALLOWED_EFFORT:
+        raise ValueError("DEEPSEEK_REASONING_EFFORT must be low, high or max")
+
+    raw = os.getenv("DEEPSEEK_REASONING_RESERVE_TOKENS", str(_DEFAULT_REASONING_RESERVE))
+    try:
+        reserve = int(raw)
+    except (ValueError, TypeError) as exc:
+        raise ValueError("DEEPSEEK_REASONING_RESERVE_TOKENS must be an integer") from exc
+    if not 0 <= reserve <= 100000:
+        raise ValueError("DEEPSEEK_REASONING_RESERVE_TOKENS must be 0-100000")
+    return effort, reserve
+
+
 def create_client():
     key = os.environ["DEEPSEEK_API_KEY"]
     return OpenAI(
@@ -31,23 +52,43 @@ def create_client():
 def complete(messages, *, model=None, max_completion_tokens=1200):
     requested = model or os.getenv("DEEPSEEK_MODEL", "deepseek-flash")
     mode = _config()
-    # DeepSeek's Chat Completions interface uses max_tokens, not Azure's
-    # max_completion_tokens. With thinking disabled, it budgets final text.
-    resp = create_client().chat.completions.create(
-        model=requested,
-        messages=messages,
-        stream=False,
-        max_tokens=int(max_completion_tokens),
-        extra_body={"thinking": {"type": mode}},
-    )
+
+    # This budget is passed in from the research endpoint, where it is sized
+    # for the final visible answer (300 words per selected paper).
+    final_text_budget = int(max_completion_tokens)
+    if final_text_budget < 1:
+        raise ValueError("max_completion_tokens must be >= 1")
+
+    effort = None
+    reserve = 0
+    if mode == "enabled":
+        effort, reserve = _reasoning_config()
+
+    # DeepSeek Chat Completions `max_tokens` includes reasoning and final text.
+    # Do not impose the original final-text-only budget on thinking mode.
+    api_max_tokens = final_text_budget + reserve
+    params = {
+        "model": requested,
+        "messages": messages,
+        "stream": False,
+        "max_tokens": api_max_tokens,
+        "extra_body": {"thinking": {"type": mode}},
+    }
+    if effort is not None:
+        params["reasoning_effort"] = effort
+
+    resp = create_client().chat.completions.create(**params)
     if not resp.choices:
         raise RuntimeError("DeepSeek returned no choices")
+
     choice = resp.choices[0]
     usage = resp.usage
     prompt_details = getattr(usage, "prompt_tokens_details", None) if usage else None
-    cached = (getattr(usage, "prompt_cache_hit_tokens", None) if usage else None)
+    completion_details = getattr(usage, "completion_tokens_details", None) if usage else None
+    cached = getattr(usage, "prompt_cache_hit_tokens", None) if usage else None
     if cached is None and prompt_details is not None:
         cached = getattr(prompt_details, "cached_tokens", None)
+
     return {
         "answer": choice.message.content or "",
         "provider": "deepseek",
@@ -55,10 +96,18 @@ def complete(messages, *, model=None, max_completion_tokens=1200):
         "model_returned": resp.model,
         "finish_reason": choice.finish_reason,
         "thinking_mode": mode,
+        "thinking_effort": effort,
+        "final_text_token_allowance": final_text_budget,
+        "reasoning_token_reserve": reserve,
+        "max_tokens_sent": api_max_tokens,
         "usage": {
             "input_tokens": usage.prompt_tokens if usage else None,
             "output_tokens": usage.completion_tokens if usage else None,
             "total_tokens": usage.total_tokens if usage else None,
             "cached_input_tokens": cached,
+            "reasoning_tokens": (
+                getattr(completion_details, "reasoning_tokens", None)
+                if completion_details is not None else None
+            ),
         },
     }