|
15 | 15 | from api.agents import AnalysisAgent, RelevancyAgent, ResponseFormatterAgent, FollowUpAgent |
16 | 16 | from api.agents.healer_agent import HealerAgent |
17 | 17 | from api.config import Config |
| 18 | +from api.config import SUPPORTED_VENDORS |
18 | 19 | from api.extensions import db |
19 | 20 | from api.graph import find, get_db_description, get_user_rules |
20 | 21 | from api.loaders.postgres_loader import PostgresLoader |
21 | 22 | from api.loaders.mysql_loader import MySQLLoader |
| 23 | +from api.loaders.snowflake_loader import SnowflakeLoader |
22 | 24 | from api.memory.graphiti_tool import MemoryTool |
23 | 25 | from api.sql_utils import SQLIdentifierQuoter, DatabaseSpecificQuoter |
24 | 26 |
|
@@ -83,6 +85,8 @@ def get_database_type_and_loader(db_url: str): |
83 | 85 | return 'postgresql', PostgresLoader |
84 | 86 | if db_url_lower.startswith('mysql://'): |
85 | 87 | return 'mysql', MySQLLoader |
| 88 | + if db_url_lower.startswith('snowflake://'): |
| 89 | + return 'snowflake', SnowflakeLoader |
86 | 90 |
|
87 | 91 | # Default to PostgresLoader for backward compatibility |
88 | 92 | return 'postgresql', PostgresLoader |
@@ -257,21 +261,17 @@ async def generate(): # pylint: disable=too-many-locals,too-many-branches,too-m |
257 | 261 | custom_model = chat_data.custom_model |
258 | 262 |
|
259 | 263 | # Validate custom model format (vendor/model) |
260 | | - supported_vendors = ("openai", "anthropic", "gemini", "azure", "ollama", "cohere") |
261 | 264 | if custom_model: |
262 | 265 | parts = custom_model.split("/", 1) |
263 | 266 | if len(parts) != 2 or not parts[0] or not parts[1]: |
264 | 267 | raise InvalidArgumentError( |
265 | 268 | "Invalid model format. Expected 'vendor/model' (e.g. 'openai/gpt-4.1')" |
266 | 269 | ) |
267 | | - if parts[0] not in supported_vendors: |
| 270 | + if parts[0] not in SUPPORTED_VENDORS: |
268 | 271 | raise InvalidArgumentError( |
269 | | - f"Unsupported vendor '{parts[0]}'. Supported: {', '.join(supported_vendors)}" |
| 272 | + f"Unsupported vendor '{parts[0]}'. Supported: {', '.join(SUPPORTED_VENDORS)}" |
270 | 273 | ) |
271 | 274 |
|
272 | | - if custom_api_key is not None and len(custom_api_key.strip()) < 10: |
273 | | - raise InvalidArgumentError("API key is too short") |
274 | | - |
275 | 275 | agent_rel = RelevancyAgent(queries_history, result_history, custom_api_key, custom_model) |
276 | 276 | agent_an = AnalysisAgent(queries_history, result_history, custom_api_key, custom_model) |
277 | 277 | follow_up_agent = FollowUpAgent(queries_history, result_history, custom_api_key, custom_model) |
|
0 commit comments