sql-assistant 1.0.0__py3-none-any.whl

sql_assistant/llm/providers/claude.py

@@ -0,0 +1,132 @@
+"""Anthropic Claude Provider"""
+
+import httpx
+from typing import AsyncGenerator, Optional
+
+from ..base import BaseLLMProvider
+from ..retry import async_retry
+from ..exceptions import LLMConnectionError, LLMResponseError
+
+
+class ClaudeProvider(BaseLLMProvider):
+    """Anthropic Claude API 实现"""
+
+    def __init__(self, api_key: str, base_url: str, model: str):
+        super().__init__(api_key, base_url, model)
+        self._client: Optional[httpx.AsyncClient] = None
+
+    async def _get_client(self) -> httpx.AsyncClient:
+        if self._client is None:
+            self._client = httpx.AsyncClient(
+                timeout=httpx.Timeout(60.0),
+                headers={
+                    "x-api-key": self.api_key,
+                    "anthropic-version": "2023-06-01",
+                    "Content-Type": "application/json",
+                },
+            )
+        return self._client
+
+    def _convert_messages(self, messages: list[dict]) -> tuple[list[dict], Optional[str]]:
+        """将 OpenAI 格式消息转为 Claude 格式"""
+        system_content = None
+        claude_messages = []
+
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+
+            if role == "system":
+                system_content = content
+            elif role == "user":
+                claude_messages.append({"role": "user", "content": content})
+            elif role == "assistant":
+                claude_messages.append({"role": "assistant", "content": content})
+
+        return claude_messages, system_content
+
+    @async_retry(max_attempts=3, base_delay=1.0, retryable_exceptions=(httpx.HTTPError,))
+    async def chat(self, messages: list[dict], temperature: float = 0.1) -> str:
+        client = await self._get_client()
+        claude_messages, system = self._convert_messages(messages)
+
+        url = "https://api.anthropic.com/v1/messages"
+        payload = {
+            "model": self.model,
+            "max_tokens": 4096,
+            "temperature": temperature,
+            "messages": claude_messages,
+        }
+        if system:
+            payload["system"] = system
+
+        response = await client.post(url, json=payload)
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            raise LLMConnectionError(f"Claude 请求失败: {e.response.status_code} - {e.response.text}", "claude")
+
+        data = response.json()
+        try:
+            return data["content"][0]["text"]
+        except (KeyError, IndexError) as e:
+            raise LLMResponseError(f"Claude 响应格式错误: {e}", "claude", data)
+
+    async def chat_stream(self, messages: list[dict], temperature: float = 0.1) -> AsyncGenerator[str, None]:
+        client = await self._get_client()
+        claude_messages, system = self._convert_messages(messages)
+
+        url = "https://api.anthropic.com/v1/messages"
+        payload = {
+            "model": self.model,
+            "max_tokens": 4096,
+            "temperature": temperature,
+            "messages": claude_messages,
+            "stream": True,
+        }
+        if system:
+            payload["system"] = system
+
+        async with client.stream("POST", url, json=payload) as response:
+            try:
+                response.raise_for_status()
+            except httpx.HTTPStatusError as e:
+                raise LLMConnectionError(f"Claude 流式请求失败: {e.response.status_code}", "claude")
+
+            async for line in response.aiter_lines():
+                if line.startswith("data: "):
+                    data_str = line[6:]
+                    try:
+                        import json
+                        data = json.loads(data_str)
+                        if data.get("type") == "content_block_delta":
+                            delta = data.get("delta", {})
+                            text = delta.get("text", "")
+                            if text:
+                                yield text
+                    except Exception:
+                        continue
+
+    async def test_connection(self) -> dict:
+        from ..exceptions import format_llm_result
+        try:
+            client = await self._get_client()
+            url = "https://api.anthropic.com/v1/messages"
+            payload = {
+                "model": self.model,
+                "max_tokens": 1,
+                "temperature": 0,
+                "messages": [{"role": "user", "content": "Hello"}],
+            }
+            response = await client.post(url, json=payload, timeout=30.0)
+            if response.status_code == 200:
+                return format_llm_result(True, data={"message": "Claude 连接测试成功"})
+            else:
+                return format_llm_result(False, error=f"连接失败: {response.status_code}", provider="claude", code="HTTP_ERROR")
+        except Exception as e:
+            return format_llm_result(False, error=f"连接失败: {str(e)}", provider="claude", code="CONNECTION_ERROR")
+
+    async def close(self):
+        if self._client:
+            await self._client.aclose()
+            self._client = None

sql_assistant/llm/providers/gemini.py

@@ -0,0 +1,127 @@
+"""Google Gemini Provider"""
+
+import httpx
+from typing import AsyncGenerator, Optional
+
+from ..base import BaseLLMProvider
+from ..retry import async_retry
+from ..exceptions import LLMConnectionError, LLMResponseError
+
+
+class GeminiProvider(BaseLLMProvider):
+    """Google Gemini API 实现"""
+
+    def __init__(self, api_key: str, base_url: str, model: str):
+        super().__init__(api_key, base_url, model)
+        self._client: Optional[httpx.AsyncClient] = None
+
+    async def _get_client(self) -> httpx.AsyncClient:
+        if self._client is None:
+            self._client = httpx.AsyncClient(timeout=httpx.Timeout(60.0))
+        return self._client
+
+    def _convert_messages(self, messages: list[dict]) -> tuple[list[dict], Optional[str]]:
+        """将 OpenAI 格式消息转为 Gemini 格式"""
+        gemini_contents = []
+        system_instruction = None
+
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+
+            if role == "system":
+                system_instruction = content
+            elif role == "user":
+                gemini_contents.append({"role": "user", "parts": [{"text": content}]})
+            elif role == "assistant":
+                gemini_contents.append({"role": "model", "parts": [{"text": content}]})
+
+        return gemini_contents, system_instruction
+
+    @async_retry(max_attempts=3, base_delay=1.0, retryable_exceptions=(httpx.HTTPError,))
+    async def chat(self, messages: list[dict], temperature: float = 0.1) -> str:
+        client = await self._get_client()
+        contents, system_instruction = self._convert_messages(messages)
+
+        url = f"https://generativelanguage.googleapis.com/v1beta/models/{self.model}:generateContent"
+        params = {"key": self.api_key}
+
+        payload = {
+            "contents": contents,
+            "generationConfig": {"temperature": temperature},
+        }
+        if system_instruction:
+            payload["systemInstruction"] = {
+                "parts": [{"text": system_instruction}]
+            }
+
+        response = await client.post(url, params=params, json=payload)
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            raise LLMConnectionError(f"Gemini 请求失败: {e.response.status_code} - {e.response.text}", "gemini")
+
+        data = response.json()
+        try:
+            return data["candidates"][0]["content"]["parts"][0]["text"]
+        except (KeyError, IndexError) as e:
+            raise LLMResponseError(f"Gemini 响应格式错误: {e}", "gemini", data)
+
+    async def chat_stream(self, messages: list[dict], temperature: float = 0.1) -> AsyncGenerator[str, None]:
+        client = await self._get_client()
+        contents, system_instruction = self._convert_messages(messages)
+
+        url = f"https://generativelanguage.googleapis.com/v1beta/models/{self.model}:streamGenerateContent"
+        params = {"key": self.api_key, "alt": "sse"}
+
+        payload = {
+            "contents": contents,
+            "generationConfig": {"temperature": temperature},
+        }
+        if system_instruction:
+            payload["systemInstruction"] = {
+                "parts": [{"text": system_instruction}]
+            }
+
+        async with client.stream("POST", url, params=params, json=payload) as response:
+            try:
+                response.raise_for_status()
+            except httpx.HTTPStatusError as e:
+                raise LLMConnectionError(f"Gemini 流式请求失败: {e.response.status_code}", "gemini")
+
+            async for line in response.aiter_lines():
+                if line.startswith("data: "):
+                    data_str = line[6:]
+                    try:
+                        import json
+                        data = json.loads(data_str)
+                        parts = data.get("candidates", [{}])[0].get("content", {}).get("parts", [])
+                        for part in parts:
+                            text = part.get("text", "")
+                            if text:
+                                yield text
+                    except Exception:
+                        continue
+
+    async def test_connection(self) -> dict:
+        from ..exceptions import format_llm_result
+        try:
+            client = await self._get_client()
+            url = f"https://generativelanguage.googleapis.com/v1beta/models/{self.model}:generateContent"
+            params = {"key": self.api_key}
+            payload = {
+                "contents": [{"role": "user", "parts": [{"text": "Hello"}]}],
+                "generationConfig": {"temperature": 0, "maxOutputTokens": 1},
+            }
+            response = await client.post(url, params=params, json=payload, timeout=30.0)
+            if response.status_code == 200:
+                return format_llm_result(True, data={"message": "Gemini 连接测试成功"})
+            else:
+                return format_llm_result(False, error=f"连接失败: {response.status_code}", provider="gemini", code="HTTP_ERROR")
+        except Exception as e:
+            return format_llm_result(False, error=f"连接失败: {str(e)}", provider="gemini", code="CONNECTION_ERROR")
+
+    async def close(self):
+        if self._client:
+            await self._client.aclose()
+            self._client = None

sql_assistant/llm/providers/openai_compatible.py

@@ -0,0 +1,103 @@
+"""OpenAI 兼容接口 Provider - 支持 DeepSeek / Doubao / Kimi / Qwen / OpenAI"""
+
+import httpx
+from typing import AsyncGenerator, Optional
+
+from ..base import BaseLLMProvider
+from ..retry import async_retry
+from ..exceptions import LLMConnectionError, LLMResponseError
+
+
+class OpenAICompatibleProvider(BaseLLMProvider):
+    """OpenAI 兼容接口实现"""
+
+    def __init__(self, api_key: str, base_url: str, model: str, provider_name: str = ""):
+        super().__init__(api_key, base_url, model)
+        self.provider_name = provider_name
+        self._client: Optional[httpx.AsyncClient] = None
+
+    async def _get_client(self) -> httpx.AsyncClient:
+        if self._client is None:
+            self._client = httpx.AsyncClient(
+                timeout=httpx.Timeout(60.0),
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                },
+            )
+        return self._client
+
+    @async_retry(max_attempts=3, base_delay=1.0, retryable_exceptions=(httpx.HTTPError,))
+    async def chat(self, messages: list[dict], temperature: float = 0.1) -> str:
+        client = await self._get_client()
+        url = f"{self.base_url}/chat/completions"
+        payload = {
+            "model": self.model,
+            "messages": messages,
+            "temperature": temperature,
+        }
+        response = await client.post(url, json=payload)
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            raise LLMConnectionError(f"LLM 请求失败: {e.response.status_code} - {e.response.text}", self.provider_name)
+
+        data = response.json()
+        try:
+            return data["choices"][0]["message"]["content"]
+        except (KeyError, IndexError) as e:
+            raise LLMResponseError(f"LLM 响应格式错误: {e}", self.provider_name, data)
+
+    async def chat_stream(self, messages: list[dict], temperature: float = 0.1) -> AsyncGenerator[str, None]:
+        client = await self._get_client()
+        url = f"{self.base_url}/chat/completions"
+        payload = {
+            "model": self.model,
+            "messages": messages,
+            "temperature": temperature,
+            "stream": True,
+        }
+        async with client.stream("POST", url, json=payload) as response:
+            try:
+                response.raise_for_status()
+            except httpx.HTTPStatusError as e:
+                raise LLMConnectionError(f"LLM 流式请求失败: {e.response.status_code}", self.provider_name)
+
+            async for line in response.aiter_lines():
+                if line.startswith("data: "):
+                    data_str = line[6:]
+                    if data_str == "[DONE]":
+                        break
+                    try:
+                        import json
+                        data = json.loads(data_str)
+                        delta = data["choices"][0].get("delta", {})
+                        content = delta.get("content", "")
+                        if content:
+                            yield content
+                    except Exception:
+                        continue
+
+    async def test_connection(self) -> dict:
+        from ..exceptions import format_llm_result
+        try:
+            client = await self._get_client()
+            url = f"{self.base_url}/chat/completions"
+            payload = {
+                "model": self.model,
+                "messages": [{"role": "user", "content": "Hello"}],
+                "temperature": 0,
+                "max_tokens": 1,
+            }
+            response = await client.post(url, json=payload, timeout=30.0)
+            if response.status_code == 200:
+                return format_llm_result(True, data={"message": "LLM 连接测试成功"})
+            else:
+                return format_llm_result(False, error=f"连接失败: {response.status_code}", provider=self.provider_name, code="HTTP_ERROR")
+        except Exception as e:
+            return format_llm_result(False, error=f"连接失败: {str(e)}", provider=self.provider_name, code="CONNECTION_ERROR")
+
+    async def close(self):
+        if self._client:
+            await self._client.aclose()
+            self._client = None

sql_assistant/llm/retry.py

@@ -0,0 +1,88 @@
+"""LLM 请求重试机制"""
+
+import asyncio
+import random
+from functools import wraps
+from typing import Callable, TypeVar, Any
+
+T = TypeVar("T")
+
+
+class RetryError(Exception):
+    """重试耗尽后的异常"""
+    def __init__(self, message: str, last_exception: Exception):
+        super().__init__(message)
+        self.last_exception = last_exception
+
+
+def async_retry(
+    max_attempts: int = 3,
+    base_delay: float = 1.0,
+    max_delay: float = 10.0,
+    exponential_base: float = 2.0,
+    jitter: bool = True,
+    retryable_exceptions: tuple = (Exception,),
+):
+    """异步重试装饰器
+
+    Args:
+        max_attempts: 最大尝试次数
+        base_delay: 基础延迟时间（秒）
+        max_delay: 最大延迟时间（秒）
+        exponential_base: 指数退避基数
+        jitter: 是否添加随机抖动
+        retryable_exceptions: 可重试的异常类型元组
+    """
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        @wraps(func)
+        async def wrapper(*args, **kwargs) -> Any:
+            last_exception = None
+
+            for attempt in range(1, max_attempts + 1):
+                try:
+                    return await func(*args, **kwargs)
+                except retryable_exceptions as e:
+                    last_exception = e
+
+                    if attempt == max_attempts:
+                        break
+
+                    delay = min(base_delay * (exponential_base ** (attempt - 1)), max_delay)
+
+                    if jitter:
+                        delay = delay * (0.5 + random.random() * 0.5)
+
+                    await asyncio.sleep(delay)
+
+            raise RetryError(
+                f"重试 {max_attempts} 次后仍然失败: {last_exception}",
+                last_exception
+            )
+
+        return wrapper
+    return decorator
+
+
+class RateLimiter:
+    """简单的速率限制器"""
+
+    def __init__(self, calls: int, period: float):
+        self.calls = calls
+        self.period = period
+        self._tokens = calls
+        self._last_update = asyncio.get_event_loop().time()
+
+    async def acquire(self) -> None:
+        loop = asyncio.get_event_loop()
+        now = loop.time()
+
+        elapsed = now - self._last_update
+        self._tokens = min(self.calls, self._tokens + elapsed * (self.calls / self.period))
+        self._last_update = now
+
+        if self._tokens < 1:
+            wait_time = (1 - self._tokens) * (self.period / self.calls)
+            await asyncio.sleep(wait_time)
+            self._tokens = 0
+        else:
+            self._tokens -= 1

sql_assistant/main.py

@@ -0,0 +1,94 @@
+"""SQL 智能助手 - 主入口"""
+
+import sys
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+from fastapi import FastAPI, Request
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
+from fastapi.middleware.cors import CORSMiddleware
+
+from .api.routes import router as api_router
+from .database.history import close_history_manager
+from .llm.manager import get_llm_manager
+from .database.manager import get_db_manager
+
+# 模板和静态文件路径
+BASE_DIR = Path(__file__).parent
+TEMPLATES_DIR = BASE_DIR / "web" / "templates"
+STATIC_DIR = BASE_DIR / "web" / "static"
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """应用生命周期管理"""
+    yield
+    await get_llm_manager().close_all()
+    await get_db_manager().close_all()
+    await close_history_manager()
+
+
+app = FastAPI(
+    title="SQL 智能助手",
+    description="自然语言转 SQL 查询工具",
+    version="1.0.0",
+    lifespan=lifespan,
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# 静态文件
+if STATIC_DIR.exists():
+    app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
+
+# Jinja2 模板
+templates = Jinja2Templates(directory=str(TEMPLATES_DIR))
+
+# API 路由
+app.include_router(api_router)
+
+
+@app.get("/")
+async def index(request: Request):
+    return templates.TemplateResponse(request, "index.html", {"request": request})
+
+
+def main():
+    """CLI 入口"""
+    import uvicorn
+
+    # 确保控制台支持 UTF-8
+    try:
+        sys.stdout.reconfigure(encoding='utf-8')
+    except Exception:
+        pass
+
+    print("=" * 60)
+    print("  SQL Assistant v1.0.0")
+    print("  Natural Language -> SQL -> Results")
+    print("=" * 60)
+    print()
+    print("  URL: http://localhost:5010")
+    print("  Docs: http://localhost:5010/docs")
+    print()
+    print("  Configure LLM and Database in Settings (gear icon)")
+    print("=" * 60)
+    print()
+
+    uvicorn.run(
+        "sql_assistant.main:app",
+        host="0.0.0.0",
+        port=5010,
+        reload=False,
+        log_level="info",
+    )
+
+
+if __name__ == "__main__":
+    main()