Adding deepseek - #35
Conversation
Add DeepSeek branch to LLMClient._structured_output_bind_kwargs that returns
{"response_format": {"type": "json_object"}} instead of strict json_schema,
since DeepSeek's live API rejects json_schema mode (HTTP 400).
Branch placement: Top of method before ChatAnthropic check, since DeepSeek
is implemented as a ChatOpenAI instance.
Tests added:
- test_structured_kwargs_deepseek_uses_json_object: Verifies DeepSeek returns json_object
- test_reasoning_effort_reaches_chat_openai: Confirms reasoning_effort flows through to ChatOpenAI
- _deepseek_client(): Helper to create test DeepSeek client
Add _maybe_inject_schema method to LLMClient that appends a JSON schema instruction as a HumanMessage when using DeepSeek with structured output. This enables DeepSeek's json_object mode to understand the desired schema, while still permitting tool use for the proposer's search functionality.
Added DeepSeek v4-pro LLM configuration to both YAML files: - configs/llms.yaml (experimental configurations) - src/ax_prover/configs/llms.yaml (bundled package copy) Configuration includes: - Model: deepseek-v4-pro - Provider: OpenAI-compatible (base_url: https://api.deepseek.com) - API key: environment variable DEEPSEEK_API_KEY - Reasoning effort: high Also added test_deepseek_llms_config_entry to verify the configuration is correctly loaded via OmegaConf with env variable resolution.
Create deepseek_local.yaml run config mirroring opus48_local.yaml structure but selecting DeepSeek V4 Pro. Includes test to verify config loads correctly with expected model, restrictions, and tool set.
…via merge_configs configs/deepseek_local.yaml referenced tools (search_lean_local, search_cslib) and a restrict_to_proof_body field that don't exist on this branch, crashing at resolution time. Mirror default.yaml's tool set instead. Rewrote the test to exercise the real production path (merge_configs, which resolves interpolations) instead of a non-resolving merge, so a dangling interpolation now fails the test.
Adds an opt-in, network-gated test that runs three concurrent LLMClient.ainvoke calls against the live DeepSeek API with output_schema, confirming endpoint wiring, json_object structured output with schema injection, and parallel-safety end to end.
… the suite deterministic
DeepSeek and qwen route through ChatOpenAI with a custom base_url, for which init_chat_model provides no model profile. profile.get( "max_input_tokens") then returns None, crashing ProverAgent.__init__ at `self.max_input_tokens < 1000` (TypeError: '<' not supported between NoneType and int) and again in _build_error_processing. Fall back to DEFAULT_MAX_INPUT_TOKENS (128k) with a warning when the profile carries no value, so any unprofiled OpenAI-compatible model runs.
DeepSeek routes structured output via response_format={"type":"json_object"}.
When tools are also bound, langchain_openai (1.1.x) routes the call through
OpenAI's .parse() helper, which client-side rejects any non-strict tool
(`ValueError: <tool> is not strict. Only strict function tools can be
auto-parsed`). DeepSeek cannot use strict tools, so json_object and tools are
incompatible on that path — this broke the proposer, which binds both the
search tools and ProverResult.
Skip the response_format bind for DeepSeek whenever tools are present
(_should_bind_structured_output). The schema is still delivered via the
prompt injection in _maybe_inject_schema (has_tools branch), and agentic_loop's
final-answer call binds no tools, so it still gets json_object enforcement.
DeepSeek returns its chain-of-thought in each choice's `reasoning_content`, but langchain_openai's base ChatOpenAI deliberately drops non-standard fields (its docstring says to use a provider-specific subclass). As a result get_reasoning() returned "", ProposalMessage.reasoning was empty, and LangSmith traces showed only the final Lean code — no reasoning anywhere. Add _DeepSeekChatOpenAI, a thin ChatOpenAI subclass that copies reasoning_content into the message's additional_kwargs in _create_chat_result. create_llm now routes DeepSeek configs to this subclass (staying on ChatOpenAI so all _is_deepseek_model detection and structured-output handling still work). get_reasoning's existing additional_kwargs fallback then picks it up, and LangSmith serializes additional_kwargs, so the reasoning is logged. Verified against the live API: reasoning_content now appears in additional_kwargs and get_reasoning() returns the full trace.
DeepSeek runs left the LangSmith Cost columns blank and made tool-usage tracking flaky. Root causes: - Cost (total/input/output) is computed natively by LangSmith as tokens x price-map; deepseek-v4-pro isn't in that map, so tokens showed but cost didn't. There was no cost evaluator either. - tool_usage queried list_runs at eval time and could throw (row shows an errored evaluator) or race the async child-run upload (undercount / zero for some theorems). Changes: - Add token_cost evaluator: input/output/total cost in USD from trace token usage x configured per-1M-token prices (prover_llm.input_token_price / output_token_price). Prefers aggregated root-run tokens, falls back to summing LLM runs; returns no feedback when prices are unset. - Harden tool_usage: never throws (guards LangSmith errors), and shares a _list_trace_runs helper that retries until the trace is populated, tolerating the flush race. - Add input_token_price/output_token_price to LLMConfig; set DeepSeek prices in llms.yaml. - Wire _token_cost into the experiment evaluators; unit tests for cost math, cost evaluator (root/fallback/error), and tool_usage hardening.
Registered a LangSmith model-price-map entry matching the bare model name our runs report (ls_model_name=deepseek-v4-pro; the existing map entry only matched the Fireworks path, so cost stayed null). Prices are DeepSeek's direct API rates: $0.435/1M input (cache miss), $0.87/1M output, with a $0.003625/1M cache-read detail. Cost now populates natively on future runs. Since native cost covers it, revert the custom token_cost evaluator added in b2ea06d (removes the duplicate input/output/total_cost feedback keys). The tool_usage hardening (never-throws + _list_trace_runs flush-race retry) is kept. LLMConfig price fields are retained as documentation, kept in sync with the price-map entry.
Issue 1 (run-terminating false positives): SearchTacticsDetectedFeedback dumped the entire declaration as its 'Locations', so the prover could not tell which tactic was flagged and stripped out innocent tactics (field_simp, apply) from a compiling proof. Report each offending tactic with its line number instead. The detection regex itself was already correct. Issue 2 (startup crashes, 0 iterations): a file that fails to compile comes back from lean_interact as a LeanError (message only, no declarations), so list_declarations_from_* blew up with an opaque AttributeError on '.declarations'. Detect LeanError and raise the actual Lean compile message.
The single AutoLeanServer subprocess is shared by all concurrent experiment samples, but run() only locked server creation, not command execution. Firing many full-Mathlib FileCommands at one REPL at once races the subprocess and trips its memory backstop, surfacing as spurious LeanErrors (the 0-iteration startup crashes). Add a run lock so commands are serialized; lake builds are unaffected (separate subprocesses).
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
BorjaRequena
left a comment
BorjaRequena
left a comment
There was a problem hiding this comment.
Thanks for looking into supporting DeepSeek models, @KrystianNowakowski!
I believe we can simplify the code by using some of the tools provided by LangSmith and LangChain direclty. Also, there are a few fixes about our pipeline that I wonder whether these are motivated by finding these bugs during the experimentation or whether they are meant to prevent issues. Maybe we could pull them out of this branch and implement them on a separate PR to prevent scope creep in this one? Although if they're meant to prevent issues, I'd rather skip them and address them once we suffer from them.
| self._run_lock = asyncio.Lock() # serialize commands: the REPL is one subprocess | ||
|
|
||
| async def run(self, command: Command) -> CommandResponse | LeanError: | ||
| "Run the command and return the response or error." | ||
| """Run the command and return the response or error. | ||
|
|
||
| Commands are serialized: the REPL is a single subprocess, and concurrent | ||
| callers (experiment samples share one server) would otherwise race on it or | ||
| trip its memory backstop, surfacing as spurious LeanErrors. | ||
| """ | ||
| server = await self._get_server() | ||
| return await server.async_run(command) | ||
| async with self._run_lock: | ||
| return await server.async_run(command) |
There was a problem hiding this comment.
Did you run into this issue of having different concurrent callers from the various samples in the experiment crash? I wonder whether this change was motivated by a real crash or whether it's a preventive change
There was a problem hiding this comment.
That was real crash of the experiment, I got some runs that died before making a single iteration. It showed up on real runs (e.g. putnam_1962_a6, putnam_1978_a1) in the tiny-Putnam set.
I am moving the fix to a separate PR, because it is not really Deepseek related.
| def _bundle_response(response: CommandResponse | LeanError, source: str) -> list[Declaration]: | ||
| """Turn a REPL response into declarations, failing loudly on a fatal Lean error. | ||
|
|
||
| A file that does not compile at all (e.g. a missing import) comes back as a `LeanError`, | ||
| which carries no declarations. Surfacing its message here turns an opaque downstream | ||
| `AttributeError` into a clear compilation failure the caller can report. | ||
| """ | ||
| if isinstance(response, LeanError): | ||
| raise RuntimeError(f"Lean failed to process {source}:\n{response.message}") | ||
| return _bundle_declarations(response.declarations, response.sorries, response.tactics) |
There was a problem hiding this comment.
Same here, was there any LeanError that appeared in the evaluation? I never managed to get one with basic testing, but I guess they could have risen when running a larger-scale experiment.
There was a problem hiding this comment.
This was the same case as above. Here's verbatim description from agent's bug report:
putnam_1962_a6_sol / putnam_1978_a1_sol crashed in <40s because a file that fails to compile comes back from lean_interact as a LeanError (which has only .message, no .declarations). list_declarations_from_file blindly read .declarations → opaque AttributeError → caught and reported as a generic "exception" with 0 iterations. That's why the error message didn't match the file content. Now it raises the actual Lean compile error:
RuntimeError: Lean failed to process …
| def _is_deepseek_model(llm: BaseChatModel) -> bool: | ||
| """True when this ChatOpenAI instance is backed by the DeepSeek endpoint. | ||
|
|
||
| DeepSeek is routed via model_provider="openai", so it is a ChatOpenAI | ||
| instance and cannot be told apart from real OpenAI by isinstance alone. | ||
| """ | ||
| if not isinstance(llm, ChatOpenAI): | ||
| return False | ||
| model = (getattr(llm, "model_name", "") or "").lower() | ||
| base_url = str(getattr(llm, "openai_api_base", "") or "").lower() | ||
| return model.startswith("deepseek") or "deepseek" in base_url | ||
|
|
||
|
|
||
| def _is_deepseek_config(config: LLMConfig) -> bool: | ||
| """True when this config selects DeepSeek via the OpenAI-compatible endpoint.""" | ||
| provider_config = config.provider_config or {} | ||
| if provider_config.get("model_provider") != "openai": | ||
| return False | ||
| model = (config.model or "").lower() | ||
| base_url = str(provider_config.get("base_url") or "").lower() | ||
| return model.startswith("deepseek") or "deepseek" in base_url | ||
|
|
||
|
|
||
| class _DeepSeekChatOpenAI(ChatOpenAI): | ||
| """ChatOpenAI variant that surfaces DeepSeek's non-standard `reasoning_content`. | ||
|
|
||
| Base ChatOpenAI intentionally drops provider-specific fields, so DeepSeek's | ||
| chain-of-thought (returned in each choice's `reasoning_content`) never reaches | ||
| the AIMessage. We copy it into `additional_kwargs`, where `get_reasoning` reads | ||
| it and LangSmith logs it alongside the rest of the message. | ||
| """ | ||
|
|
||
| def _create_chat_result(self, response, generation_info=None): | ||
| result = super()._create_chat_result(response, generation_info) | ||
| response_dict = response if isinstance(response, dict) else response.model_dump() | ||
| choices = response_dict.get("choices") or [] | ||
| for generation, choice in zip(result.generations, choices, strict=False): | ||
| reasoning = (choice.get("message") or {}).get("reasoning_content") | ||
| if reasoning: | ||
| generation.message.additional_kwargs["reasoning_content"] = reasoning | ||
| return result |
There was a problem hiding this comment.
Maybe we could try using langsmith-deepseek instead of creating our own custom class? The additional reasoning extraction logic could still be added into the get_reasoning function. I believe we don't want to be adding extra specific stuff for every provider and checking against them every time. Hopefully, that's what we're aiming to avoid by using LangSmith, even if it sucks some times
There was a problem hiding this comment.
Oh, I thought Langsmith had no deepseek support, we spoke about that in the meeting and someone suggested using the Anthropic or OpenAI client was suggested.
Of course, better to use their stuff than our own custom class. Gonna add that.
| messages = self._maybe_inject_schema(messages, output_schema, has_tools=bool(tools)) | ||
| runnable = self._get_runnable( | ||
| tools=tools, output_schema=output_schema, retry_config=effective_retry | ||
| ) | ||
| return await runnable.ainvoke(messages) | ||
|
|
||
| def _maybe_inject_schema( | ||
| self, | ||
| messages: LanguageModelInput, | ||
| output_schema: type[BaseModel] | None, | ||
| has_tools: bool = False, | ||
| ) -> LanguageModelInput: | ||
| """Append a JSON-schema instruction for DeepSeek's json_object mode. |
There was a problem hiding this comment.
Once we have the ChatDeepSeek from langsmith, then we can check out whether we need to do this schema injection and how it works with the model having tools
There was a problem hiding this comment.
Possibly, testing.
There was a problem hiding this comment.
Nope, thinking mode doesn't work with structured output. (Non-thinking mode does though)
| # LangSmith uploads child runs (LLM calls, tool calls) from a background thread, | ||
| # so they may not be queryable the instant an evaluator fires. Retry listing the | ||
| # trace until it looks populated (an LLM run is always present in a proof run). | ||
| _TRACE_LIST_RETRIES = 4 | ||
| _TRACE_LIST_RETRY_WAIT_S = 1.5 | ||
|
|
||
|
|
||
| def _list_trace_runs(client: Client, trace_id) -> list[Run]: | ||
| """List a trace's runs, retrying until the trace looks populated. | ||
|
|
||
| Tolerates LangSmith's asynchronous run upload: returns as soon as an LLM run | ||
| appears (every proof run makes at least one LLM call), otherwise falls back | ||
| to the last attempt's result. | ||
| """ |
There was a problem hiding this comment.
Did you encounter this issue? I believe we've never had this problem (?)
There was a problem hiding this comment.
I don't remember clearly. I think so, but I asked a friend to check the traces again from some test runs and it found no evidence. I removed it, it is easy to reimplement if we hit a problem.
| # Prices in USD per 1M tokens, documenting the model's cost near its config. | ||
| # Cost is tracked natively by LangSmith via its model price map (keyed on the | ||
| # model name); keep these in sync with that entry. Leave null if unknown. | ||
| input_token_price: float | None = None | ||
| output_token_price: float | None = None |
There was a problem hiding this comment.
I believe it's better if we add this info on LangSmith. We can add the corresponding input and output token prices for every model that we want. Then LangSmith will match the model name we provide and it'll count it accordingly. This will make everything simpler, as we can also remove it form the config files
There was a problem hiding this comment.
Yes, that's better! Will do.
Drop the count-stabilization retry and never-throws guard so we can empirically check whether undercounts or eval-time throws actually occur.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 2 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit c7a8321. Configure here.
|
@BorjaRequena, I addressed your comments, would you like to take another look? |
2 participants
Add DeepSeek V4 Pro support
Adds DeepSeek V4 Pro model. DeepSeek is served over an OpenAI-compatible endpoint, which needs different handling than the native providers — structured output, tool calling, reasoning capture, and cost tracking all required provider-specific plumbing.
What's included
Provider integration (
utils/llm.py)ChatOpenAIand route it through a_DeepSeekChatOpenAIsubclass that preservesreasoning_contentjson_objectmode with the schema injected into the prompt (DeepSeek has no strictjson_schemasupport)get_reasoningfalls back toreasoning_contentso DeepSeek's thinking is captured and loggedConfiguration (
configs/,config.py)deepseek_v4_proLLM config anddeepseek_localrun configDEEPSEEK_API_KEYresolver defaults tonullso other configs still resolve when it's unsetmax_input_tokensfield onLLMConfigfor providers LangChain can't profile (removes a startup warning and lets the trim budget be set per model)Cost tracking (
evaluators.py)Robustness fixes (
utils/lean_interact.py,utils/lean_parsing.py,prover/agent.py)LeanErrorcrashes from concurrent experiment samples racing the shared subprocess)LeanErrorin declaration parsing instead of crashing withAttributeErrorTests
Added coverage for the DeepSeek LLM path, evaluators, REPL locking, Lean-error parsing, and cheat detection. Full unit suite passes.
Notes
AX_PROVER_LIVE_TESTSto keep the default suite deterministic.Note
Medium Risk
Touches core prover LLM invocation, structured output, and shared Lean REPL concurrency; mistakes could affect all models' tool loops or experiment stability under parallel runs.
Overview
Adds DeepSeek V4 Pro as a prover model via a new
deepseek_v4_proLLM entry,deepseek_localrun config, andlangchain-deepseek(ChatDeepSeekwithDEEPSEEK_API_KEYoptional in merged configs).LLM layer routes DeepSeek configs through
ChatDeepSeek, usesjson_objectplus prompt-injected JSON schema (no strictjson_schema/ strict tools), and adjustsagentic_loopso the final answer is a tool-free call when needed. Reasoning is read fromreasoning_contentwhen content blocks are empty.Prover gains configurable
max_input_tokensonLLMConfig, resolves context budget from config/profile/default, and trims long build output with a ~2 chars/token budget instead of treating tokens as characters.Lean robustness: serializes REPL
run()with a lock for concurrent samples; declaration listing raises a clearRuntimeErroron fatalLeanError; search-tactic cheat feedback lists offending lines instead of the full proof.Tests cover tool usage counting, REPL serialization, and Lean error surfacing.
Reviewed by Cursor Bugbot for commit 6ef4385. Bugbot is set up for automated code reviews on this repo. Configure here.