Skip to content

Commit ab76312

Browse files
remi-tdremi-td
authored
Merge pull request #243 from asmirnov-tba/Chat-Completion-tools
New tools for Chat Completion CE in-db function (issue #242)
2 parents eb5bb1c + b0a8cab commit ab76312

File tree

10 files changed

+1019
-1
lines changed

10 files changed

+1019
-1
lines changed

‎src/teradata_mcp_server/app.py‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ def create_mcp_app(settings: Settings):
6464
enableEFS = True if any(re.match(pattern, 'fs_*') for pattern in config.get('tool', [])) else False
6565
enableTDVS = True if any(re.match(pattern, 'tdvs_*') for pattern in config.get('tool', [])) else False
6666
enableBAR = True if any(re.match(pattern, 'bar_*') for pattern in config.get('tool', [])) else False
67+
enableChatCmplt = True if any(re.match(pattern, 'chat_cmplt_*') for pattern in config.get('tool', [])) else False
6768

6869
# Initialize TD connection and optional teradataml/EFS context
6970
# Pass settings object to TDConn instead of just connection_url
@@ -128,6 +129,116 @@ def create_mcp_app(settings: Settings):
128129
logger.warning(f"BAR system dependencies not available - disabling BAR functionality: {e}")
129130
enableBAR = False
130131

132+
# Chat Completion module validation (optional)
133+
if enableChatCmplt:
134+
try:
135+
from teradata_mcp_server.tools.chat_cmplt.chat_cmplt_tools import load_chat_cmplt_config
136+
137+
# Test 1: Check if base_url and model are set in chat_cmplt_config.yml
138+
chat_cmplt_config = load_chat_cmplt_config()
139+
base_url = chat_cmplt_config.get("base_url", "").strip()
140+
model = chat_cmplt_config.get("model", "").strip()
141+
function_db = chat_cmplt_config.get("databases", {}).get("function_db", "").strip()
142+
143+
if not base_url or not model:
144+
logger.warning(
145+
f"Chat completion config missing required parameters "
146+
f"(base_url: {'set' if base_url else 'not set'}, "
147+
f"model: {'set' if model else 'not set'}) - "
148+
f"disabling chat completion functionality"
149+
)
150+
enableChatCmplt = False
151+
elif not function_db:
152+
logger.warning(
153+
"Chat completion config missing function database "
154+
"(databases.function_db not set) - disabling chat completion functionality"
155+
)
156+
enableChatCmplt = False
157+
else:
158+
# Tests 2 & 3: Check database function existence and permissions
159+
# Only perform these if we can establish a connection
160+
try:
161+
# Check if connection is available
162+
if not getattr(tdconn, "engine", None):
163+
logger.info(
164+
f"Chat completion module config validated (base_url, model, function_db set). "
165+
f"Database checks (function existence and permissions) will be skipped in stdio mode - "
166+
f"they will be validated on first tool use."
167+
)
168+
else:
169+
with tdconn.engine.connect() as conn:
170+
from sqlalchemy import text
171+
172+
# Test 2: Check if CompleteChat function exists in configured database
173+
check_function_sql = text(f"""
174+
SELECT 1
175+
FROM DBC.FunctionsV
176+
WHERE DatabaseName = '{function_db}'
177+
AND FunctionName = 'CompleteChat'
178+
""")
179+
result = conn.execute(check_function_sql)
180+
function_exists = result.fetchone() is not None
181+
182+
if not function_exists:
183+
logger.warning(
184+
f"CompleteChat function not found in database '{function_db}' - "
185+
f"disabling chat completion functionality"
186+
)
187+
enableChatCmplt = False
188+
else:
189+
# Test 3: Check if current user has execute permission on CompleteChat
190+
# This includes: direct function grants, database-level grants, and role-based grants
191+
192+
# First, get current username
193+
username_result = conn.execute(text("SELECT USER"))
194+
current_user = username_result.fetchone()[0]
195+
196+
check_permission_sql = text(f"""
197+
SELECT 1
198+
FROM DBC.AllRightsV
199+
WHERE UPPER(UserName) = UPPER('{current_user}')
200+
AND UPPER(DatabaseName) = UPPER('{function_db}')
201+
AND (
202+
-- Case 1: Direct grant on the function itself
203+
(UPPER(TableName) = UPPER('CompleteChat') AND AccessRight = 'EF')
204+
OR
205+
-- Case 2: Database-level execute function grant
206+
(TableName = 'All' AND AccessRight = 'EF')
207+
)
208+
""")
209+
result = conn.execute(check_permission_sql)
210+
has_permission = result.fetchone() is not None
211+
212+
if not has_permission:
213+
logger.warning(
214+
f"User '{current_user}' does not have EXECUTE FUNCTION permission "
215+
f"on {function_db}.CompleteChat (checked direct grants, database-level grants, and role-based grants) - "
216+
f"disabling chat completion functionality"
217+
)
218+
enableChatCmplt = False
219+
else:
220+
logger.info(
221+
f"Chat completion module validated successfully "
222+
f"(user: {current_user}, base_url: {base_url[:30]}..., model: {model}, "
223+
f"function: {function_db}.CompleteChat)"
224+
)
225+
except (AttributeError, Exception) as db_error:
226+
# In stdio mode, connection might not be available at startup
227+
# Log info instead of warning and allow tools to load
228+
# They will fail at runtime if there are actual permission issues
229+
logger.info(
230+
f"Chat completion config validated (base_url, model, function_db set). "
231+
f"Database validation skipped (connection not available at startup): {db_error}. "
232+
f"Function existence and permissions will be validated on first tool use."
233+
)
234+
235+
except (AttributeError, ImportError, ModuleNotFoundError) as e:
236+
logger.warning(f"Chat completion module not available - disabling chat completion functionality: {e}")
237+
enableChatCmplt = False
238+
except Exception as e:
239+
logger.warning(f"Error loading chat completion config - disabling chat completion functionality: {e}")
240+
enableChatCmplt = False
241+
131242
# Middleware (auth + request context)
132243
from teradata_mcp_server.tools.auth_cache import SecureAuthCache
133244
auth_cache = SecureAuthCache(ttl_seconds=settings.auth_cache_ttl)
@@ -358,6 +469,10 @@ def executor(**kwargs):
358469
if tool_name.startswith("bar_") and not enableBAR:
359470
logger.info(f"Skipping BAR tool: {tool_name} (BAR functionality disabled)")
360471
continue
472+
# Skip chat completion tools if chat completion functionality is disabled
473+
if tool_name.startswith("chat_cmplt_") and not enableChatCmplt:
474+
logger.info(f"Skipping chat completion tool: {tool_name} (chat completion functionality disabled)")
475+
continue
361476
wrapped = make_tool_wrapper(func)
362477
mcp.tool(name=tool_name, description=wrapped.__doc__)(wrapped)
363478
logger.info(f"Created tool: {tool_name}")

‎src/teradata_mcp_server/config/chat_cmplt_config.yml‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
# Configuration for Chat Completion Tools
2+
# Used to connect Teradata with OpenAI-compatible LLM inference servers
3+
4+
# Connection Configuration
5+
# base_url (REQUIRED): Base URL of the inference server
6+
# Will log error and disable the tool if not set
7+
# Example: "https://api.openai.com"
8+
base_url: ""
9+
10+
# Model Configuration (REQUIRED)
11+
# model: The model name/ID to use for chat completion
12+
# Example: "qwen2.5-coder:7b", "gpt-4", "claude-3-opus"
13+
model: "" # REQUIRED - will log error and disable tool if not set
14+
15+
# TLS/HTTPS Configuration
16+
# IgnoreHTTPSVerification: Skip TLS certificate verification (dev/test only)
17+
# HTTP proxy configuration
18+
# http_proxy: Optional HTTP proxy in "host:port" format (e.g. "proxy.example.com:8080")
19+
# If omitted or empty, no proxy is used.
20+
http_proxy: ""
21+
22+
# TLS/HTTPS Configuration
23+
# ignore_https_verification: Skip TLS certificate verification (dev/test only)
24+
# Default: false
25+
ignore_https_verification: false
26+
27+
# HTTP Headers Configuration
28+
# custom_headers: List of custom HTTP headers to add to requests
29+
# Format: key-value pairs
30+
# Example:
31+
# - key: "x-custom-header"
32+
# value: "custom-value"
33+
custom_headers: []
34+
35+
# Request Body Parameters
36+
# body_parameters: Additional parameters to include in JSON request body
37+
# Values can be JSON objects, arrays, strings, numbers, or booleans
38+
# Format: key-value pairs
39+
# Example:
40+
# - key: "temperature"
41+
# value: "0.1"
42+
# - key: "max_tokens"
43+
# value: "100"
44+
# - key: "response_format"
45+
# value: '{"type":"json_object"}'
46+
body_parameters: []
47+
48+
# Optional PEM file selector SQL
49+
# pem_file_sql: SQL statement to select PEM file for the second input (TLS cert table).
50+
# If empty or not provided, it will not be added to generated SQL statements.
51+
# Example:
52+
# pem_file_sql: "SELECT cert FROM mydb.pem_files WHERE id = 'cert'"
53+
pem_file_sql: ""
54+
55+
# Rate Limiting Configuration
56+
# delays: Comma-separated list of retry delays in milliseconds for HTTP 429
57+
# Default: "500"
58+
delays: "500"
59+
60+
# retries_number: Number of retries after hitting rate limit (HTTP 429)
61+
# Default: 0
62+
retries_number: 0
63+
64+
# throw_error_on_rate_limit: Whether to throw error after retries exhausted
65+
# If false, returns NULL/empty response with diagnostics
66+
# Default: false
67+
throw_error_on_rate_limit: false
68+
69+
# Output Configuration
70+
# output_text_length: Maximum length of output text field (VARCHAR size)
71+
# Range: 2-32000
72+
# Default: 16000
73+
output_text_length: 16000
74+
75+
# remove_deepseek_thinking: Strip <think>...</think> blocks from DeepSeek responses
76+
# Default: false
77+
remove_deepseek_thinking: false
78+
79+
# Database Configuration
80+
databases:
81+
# function_db (REQUIRED) - Database for CompleteChat function
82+
function_db: "openai_client"
83+
84+
# Output Options
85+
output:
86+
# include_diagnostics: Include diagnostic columns (retries_made, last_attempt_duration, etc.)
87+
include_diagnostics: false
88+
# include_tachyon_headers: Include Tachyon call-level headers (x_request_id, x_correlation_id, etc.)
89+
include_tachyon_headers: false

‎src/teradata_mcp_server/config/profiles.yml‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,4 +53,11 @@ bar:
5353
prompt:
5454
- ^bar_*
5555
resource:
56-
- ^bar_*
56+
- ^bar_*
57+
58+
llmUser:
59+
tool:
60+
- ^base_*
61+
- ^chat_cmplt_*
62+
prompt:
63+
- ^chat_cmplt_*

0 commit comments

Comments
 (0)