From 65b101d4d7bf62350fd111cf329f68a54cf39638 Mon Sep 17 00:00:00 2001 From: sohail-eng Date: Fri, 9 Oct 2026 02:38:27 +0500 Subject: [PATCH] fix(wren-ai-service): support Ollama-hosted LLM models and fix JSON field 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. --- docker/.env.example | 4 + docker/README.md | 45 +++++++ docker/config.example.yaml | 19 +++ .../src/pipelines/generation/utils/sql.py | 16 +-- .../retrieval/db_schema_retrieval.py | 17 ++- wren-ai-service/src/providers/llm/litellm.py | 119 ++++++++++++------ 6 files changed, 168 insertions(+), 52 deletions(-) diff --git a/docker/.env.example b/docker/.env.example index 8886a688fa..f7154b9441 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -17,6 +17,10 @@ SHOULD_FORCE_DEPLOY=1 # vendor keys OPENAI_API_KEY= +# OPTIONAL: for Ollama hosted models (https://ollama.com/v1) +OLLAMA_API_KEY= +# OPTIONAL: for Gemini embedder +GEMINI_API_KEY= # version # CHANGE THIS TO THE LATEST VERSION diff --git a/docker/README.md b/docker/README.md index d62873d974..b354733d14 100644 --- a/docker/README.md +++ b/docker/README.md @@ -40,3 +40,48 @@ that we created on the previous step. After modifying the file, you can restart For detailed information on how to modify the configuration for different LLM providers and models, please refer to the [AI Service Configuration](../wren-ai-service/docs/configuration.md). This guide provides comprehensive instructions on setting up various LLM providers, embedders, and other components of the AI service. + +## How to start with Ollama + +Wren AI supports Ollama-hosted LLM models (e.g. `gpt-oss`, `nemotron`) through the OpenAI-compatible API at `https://ollama.com/v1`. + +1. Copy `.env.example` to `.env` and set your `OLLAMA_API_KEY`. + ```env + OLLAMA_API_KEY=your_ollama_api_key + ``` +2. Copy `config.example.yaml` to `config.yaml` and configure the LLM to use Ollama: + ```yaml + type: llm + provider: litellm_llm + timeout: 120 + models: + - alias: default + model: openai/gpt-oss:120b-cloud + api_base: https://ollama.com/v1 + api_key_name: OLLAMA_API_KEY + context_window_size: 131072 + kwargs: + max_completion_tokens: 32768 + n: 1 + temperature: 0.0 + ``` + Note: Ollama does not provide an `/v1/embeddings` endpoint, so you must use a separate embedder (e.g. Gemini). + Configure the embedder in the same `config.yaml`: + ```yaml + type: embedder + provider: litellm_embedder + models: + - model: gemini/gemini-embedding-001 + alias: default + timeout: 120 + ``` + And set `GEMINI_API_KEY` in `.env`. +3. Update `GENERATION_MODEL` in `.env` to match your model name (e.g. `gpt-oss:120b-cloud`). +4. Update `embedding_model_dim` in `config.yaml` to match your embedder's output dimension (e.g. `3072` for Gemini). +5. Start all services: `docker-compose --env-file .env up -d`. + +### Notes on Ollama reasoning models + +- Ollama reasoning models do not support OpenAI's `response_format: json_schema` structured output. Wren AI automatically strips this parameter when an Ollama API base is detected and retries the request if the model returns invalid JSON. +- Reasoning models may occasionally return error messages (e.g. `{"error": "Insufficient information..."}`) instead of the expected JSON. Wren AI retries up to 3 times in such cases. +- `max_completion_tokens` should be set high enough (e.g. `32768`) to accommodate the reasoning tokens. diff --git a/docker/config.example.yaml b/docker/config.example.yaml index 2a22e788e9..efd17a9437 100644 --- a/docker/config.example.yaml +++ b/docker/config.example.yaml @@ -2,6 +2,7 @@ type: llm provider: litellm_llm timeout: 120 models: + # --- OpenAI models --- - alias: default model: gpt-4.1-nano-2025-04-14 context_window_size: 1000000 @@ -45,6 +46,18 @@ models: n: 1 seed: 0 reasoning_effort: minimal + # --- Ollama hosted models (uncomment to use) --- + # Requires: OLLAMA_API_KEY in .env + # Note: Ollama does not support /v1/embeddings, use a separate embedder (e.g. Gemini) + # - alias: default + # model: openai/gpt-oss:120b-cloud + # api_base: https://ollama.com/v1 + # api_key_name: OLLAMA_API_KEY + # context_window_size: 131072 + # kwargs: + # max_completion_tokens: 32768 + # n: 1 + # temperature: 0.0 --- type: embedder provider: litellm_embedder @@ -52,6 +65,12 @@ models: - model: text-embedding-3-large alias: default timeout: 120 + # --- Gemini embedder (for use with Ollama LLM) --- + # Requires: GEMINI_API_KEY in .env + # Note: set embedding_model_dim below to 3072 to match + # - model: gemini/gemini-embedding-001 + # alias: default + # timeout: 120 --- type: engine diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 088282574e..81148f9454 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -364,11 +364,11 @@ async def _classify_generation_result( _DEFAULT_JSON_FIELD_INSTRUCTIONS = """ #### Instructions for JSON related functions #### - ONLY USE JSON_QUERY for querying fields if "json_type":"JSON" is identified in the columns comment, NOT the deprecated JSON_EXTRACT_SCALAR function. - - DON'T USE CAST for JSON fields, ONLY USE the following funtions: - - LAX_BOOL for boolean fields - - LAX_FLOAT64 for double and float fields - - LAX_INT64 for bigint fields - - LAX_STRING for varchar fields + - USE CAST to convert JSON_QUERY results to the appropriate type: + - CAST(... AS BOOLEAN) for boolean fields + - CAST(... AS DOUBLE) for double and float fields + - CAST(... AS BIGINT) for bigint fields + - CAST(... AS VARCHAR) for varchar fields - For Example: DATA SCHEMA: `/* {"alias":"users","description":"A model representing the users data."} */ @@ -377,7 +377,7 @@ async def _classify_generation_result( address JSON )` To get the city of address in user table use SQL: - `SELECT LAX_STRING(JSON_QUERY(u.address, '$.city')) FROM user as u` + `SELECT CAST(JSON_QUERY(u.address, '$.city') AS VARCHAR) FROM user as u` - ONLY USE JSON_QUERY_ARRAY for querying "json_type":"JSON_ARRAY" is identified in the comment of the column, NOT the deprecated JSON_EXTRACT_ARRAY. - USE UNNEST to analysis each item individually in the ARRAY. YOU MUST SELECT FROM the parent table ahead of the UNNEST ARRAY. - The alias of the UNNEST(ARRAY) should be in the format `unnest_table_alias(individual_item_alias)` @@ -391,11 +391,11 @@ async def _classify_generation_result( elements JSON )` To get the number of elements in my_table table use SQL: - `SELECT LAX_INT64(JSON_QUERY(element, '$.number')) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE LAX_FLOAT64(JSON_QUERY(element, '$.value')) > 3.5` + `SELECT CAST(JSON_QUERY(element, '$.number') AS BIGINT) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE CAST(JSON_QUERY(element, '$.value') AS DOUBLE) > 3.5` - To JOIN ON the fields inside UNNEST(ARRAY), YOU MUST SELECT FROM the parent table ahead of the UNNEST syntax, and the alias of the UNNEST(ARRAY) SHOULD BE IN THE FORMAT unnest_table_alias(individual_item_alias) - For Example: `SELECT p.column_1, j.column_2 FROM parent_table AS p, join_table AS j JOIN UNNEST(p.array_column) AS unnested(array_item) ON j.id = array_item.id` - DON'T USE JSON_QUERY and JSON_QUERY_ARRAY when "json_type":"". -- DON'T USE LAX_BOOL, LAX_FLOAT64, LAX_INT64, LAX_STRING when "json_type":"". +- DON'T USE LAX_BOOL, LAX_FLOAT64, LAX_INT64, LAX_STRING. Use CAST instead. """ sql_samples_instructions = """ diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6c8dd7bbe3..f8e6990e54 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -346,9 +346,20 @@ def construct_retrieval_results( dbschema_retrieval: list[Document], ) -> dict[str, Any]: if filter_columns_in_tables: - columns_and_tables_needed = orjson.loads( - filter_columns_in_tables["replies"][0] - )["results"] + try: + columns_and_tables_needed = orjson.loads( + filter_columns_in_tables["replies"][0] + )["results"] + except (KeyError, orjson.JSONDecodeError, TypeError): + retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] + return { + "retrieval_results": retrieval_results, + "has_calculated_field": check_using_db_schemas_without_pruning[ + "has_calculated_field" + ], + "has_metric": check_using_db_schemas_without_pruning["has_metric"], + "has_json_field": check_using_db_schemas_without_pruning["has_json_field"], + } # we need to change the below code to match the new schema of structured output # the objective of this loop is to change the structure of JSON to match the needed format diff --git a/wren-ai-service/src/providers/llm/litellm.py b/wren-ai-service/src/providers/llm/litellm.py index e94918b8c5..43a06fab90 100644 --- a/wren-ai-service/src/providers/llm/litellm.py +++ b/wren-ai-service/src/providers/llm/litellm.py @@ -104,53 +104,90 @@ async def _run( **(generation_kwargs or {}), } + import json as _json + + _expected_keys = set() + _rf = generation_kwargs.get("response_format") + if isinstance(_rf, dict) and _rf.get("type") == "json_schema": + _schema = _rf.get("json_schema", {}).get("schema", {}) + _expected_keys = set(_schema.get("required", [])) + + if self._api_base and "ollama" in self._api_base: + generation_kwargs.pop("response_format", None) + allowed_openai_params = generation_kwargs.get( "allowed_openai_params", [] ) + (["reasoning_effort"] if self._model.startswith("gpt-5") else []) - if self._has_fallbacks: - completion = await self._router.acompletion( - model=self._model, - messages=openai_formatted_messages, - stream=streaming_callback is not None, - allowed_openai_params=allowed_openai_params, - mock_testing_fallbacks=self._enable_fallback_testing, - **generation_kwargs, - ) - else: - completion = await acompletion( - model=self._model, - api_key=self._api_key, - api_base=self._api_base, - api_version=self._api_version, - timeout=self._timeout, - messages=openai_formatted_messages, - stream=streaming_callback is not None, - allowed_openai_params=allowed_openai_params, - **generation_kwargs, - ) - + _max_content_retries = 3 completions: List[ChatMessage] = [] - if streaming_callback is not None: - num_responses = generation_kwargs.pop("n", 1) - if num_responses > 1: - raise ValueError( - "Cannot stream multiple responses, please set n=1." + for _content_attempt in range(_max_content_retries): + if self._has_fallbacks: + completion = await self._router.acompletion( + model=self._model, + messages=openai_formatted_messages, + stream=streaming_callback is not None, + allowed_openai_params=allowed_openai_params, + mock_testing_fallbacks=self._enable_fallback_testing, + **generation_kwargs, ) - chunks: List[StreamingChunk] = [] - - async for chunk in completion: - if chunk.choices and streaming_callback: - chunk_delta: StreamingChunk = build_chunk(chunk) - chunks.append(chunk_delta) - streaming_callback( - chunk_delta, query_id - ) # invoke callback with the chunk_delta - completions = [connect_chunks(chunk, chunks)] - else: - completions = [ - build_message(completion, choice) for choice in completion.choices - ] + else: + completion = await acompletion( + model=self._model, + api_key=self._api_key, + api_base=self._api_base, + api_version=self._api_version, + timeout=self._timeout, + messages=openai_formatted_messages, + stream=streaming_callback is not None, + allowed_openai_params=allowed_openai_params, + **generation_kwargs, + ) + + if streaming_callback is not None: + num_responses = generation_kwargs.pop("n", 1) + if num_responses > 1: + raise ValueError( + "Cannot stream multiple responses, please set n=1." + ) + chunks: List[StreamingChunk] = [] + + async for chunk in completion: + if chunk.choices and streaming_callback: + chunk_delta: StreamingChunk = build_chunk(chunk) + chunks.append(chunk_delta) + streaming_callback( + chunk_delta, query_id + ) + completions = [connect_chunks(chunk, chunks)] + else: + completions = [ + build_message(completion, choice) for choice in completion.choices + ] + + _all_valid = True + for _msg in completions: + if _msg.meta: + 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)) + _extracted = extract_braces_content(_msg.content or "") + if not _extracted.strip(): + _all_valid = False + break + try: + _parsed = _json.loads(_extracted) + except (ValueError, TypeError): + _all_valid = False + break + if _expected_keys and isinstance(_parsed, dict): + if not _expected_keys.issubset(_parsed.keys()): + _all_valid = False + break + + if _all_valid or _content_attempt == _max_content_retries - 1: + break # before returning, do post-processing of the completions for response in completions: