langchain-timbr 1.5.2__tar.gz → 1.5.4__tar.gz

{langchain_timbr-1.5.2 → langchain_timbr-1.5.4}/.github/workflows/publish.yml

@@ -16,7 +16,7 @@ jobs:
         with: { fetch-depth: 0 }  # IMPORTANT: tags available to hatch-vcs
       - uses: actions/setup-python@v5
         with: { python-version: "3.12" }
-      - run: python -m pip install --upgrade build twine
+      - run: python -m pip install --upgrade build
       - run: python -m build               # hatch-vcs injects version from tag
       - name: Publish to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1

{langchain_timbr-1.5.2 → langchain_timbr-1.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langchain-timbr
-Version: 1.5.2
+Version: 1.5.4
 Summary: LangChain & LangGraph extensions that parse LLM prompts into Timbr semantic SQL and execute them.
 Project-URL: Homepage, https://github.com/WPSemantix/langchain-timbr
 Project-URL: Documentation, https://docs.timbr.ai/doc/docs/integration/langchain-sdk/
@@ -28,15 +28,18 @@ Requires-Dist: langgraph>=0.3.20
 Requires-Dist: pydantic==2.10.4
 Requires-Dist: pytimbr-api>=2.0.0
 Requires-Dist: tiktoken==0.8.0
-Requires-Dist: transformers>=4.51.3
+Requires-Dist: transformers>=4.53
 Provides-Extra: all
 Requires-Dist: anthropic==0.42.0; extra == 'all'
+Requires-Dist: databricks-langchain==0.3.0; (python_version < '3.10') and extra == 'all'
+Requires-Dist: databricks-langchain==0.7.1; (python_version >= '3.10') and extra == 'all'
+Requires-Dist: databricks-sdk==0.64.0; extra == 'all'
 Requires-Dist: google-generativeai==0.8.4; extra == 'all'
 Requires-Dist: langchain-anthropic>=0.3.1; extra == 'all'
 Requires-Dist: langchain-google-genai>=2.0.9; extra == 'all'
 Requires-Dist: langchain-openai>=0.3.16; extra == 'all'
 Requires-Dist: langchain-tests>=0.3.20; extra == 'all'
-Requires-Dist: openai==1.77.0; extra == 'all'
+Requires-Dist: openai>=1.77.0; extra == 'all'
 Requires-Dist: pyarrow<19.0.0; extra == 'all'
 Requires-Dist: pytest==8.3.4; extra == 'all'
 Requires-Dist: snowflake-snowpark-python>=1.6.0; extra == 'all'
@@ -45,6 +48,13 @@ Requires-Dist: uvicorn==0.34.0; extra == 'all'
 Provides-Extra: anthropic
 Requires-Dist: anthropic==0.42.0; extra == 'anthropic'
 Requires-Dist: langchain-anthropic>=0.3.1; extra == 'anthropic'
+Provides-Extra: azure-openai
+Requires-Dist: langchain-openai>=0.3.16; extra == 'azure-openai'
+Requires-Dist: openai>=1.77.0; extra == 'azure-openai'
+Provides-Extra: databricks
+Requires-Dist: databricks-langchain==0.3.0; (python_version < '3.10') and extra == 'databricks'
+Requires-Dist: databricks-langchain==0.7.1; (python_version >= '3.10') and extra == 'databricks'
+Requires-Dist: databricks-sdk==0.64.0; extra == 'databricks'
 Provides-Extra: dev
 Requires-Dist: langchain-tests>=0.3.20; extra == 'dev'
 Requires-Dist: pyarrow<19.0.0; extra == 'dev'
@@ -55,7 +65,7 @@ Requires-Dist: google-generativeai==0.8.4; extra == 'google'
 Requires-Dist: langchain-google-genai>=2.0.9; extra == 'google'
 Provides-Extra: openai
 Requires-Dist: langchain-openai>=0.3.16; extra == 'openai'
-Requires-Dist: openai==1.77.0; extra == 'openai'
+Requires-Dist: openai>=1.77.0; extra == 'openai'
 Provides-Extra: snowflake
 Requires-Dist: snowflake-snowpark-python>=1.6.0; extra == 'snowflake'
 Requires-Dist: snowflake>=0.8.0; extra == 'snowflake'
@@ -80,15 +90,23 @@ Timbr LangChain LLM SDK is a Python SDK that extends LangChain and LangGraph wit
 
 ## Dependencies
 - Access to a timbr-server
-- Python from 3.9.13 or newer
+- Python 3.9.13 or newer
 
 ## Installation
 
 ### Using pip
+
 ```bash
 python -m pip install langchain-timbr
 ```
 
+### Install with selected LLM providers
+#### One of: openai, anthropic, google, azure_openai, snowflake, databricks (or 'all')
+
+```bash
+python -m pip install 'langchain-timbr[<your selected providers, separated by comma w/o space]'
+```
+
 ### Using pip from github
 ```bash
 pip install git+https://github.com/WPSemantix/langchain-timbr

{langchain_timbr-1.5.2 → langchain_timbr-1.5.4}/README.md

@@ -17,15 +17,23 @@ Timbr LangChain LLM SDK is a Python SDK that extends LangChain and LangGraph wit
 
 ## Dependencies
 - Access to a timbr-server
-- Python from 3.9.13 or newer
+- Python 3.9.13 or newer
 
 ## Installation
 
 ### Using pip
+
 ```bash
 python -m pip install langchain-timbr
 ```
 
+### Install with selected LLM providers
+#### One of: openai, anthropic, google, azure_openai, snowflake, databricks (or 'all')
+
+```bash
+python -m pip install 'langchain-timbr[<your selected providers, separated by comma w/o space]'
+```
+
 ### Using pip from github
 ```bash
 pip install git+https://github.com/WPSemantix/langchain-timbr

{langchain_timbr-1.5.2 → langchain_timbr-1.5.4}/pyproject.toml

@@ -32,15 +32,21 @@ dependencies = [
   "pydantic==2.10.4",
   "pytimbr-api>=2.0.0",
   "tiktoken==0.8.0",
-  "transformers>=4.51.3"
+  "transformers>=4.53"
 ]
 
 [project.optional-dependencies]
 # LLM providers
-openai = ["openai==1.77.0", "langchain-openai>=0.3.16"]
+openai = ["openai>=1.77.0", "langchain-openai>=0.3.16"]
+azure_openai = ["openai>=1.77.0", "langchain-openai>=0.3.16"]
 anthropic = ["anthropic==0.42.0", "langchain-anthropic>=0.3.1"]
 google = ["langchain-google-genai>=2.0.9", "google-generativeai==0.8.4"]
 snowflake = ["snowflake>=0.8.0", "snowflake-snowpark-python>=1.6.0"]
+databricks = [
+  "databricks-langchain==0.3.0; python_version < '3.10'",
+  "databricks-langchain==0.7.1; python_version >= '3.10'",
+  "databricks-sdk==0.64.0"
+]
 
 # Development and testing
 dev = [
@@ -55,7 +61,7 @@ all = [
   "anthropic==0.42.0",
   "google-generativeai==0.8.4",
   "langchain-anthropic>=0.3.1",
-  "openai==1.77.0", 
+  "openai>=1.77.0",
   "langchain-openai>=0.3.16",
   "langchain-google-genai>=2.0.9",
   "snowflake>=0.8.0",
@@ -63,7 +69,10 @@ all = [
   "pytest==8.3.4",
   "langchain-tests>=0.3.20",
   "pyarrow<19.0.0",
-  "uvicorn==0.34.0"
+  "uvicorn==0.34.0",
+  "databricks-langchain==0.3.0; python_version < '3.10'",
+  "databricks-langchain==0.7.1; python_version >= '3.10'",
+  "databricks-sdk==0.64.0"
 ]
 
 [project.urls]

{langchain_timbr-1.5.2 → langchain_timbr-1.5.4}/requirements.txt

@@ -1,5 +1,8 @@
 anthropic==0.42.0
 cryptography>=44.0.3
+databricks-langchain==0.3.0; python_version < '3.10'
+databricks-langchain==0.7.1; python_version >= '3.10'
+databricks-sdk==0.64.0
 google-generativeai==0.8.4
 langchain>=0.3.25
 langchain-anthropic>=0.3.1
@@ -9,7 +12,7 @@ langchain-google-genai>=2.0.9
 langchain-openai>=0.3.16
 langchain-tests>=0.3.20
 langgraph>=0.3.20
-openai==1.77.0
+openai>=1.77.0
 pyarrow<19.0.0
 pydantic==2.10.4
 pytest==8.3.4
@@ -17,5 +20,5 @@ pytimbr-api>=2.0.0
 snowflake>=0.8.0
 snowflake-snowpark-python>=1.6.0
 tiktoken==0.8.0
-transformers>=4.51.3
+transformers>=4.53
 uvicorn==0.34.0

{langchain_timbr-1.5.2 → langchain_timbr-1.5.4}/src/langchain_timbr/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '1.5.2'
-__version_tuple__ = version_tuple = (1, 5, 2)
+__version__ = version = '1.5.4'
+__version_tuple__ = version_tuple = (1, 5, 4)
 
 __commit_id__ = commit_id = None

{langchain_timbr-1.5.2 → langchain_timbr-1.5.4}/src/langchain_timbr/llm_wrapper/llm_wrapper.py

@@ -12,13 +12,14 @@ class LlmTypes(Enum):
   Google = 'chat-google-generative-ai'
   AzureOpenAI = 'azure-openai-chat'
   Snowflake = 'snowflake-cortex'
+  Databricks = 'chat-databricks'
   Timbr = 'timbr'
 
 
 class LlmWrapper(LLM):
   """
   LlmWrapper is a unified interface for connecting to various Large Language Model (LLM) providers
-  (OpenAI, Anthropic, Google, Azure OpenAI, Snowflake Cortex, etc.) using LangChain. It abstracts
+  (OpenAI, Anthropic, Google, Azure OpenAI, Snowflake Cortex, Databricks, etc.) using LangChain. It abstracts
   the initialization and connection logic for each provider, allowing you to switch between them
   with a consistent API.
   """
@@ -95,12 +96,14 @@ class LlmWrapper(LLM):
         **params,
       )
     elif is_llm_type(llm_type, LlmTypes.Snowflake):
-      from langchain_community.chat_models import ChatSnowflakeCortex      
+      from langchain_community.chat_models import ChatSnowflakeCortex
       llm_model = model or "openai-gpt-4.1"
       params = self._add_temperature(LlmTypes.Snowflake.name, llm_model, **llm_params)
+      snowflake_password = params.pop('snowflake_api_key', params.pop('snowflake_password', api_key))
       
       return ChatSnowflakeCortex(
         model=llm_model,
+        snowflake_password=snowflake_password,
         **params,
       )
     elif is_llm_type(llm_type, LlmTypes.AzureOpenAI):
@@ -116,6 +119,19 @@ class LlmWrapper(LLM):
         openai_api_version=azure_api_version,
         **params,
       )
+    elif is_llm_type(llm_type, LlmTypes.Databricks):
+      from databricks.sdk import WorkspaceClient
+      from databricks_langchain import ChatDatabricks
+      llm_model = model or "databricks-claude-sonnet-4"
+      params = self._add_temperature(LlmTypes.Databricks.name, llm_model, **llm_params)
+
+      host = params.pop('databricks_host', params.pop('host', None))
+      w = WorkspaceClient(host=host, token=api_key)
+      return ChatDatabricks(
+        endpoint=llm_model,
+        workspace_client=w,  # Using authenticated client
+        **params,
+      )
     else:
       raise ValueError(f"Unsupported LLM type: {llm_type}")
 
@@ -163,12 +179,16 @@ class LlmWrapper(LLM):
           "llama3.1-70b",
           "llama3.1-405b"
         ]
+      elif is_llm_type(self._llm_type, LlmTypes.Databricks):
+        w = self.client.workspace_client
+        models = [ep.name for ep in w.serving_endpoints.list()]
+        
       # elif self._is_llm_type(self._llm_type, LlmTypes.Timbr):
         
     except Exception as e:
       models = []
     
-    return models
+    return sorted(models)
 
 
   def _call(self, prompt, **kwargs):

{langchain_timbr-1.5.2 → langchain_timbr-1.5.4}/src/langchain_timbr/utils/general.py

@@ -44,7 +44,8 @@ def is_llm_type(llm_type, enum_value):
             llm_type_lower == enum_name_lower or
             llm_type_lower == enum_value_lower or
             llm_type_lower.startswith(enum_name_lower) or # Usecase for snowflake which its type is the provider name + the model name
-            llm_type_lower.startswith(enum_value_lower)
+            llm_type_lower.startswith(enum_value_lower) or
+            llm_type_lower in enum_value_lower # Check if the enum value includes the llm type - when providing partial name
         )
 
     return False

{langchain_timbr-1.5.2 → langchain_timbr-1.5.4}/src/langchain_timbr/utils/prompt_service.py

@@ -1,32 +1,45 @@
 import requests
 from typing import Dict, Any, Optional, List, Union
-from langchain.schema import SystemMessage, HumanMessage
 from langchain.prompts.chat import ChatPromptTemplate
 from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate
 import json
 import logging
-
-from ..config import url, token as default_token, is_jwt, jwt_tenant_id as default_jwt_tenant_id, llm_timeout
+from ..config import url as default_url, token as default_token, is_jwt, jwt_tenant_id as default_jwt_tenant_id, llm_timeout
 
 logger = logging.getLogger(__name__)
 
 # Global template cache shared across all PromptService instances
-_global_template_cache = {}
+_global_template_cache: dict[Any, Any] = {}
 
 class PromptService:
     def __init__(
         self, 
-        base_url: Optional[str] = url, 
-        token: Optional[str] = default_token,
-        is_jwt: Optional[bool] = is_jwt,
-        jwt_tenant_id: Optional[str] = default_jwt_tenant_id,
-        timeout: Optional[int] = llm_timeout,
+        conn_params: Optional[Dict[str, Any]] = None,
+        **kwargs
     ):
-        self.base_url = base_url.rstrip('/')
-        self.token = token
-        self.is_jwt = is_jwt
-        self.jwt_tenant_id = jwt_tenant_id
-        self.timeout = timeout
+        """
+        Initialize PromptService with connection parameters.
+        
+        Args:
+            conn_params: Dictionary containing connection parameters
+            **kwargs: Additional parameters for backward compatibility
+        """
+        # Extract relevant parameters from conn_params or use defaults
+        if conn_params:
+            url_value = conn_params.get('url') or default_url
+            self.base_url = url_value.rstrip('/') if url_value else ''
+            self.token = conn_params.get('token') or default_token
+            self.is_jwt = conn_params.get('is_jwt', is_jwt)
+            self.jwt_tenant_id = conn_params.get('jwt_tenant_id') or default_jwt_tenant_id
+            self.timeout = conn_params.get('timeout') or llm_timeout
+        else:
+            # Fallback to kwargs for backward compatibility
+            url_value = kwargs.get('url') or default_url
+            self.base_url = url_value.rstrip('/') if url_value else ''
+            self.token = str(kwargs.get('token') or default_token)
+            self.is_jwt = kwargs.get('is_jwt', is_jwt)
+            self.jwt_tenant_id = kwargs.get('jwt_tenant_id') or default_jwt_tenant_id
+            self.timeout = kwargs.get('timeout') or llm_timeout
 
     
     def _get_headers(self) -> Dict[str, str]:
@@ -220,99 +233,67 @@ class PromptTemplateWrapper:
 
 # Individual prompt template getter functions
 def get_determine_concept_prompt_template(
-    token: Optional[str] = None,
-    is_jwt: Optional[bool] = None,
-    jwt_tenant_id: Optional[str] = None
+    conn_params: Optional[dict] = None
 ) -> PromptTemplateWrapper:
     """
     Get determine concept prompt template wrapper
     
     Args:
-        token: Authentication token
-        is_jwt: Whether the token is a JWT
-        jwt_tenant_id: JWT tenant ID
-        
+        conn_params: Connection parameters including url, token, is_jwt, and jwt_tenant_id
+
     Returns:
         PromptTemplateWrapper for determine concept
     """
-    prompt_service = PromptService(
-        token=token,
-        is_jwt=is_jwt,
-        jwt_tenant_id=jwt_tenant_id
-    )
+    prompt_service = PromptService(conn_params=conn_params)
     return PromptTemplateWrapper(prompt_service, "get_identify_concept_template")
 
 
 def get_generate_sql_prompt_template(
-    token: Optional[str] = None,
-    is_jwt: Optional[bool] = None,
-    jwt_tenant_id: Optional[str] = None
+    conn_params: Optional[dict] = None
 ) -> PromptTemplateWrapper:
     """
     Get generate SQL prompt template wrapper
     
     Args:
-        token: Authentication token
-        is_jwt: Whether the token is a JWT
-        jwt_tenant_id: JWT tenant ID
+        conn_params: Connection parameters including url, token, is_jwt, and jwt_tenant_id
         
     Returns:
         PromptTemplateWrapper for generate SQL
     """
-    prompt_service = PromptService(
-        token=token,
-        is_jwt=is_jwt,
-        jwt_tenant_id=jwt_tenant_id
-    )
+    prompt_service = PromptService(conn_params=conn_params)
     return PromptTemplateWrapper(prompt_service, "get_generate_sql_template")
 
 
 def get_qa_prompt_template(
-    token: Optional[str] = None,
-    is_jwt: Optional[bool] = None,
-    jwt_tenant_id: Optional[str] = None
+    conn_params: Optional[dict] = None
 ) -> PromptTemplateWrapper:
     """
     Get QA prompt template wrapper
     
     Args:
-        token: Authentication token
-        is_jwt: Whether the token is a JWT
-        jwt_tenant_id: JWT tenant ID
-        
+        conn_params: Connection parameters including url, token, is_jwt, and jwt_tenant_id
+
     Returns:
         PromptTemplateWrapper for QA
     """
-    prompt_service = PromptService(
-        token=token,
-        is_jwt=is_jwt,
-        jwt_tenant_id=jwt_tenant_id
-    )
+    prompt_service = PromptService(conn_params=conn_params)
     return PromptTemplateWrapper(prompt_service, "get_generate_answer_template")
 
 
 # Global prompt service instance (updated signature)
 def get_prompt_service(
-    token: str = None,
-    is_jwt: bool = None, 
-    jwt_tenant_id: str = None
+    conn_params: Optional[dict] = None
 ) -> PromptService:
     """
     Get or create a prompt service instance
     
     Args:
-        token: Authentication token (API key or JWT token)
-        is_jwt: Whether the token is a JWT
-        jwt_tenant_id: JWT tenant ID
-        
+        conn_params: Connection parameters including url, token, is_jwt, and jwt_tenant_id
+
     Returns:
         PromptService instance
     """
-    return PromptService(
-        token=token,
-        is_jwt=is_jwt,
-        jwt_tenant_id=jwt_tenant_id
-    )
+    return PromptService(conn_params=conn_params)
 
 
 # Global cache management functions

{langchain_timbr-1.5.2 → langchain_timbr-1.5.4}/src/langchain_timbr/utils/temperature_supported_models.json

@@ -2,61 +2,69 @@
   "OpenAI": [
     "gpt-4",
     "gpt-4-turbo",
-    "gpt-4o"
+    "gpt-4o",
+    "gpt-5",
+    "gpt-5-chat-latest",
+    "gpt-5-mini"
   ],
   "Anthropic": [
-    "claude-opus-4-20250514",
-    "claude-sonnet-4-20250514",
-    "claude-3-7-sonnet-20250219",
-    "claude-3-5-sonnet-20241022",
     "claude-3-5-haiku-20241022",
     "claude-3-5-sonnet-20240620",
+    "claude-3-5-sonnet-20241022",
+    "claude-3-7-sonnet-20250219",
     "claude-3-haiku-20240307",
     "claude-3-opus-20240229",
-    "claude-3-sonnet-20240229",
-    "claude-2.1",
-    "claude-2.0"
+    "claude-opus-4-20250514",
+    "claude-sonnet-4-20250514"
   ],
   "Google": [
-    "gemini-1.5-flash-latest",
     "gemini-1.5-flash",
     "gemini-1.5-flash-002",
     "gemini-1.5-flash-8b",
     "gemini-1.5-flash-8b-001",
     "gemini-1.5-flash-8b-latest",
-    "gemini-2.5-flash-preview-04-17",
-    "gemini-2.5-flash-preview-05-20",
-    "gemini-2.5-flash",
-    "gemini-2.5-flash-preview-04-17-thinking",
-    "gemini-2.5-flash-lite-preview-06-17",
-    "gemini-2.5-pro",
-    "gemini-2.0-flash-exp",
+    "gemini-1.5-flash-latest",
     "gemini-2.0-flash",
     "gemini-2.0-flash-001",
+    "gemini-2.0-flash-exp",
     "gemini-2.0-flash-exp-image-generation",
-    "gemini-2.0-flash-lite-001",
     "gemini-2.0-flash-lite",
-    "gemini-2.0-flash-lite-preview-02-05",
+    "gemini-2.0-flash-lite-001",
     "gemini-2.0-flash-lite-preview",
-    "gemini-2.0-flash-thinking-exp-01-21",
+    "gemini-2.0-flash-lite-preview-02-05",
     "gemini-2.0-flash-thinking-exp",
+    "gemini-2.0-flash-thinking-exp-01-21",
     "gemini-2.0-flash-thinking-exp-1219",
-    "learnlm-2.0-flash-experimental",
-    "gemma-3-1b-it",
-    "gemma-3-4b-it",
+    "gemini-2.5-flash",
+    "gemini-2.5-flash-lite",
+    "gemini-2.5-flash-lite-preview-06-17",
+    "gemini-2.5-flash-preview-05-20",
     "gemma-3-12b-it",
+    "gemma-3-1b-it",
     "gemma-3-27b-it",
-    "gemma-3n-e4b-it",
-    "gemma-3n-e2b-it"
+    "gemma-3-4b-it",
+    "gemma-3n-e2b-it",
+    "gemma-3n-e4b-it"
   ],
   "AzureOpenAI": [
     "gpt-4o"
   ],
   "Snowflake": [
-    "openai-gpt-4.1",
-    "mistral-large2",
+    "llama3.1-405b",
     "llama3.1-70b",
-    "llama3.1-405b"
+    "mistral-large2",
+    "openai-gpt-4.1"
+  ],
+  "Databricks": [
+    "databricks-claude-3-7-sonnet",
+    "databricks-claude-sonnet-4",
+    "databricks-gemma-3-12b",
+    "databricks-gpt-oss-120b",
+    "databricks-gpt-oss-20b",
+    "databricks-llama-4-maverick",
+    "databricks-meta-llama-3-1-405b-instruct",
+    "databricks-meta-llama-3-1-8b-instruct",
+    "databricks-meta-llama-3-3-70b-instruct"
   ],
   "Timbr": []
 }

{langchain_timbr-1.5.2 → langchain_timbr-1.5.4}/src/langchain_timbr/utils/timbr_llm_utils.py

@@ -165,6 +165,40 @@ def _calculate_token_count(llm: LLM, prompt: str) -> int:
     return token_count
     
 
+def _get_response_text(response: Any) -> str:
+    if hasattr(response, "content"):
+        response_text = response.content
+
+        # Handle Databricks gpt-oss type of responses (having list of dicts with type + summary for reasoning or type + text for result)
+        if isinstance(response_text, list):
+            response_text = next(filter(lambda x: x.get('type') == 'text', response.content), None)
+        if isinstance(response_text, dict):
+            response_text = response_text.get('text', '')
+    elif isinstance(response, str):
+        response_text = response
+    else:
+        raise ValueError("Unexpected response format from LLM.")
+
+    return response_text
+
+def _extract_usage_metadata(response: Any) -> dict:
+    usage_metadata = response.response_metadata
+
+    if usage_metadata and 'usage' in usage_metadata:
+        usage_metadata = usage_metadata['usage']
+
+    if not usage_metadata and 'usage_metadata' in response:
+        usage_metadata = response.usage_metadata
+        if usage_metadata and 'usage' in usage_metadata:
+            usage_metadata = usage_metadata['usage']
+
+    if not usage_metadata and 'usage' in response:
+        usage_metadata = response.usage
+        if usage_metadata and 'usage' in usage_metadata:
+            usage_metadata = usage_metadata['usage']
+
+    return usage_metadata
+
 def determine_concept(
     question: str,
     llm: LLM,
@@ -178,7 +212,7 @@ def determine_concept(
     note: Optional[str] = '',
     debug: Optional[bool] = False,
     timeout: Optional[int] = None,
-) -> dict[str, any]:
+) -> dict[str, Any]:
     usage_metadata = {}
     determined_concept_name = None
     schema = 'dtimbr'
@@ -187,7 +221,7 @@ def determine_concept(
     if timeout is None:
         timeout = llm_timeout
     
-    determine_concept_prompt = get_determine_concept_prompt_template(conn_params["token"], conn_params["is_jwt"], conn_params["jwt_tenant_id"])
+    determine_concept_prompt = get_determine_concept_prompt_template(conn_params)
     tags = get_tags(conn_params=conn_params, include_tags=include_tags)
     concepts = get_concepts(
         conn_params=conn_params,
@@ -253,20 +287,12 @@ def determine_concept(
                 continue
             usage_metadata['determine_concept'] = {
                 "approximate": apx_token_count,
-                # **(response.usage_metadata or response.usage or {}),
-                **(response.usage_metadata or {}),
+                **_extract_usage_metadata(response),
             }
             if debug:
                 usage_metadata['determine_concept']["p_hash"] = encrypt_prompt(prompt)
 
-            if hasattr(response, "content"):
-                response_text = response.content
-            elif isinstance(response, str):
-                response_text = response
-            else:
-                raise ValueError("Unexpected response format from LLM.")
-
-
+            response_text = _get_response_text(response)
             candidate = response_text.strip()
             if should_validate and candidate not in concepts.keys():
                 error = f"Concept '{determined_concept_name}' not found in the list of concepts."
@@ -351,13 +377,7 @@ def _build_rel_columns_str(relationships: list[dict], columns_tags: Optional[dic
 
 
 def _parse_sql_from_llm_response(response: Any) -> str:
-    if hasattr(response, "content"):
-        response_text = response.content
-    elif isinstance(response, str):
-        response_text = response
-    else:
-        raise ValueError("Unexpected response format from LLM.")
-    
+    response_text = _get_response_text(response)
     return (response_text
             .replace("```sql", "")
             .replace("```", "")
@@ -398,7 +418,7 @@ def generate_sql(
     if timeout is None:
         timeout = llm_timeout
     
-    generate_sql_prompt = get_generate_sql_prompt_template(conn_params["token"], conn_params["is_jwt"], conn_params["jwt_tenant_id"])
+    generate_sql_prompt = get_generate_sql_prompt_template(conn_params)
    
     if concept and concept != "" and (schema is None or schema != "vtimbr"):
         concepts_list = [concept]
@@ -497,8 +517,7 @@ def generate_sql(
 
         usage_metadata['generate_sql'] = {
             "approximate": apx_token_count,
-            # **(response.usage_metadata or response.usage or {}),
-            **(response.usage_metadata or {}),
+            **_extract_usage_metadata(response),
         }
         if debug:
             usage_metadata['generate_sql']["p_hash"] = encrypt_prompt(prompt)
@@ -531,7 +550,7 @@ def answer_question(
     if timeout is None:
         timeout = llm_timeout
 
-    qa_prompt = get_qa_prompt_template(conn_params["token"], conn_params["is_jwt"], conn_params["jwt_tenant_id"])
+    qa_prompt = get_qa_prompt_template(conn_params)
 
     prompt = qa_prompt.format_messages(
         question=question,
@@ -561,8 +580,7 @@ def answer_question(
     usage_metadata = {
         "answer_question": {
             "approximate": apx_token_count,
-            # **(response.usage_metadata or response.usage or {}),
-            **(response.usage_metadata or {}),
+            **_extract_usage_metadata(response),
         },
     }
     if debug:

langchain_timbr-1.5.4/tests/integration/test_azure_databricks_provider.py

@@ -0,0 +1,42 @@
+from langchain_timbr import LlmWrapper, LlmTypes, ExecuteTimbrQueryChain
+import os
+
+class TestAzureDatabricksProvider:
+    """Test suite for Azure Databricks provider integration with Timbr chains."""
+    
+    def skip_test_databricks_connection(self, llm, config):
+        DATABRICKS_LLM_HOST = os.getenv('databricks_host')
+        AZURE_TOKEN = os.getenv('databricks_token')
+        
+        llm_models_test = [
+          "databricks-claude-sonnet-4",
+          "databricks-gpt-oss-20b",
+        ]
+
+        for llm_model in llm_models_test:
+          llm_instance = LlmWrapper(
+              llm_type=LlmTypes.Databricks,
+              api_key=AZURE_TOKEN,
+              databricks_host=DATABRICKS_LLM_HOST,
+              model=llm_model,
+          )
+
+          chain = ExecuteTimbrQueryChain(
+              llm=llm_instance,
+              url=config["timbr_url"],
+              token=config["timbr_token"],
+              ontology=config["timbr_ontology"],
+              verify_ssl=config["verify_ssl"],
+          )
+
+          inputs = {
+              "prompt": config["test_prompt"],
+          }
+          result = chain.invoke(inputs)
+
+          print("ExecuteTimbrQueryChain result:", result)
+          assert "rows" in result, "Result should contain 'rows'"
+          assert isinstance(result["rows"], list), "'rows' should be a list"
+          assert result["sql"], "SQL should be present in the result"
+
+