fix(wren-ai-service): support Ollama-hosted LLM models and fix JSON field instructions for Wren Engine compatibility

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

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.

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,13 +46,31 @@ 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
 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

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 = """

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

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: