Repository navigation
fix(wren-ai-service): support Ollama-hosted LLM models and fix JSON field instructions for Wren Engine compatibility - #2791
Open
sohail-eng wants to merge 1 commit into
Conversation
…ield instructions for Wren Engine compatibility - litellm.py: Strip response_format json_schema when Ollama API base detected; extract expected keys from schema and retry up to 3 times if LLM returns invalid/empty JSON or missing keys; sanitize meta via json.loads(json.dumps(..., default=str)) for langfuse compatibility - db_schema_retrieval.py: Add try/except fallback on orjson.loads filter_columns_in_tables results to handle non-JSON LLM responses gracefully - sql.py: Replace BigQuery-specific LAX_STRING/LAX_BOOL/LAX_FLOAT64/LAX_INT64 with DataFusion-compatible CAST expressions in _DEFAULT_JSON_FIELD_INSTRUCTIONS - docker/.env.example: Add OLLAMA_API_KEY and GEMINI_API_KEY vendor keys - docker/config.example.yaml: Add commented Ollama LLM and Gemini embedder examples - docker/README.md: Add 'How to start with Ollama' section with setup instructions Tested end-to-end with openai/gpt-oss:120b-cloud via https://ollama.com/v1 and gemini/gemini-embedding-001 embedder.
Contributor
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configuration
You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
github-actions
Bot
added
the
documentation
Author
SummaryHi @goldmedal @yichieh-lu @chilijung 👋 I've opened this PR to add Ollama-hosted LLM model support to Wren AI, and to fix a JSON field SQL instruction issue that was causing runtime errors with the Wren Engine. What this PR does1. Ollama LLM Support (
2. DB Schema Retrieval Fallback (
3. JSON Field SQL Instructions Fix (
4. Documentation (
Why this matters
Test resultsTested end-to-end with Notes
|
goldmedal
requested review from
cyyeh and
yichieh-lu
and removed request for
cyyeh
October 13, 2026 02:52
There was a problem hiding this comment.
🟡 Changes recommended
The retry flow can repeat streamed responses, and malformed pruning replies can still fail or leave SQL generation without schemas.
5 open findings
What changed in this PR
This PR adds Ollama-hosted LLM setup to the Wren AI service and updates JSON-field SQL guidance for Wren Engine.
Changes:
- Adds Ollama configuration examples and setup instructions, including a separate Gemini embedder.
- Adds LLM response validation and a schema-retrieval fallback, and replaces BigQuery-specific JSON conversion guidance with
CAST.
| File | Description |
|---|---|
wren-ai-service/src/providers/llm/litellm.py |
Adds Ollama response handling, retries, and metadata sanitization. |
wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py |
Adds a fallback for failed pruning responses. |
wren-ai-service/src/pipelines/generation/utils/sql.py |
Updates JSON-field SQL instructions to use CAST. |
docker/README.md |
Documents Ollama setup. |
docker/config.example.yaml |
Adds Ollama and Gemini configuration examples. |
docker/.env.example |
Adds example API-key variables. |
🧠 Review effort: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| **generation_kwargs, | ||
| ) | ||
|
|
||
| _max_content_retries = 3 |
Comment on lines
+349
to
+353
| try: | ||
| columns_and_tables_needed = orjson.loads( | ||
| filter_columns_in_tables["replies"][0] | ||
| )["results"] | ||
| except (KeyError, orjson.JSONDecodeError, TypeError): |
| filter_columns_in_tables["replies"][0] | ||
| )["results"] | ||
| except (KeyError, orjson.JSONDecodeError, TypeError): | ||
| retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] |
| try: | ||
| _msg.meta = _json.loads(_json.dumps(_msg.meta, default=str)) | ||
| except (TypeError, ValueError): | ||
| _msg.meta = _json.loads(_json.dumps(_msg.meta, default=str, ignore_nan=True)) |
Comment on lines
+184
to
+187
| if _expected_keys and isinstance(_parsed, dict): | ||
| if not _expected_keys.issubset(_parsed.keys()): | ||
| _all_valid = False | ||
| break |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds support for Ollama-hosted LLM models (e.g.
gpt-oss,nemotron) to the Wren AI service, and fixes JSON field SQL instructions for Wren Engine (DataFusion) compatibility.Changes
wren-ai-service/src/providers/llm/litellm.pyresponse_formatfor Ollama: Ollama reasoning models do not support OpenAI'sresponse_format: json_schemastructured output. Whenapi_basecontains "ollama", theresponse_formatparameter is removed before the LLM call.response_format, the expected keys are extracted from thejson_schemadefinition so they can be validated post-hoc.metadict is sanitized viajson.loads(json.dumps(meta, default=str))to handle non-serializable values (e.g.Noneintegers) that cause langfuseUpdateGenerationBodyvalidation errors.wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.pytry/exceptaroundorjson.loads(filter_columns_in_tables["replies"][0])["results"]catchingKeyError,JSONDecodeError, andTypeError. Falls back to unpruned db schemas when the LLM returns non-JSON or unexpected content.wren-ai-service/src/pipelines/generation/utils/sql.py_DEFAULT_JSON_FIELD_INSTRUCTIONSprompt previously instructed the LLM to useLAX_STRING,LAX_BOOL,LAX_FLOAT64,LAX_INT64— these are BigQuery-specific functions not supported by the Wren Engine (DataFusion). Replaced with DataFusion-compatibleCAST(... AS VARCHAR),CAST(... AS BOOLEAN),CAST(... AS DOUBLE),CAST(... AS BIGINT).docker/.env.exampleOLLAMA_API_KEYandGEMINI_API_KEYvendor key entries.docker/config.example.yamlopenai/gpt-oss:120b-cloudviahttps://ollama.com/v1) and Gemini embedder (gemini/gemini-embedding-001).docker/README.mdTest Results
Tested end-to-end with:
openai/gpt-oss:120b-cloudviahttps://ollama.com/v1gemini/gemini-embedding-001(3072 dimensions)lax_stringor structured output errors.Notes
/v1/embeddingsendpoint, so a separate embedder (e.g. Gemini) is required.max_completion_tokensshould be set high (e.g.32768) to accommodate reasoning tokens.enable_column_pruning: falseis recommended when the LLM's column pruning responses are unreliable.