Merge pull request #234 from thrcle/feat/enrich-rowlimit-session

Feat/enrich rowlimit session

Author: seyoung4503

.env.example

@@ -4,11 +4,24 @@
 # Required to run `lang2sql-bot`; the bot exits with a clear error if it's unset.
 DISCORD_BOT_TOKEN=
 
+# Set to true only when you add/remove slash commands to re-sync with Discord.
+# Leave unset (or false) during normal development to avoid rate limits.
+LANG2SQL_SYNC_COMMANDS=
+
 # OpenAI API key. Optional: when set, the agent uses gpt-4.1-mini. When unset,
 # it falls back to the offline FakeLLM (deterministic canned tool cycles — fine
 # for a smoke run, not for real answers).
 OPENAI_API_KEY=
 
+# ── Local LLM (vLLM / Ollama) ────────────────────────────────────────────
+# When LANG2SQL_LLM_BASE_URL is set it takes priority over OPENAI_API_KEY.
+# Point to the base URL of any OpenAI-compatible server.
+#   vLLM:  http://localhost:8000
+#   Ollama: http://localhost:11434
+LANG2SQL_LLM_BASE_URL=
+# Model name as the server expects it (e.g. Qwen/Qwen3-14B-AWQ for vLLM).
+LANG2SQL_LLM_MODEL=
+
 # Fernet key used to encrypt stored secrets (DSNs / API keys) at rest. Optional:
 # if unset, a key is auto-generated and persisted in the SQLite kv table. Set it
 # in production so secrets decrypt across restarts and machines. Generate one:

dev.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+# Auto-reload dev runner for lang2sql-bot.
+# Watches src/ for .py changes and restarts the bot automatically.
+
+set -a
+source "$(dirname "$0")/.env"
+set +a
+
+REF=$(mktemp)
+
+restart_bot() {
+    if [ -n "$BOT_PID" ] && kill -0 "$BOT_PID" 2>/dev/null; then
+        echo "[watch] stopping PID $BOT_PID..."
+        kill "$BOT_PID"
+        wait "$BOT_PID" 2>/dev/null
+    fi
+    echo "[watch] starting bot..."
+    .venv/bin/lang2sql-bot &
+    BOT_PID=$!
+    touch "$REF"
+    echo "[watch] PID $BOT_PID"
+}
+
+trap 'kill $BOT_PID 2>/dev/null; rm -f $REF; exit' INT TERM
+
+restart_bot
+
+while true; do
+    sleep 2
+    if find src/ -name "*.py" -newer "$REF" | grep -q .; then
+        CHANGED=$(find src/ -name "*.py" -newer "$REF" | head -3 | tr '\n' ' ')
+        echo "[watch] changed: $CHANGED"
+        restart_bot
+    elif ! kill -0 "$BOT_PID" 2>/dev/null; then
+        echo "[watch] bot crashed, restarting in 2s..."
+        sleep 2
+        restart_bot
+    fi
+done

src/lang2sql/adapters/db/dsn_builder.py

@@ -10,7 +10,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from urllib.parse import quote_plus
+from urllib.parse import quote_plus, urlsplit
 
 
 @dataclass
@@ -37,8 +37,13 @@ def _quote(s: str) -> str:
 
 
 def build_postgresql(*, host: str, port: str, database: str, user: str, password: str) -> ConnectionSpec:
+    # User may paste a full URL (e.g. "host/db?sslmode=require") into the host field.
+    # Extract just the hostname to avoid corrupting the assembled DSN.
+    parsed = urlsplit("//" + host)
+    clean_host = parsed.hostname or host
     p = int(port) if port else 5432
-    dsn = f"postgresql+psycopg://{_quote(user)}:{_quote(password)}@{host}:{p}/{database}"
+    suffix = "?sslmode=require" if clean_host.endswith(".neon.tech") else ""
+    dsn = f"postgresql+psycopg://{_quote(user)}:{_quote(password)}@{clean_host}:{p}/{database}{suffix}"
     return ConnectionSpec(dsn=dsn, extras={})
 
 

src/lang2sql/adapters/db/factory.py

@@ -55,6 +55,10 @@ def build_explorer(
             token=extras.get("d1_token"),
         )
 
+    # Normalize bare postgresql:// → postgresql+psycopg:// (psycopg3 is installed).
+    if scheme == "postgresql":
+        connection = "postgresql+psycopg" + connection[len("postgresql"):]
+
     # Anything else is assumed to be a SQLAlchemy URL (driver loaded lazily).
     return SqlAlchemyExplorer(connection, schema=schema)
 
@@ -67,7 +71,8 @@ def explorer_from_env() -> ExplorerPort | None:
     """
     url = os.environ.get("LANG2SQL_DB_URL")
     if url:
-        return build_explorer(url, schema=os.environ.get("LANG2SQL_DB_SCHEMA"))
+        schema = os.environ.get("LANG2SQL_DB_SCHEMA") or None
+        return build_explorer(url, schema=schema)
 
     account = os.environ.get("CLOUDFLARE_D1_ACCOUNT_ID")
     database = os.environ.get("CLOUDFLARE_D1_DATABASE_ID")

src/lang2sql/adapters/db/sqlalchemy_explorer.py

@@ -55,10 +55,15 @@ async def execute(self, sql: str, limit: int = 1000) -> list[dict]:
     def _list_tables_sync(self) -> list[Table]:
         from sqlalchemy import inspect
 
-        insp = inspect(self._get_engine())
-        schema = self._schema or insp.default_schema_name
+        engine = self._get_engine()
+        engine.dispose()  # flush stale pool connections so schema changes are visible
+        insp = inspect(engine)
+        default = insp.default_schema_name
+        effective = self._schema or default
+        # Omit schema when it's the connection default so SQL stays unqualified.
+        display_schema = "" if (not self._schema or self._schema == default) else effective
         return [
-            Table(name=t, schema=schema or "")
+            Table(name=t, schema=display_schema)
             for t in insp.get_table_names(schema=self._schema)
         ]
 

src/lang2sql/adapters/llm/openai_.py

@@ -11,8 +11,10 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 import os
+import re
 import urllib.error
 import urllib.request
 from typing import Any, Sequence
@@ -27,15 +29,16 @@ class OpenAILLM:
 
     def __init__(
         self,
-        model: str = "gpt-4.1-mini",
+        model: str = "gpt-4o-mini",
         api_key: str | None = None,
         *,
         base_url: str = _DEFAULT_URL,
         timeout: float = 60.0,
     ) -> None:
         self.model = model
         # Resolve lazily-ish: read env now, but tolerate absence until complete().
-        self._api_key = api_key if api_key is not None else os.environ.get("OPENAI_API_KEY")
+        raw_key = api_key if api_key is not None else os.environ.get("OPENAI_API_KEY")
+        self._api_key = raw_key.strip() if raw_key else raw_key
         self._base_url = base_url
         self._timeout = timeout
 
@@ -54,7 +57,7 @@ async def complete(
         if tools:
             payload["tools"] = [_encode_tool(t) for t in tools]
 
-        raw = self._post(payload)
+        raw = await asyncio.to_thread(self._post, payload)
         return _decode_completion(raw)
 
     def _post(self, payload: dict[str, Any]) -> dict[str, Any]:
@@ -83,11 +86,19 @@ def _post(self, payload: dict[str, Any]) -> dict[str, Any]:
             raise RuntimeError(f"OpenAI returned non-JSON response: {text[:200]!r}") from exc
 
 
+def _strip_thinking(text: str) -> str:
+    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
+
+
 def _encode_message(m: Message) -> dict[str, Any]:
     """Core :class:`Message` → an OpenAI chat message dict."""
     out: dict[str, Any] = {"role": m.role.value}
-    # OpenAI wants content present (may be null when only tool_calls are set).
-    out["content"] = m.content or None
+    # OpenAI allows null content only when tool_calls are present.
+    # For plain assistant messages (after session compress), force empty string.
+    if m.role == Role.ASSISTANT and not m.tool_calls:
+        out["content"] = m.content or ""
+    else:
+        out["content"] = m.content or None
     if m.role == Role.TOOL:
         out["tool_call_id"] = m.tool_call_id
         if m.name:
@@ -141,7 +152,7 @@ def _decode_completion(raw: dict[str, Any]) -> Completion:
         )
 
     return Completion(
-        content=msg.get("content") or "",
+        content=_strip_thinking(msg.get("content") or ""),
         tool_calls=tool_calls,
         finish_reason=choice.get("finish_reason"),
     )

src/lang2sql/adapters/storage/sqlite_store.py

@@ -130,6 +130,15 @@ def kv_delete(self, scope: str, key: str) -> None:
         )
         self._conn.commit()
 
+    def kv_delete_prefix(self, scope: str, prefix: str) -> int:
+        """Delete all keys under scope that start with prefix. Returns count deleted."""
+        cur = self._conn.execute(
+            "DELETE FROM kv WHERE scope = ? AND key LIKE ?",
+            (scope, prefix + "%"),
+        )
+        self._conn.commit()
+        return cur.rowcount
+
 
 # -- Session (de)serialization ------------------------------------------
 

src/lang2sql/frontends/discord/bot.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import io
+import logging
 import os
 
 import discord
@@ -25,6 +26,8 @@
 from .commands import CommandHandlers
 from .session_router import InteractionContext, to_identity
 
+logger = logging.getLogger(__name__)
+
 TOKEN_ENV = "DISCORD_BOT_TOKEN"
 
 
@@ -102,8 +105,11 @@ def __init__(self, handlers: CommandHandlers) -> None:
         self._register_commands()
 
     async def setup_hook(self) -> None:
-        # Sync slash commands with Discord on startup.
-        await self.tree.sync()
+        # Sync only when LANG2SQL_SYNC_COMMANDS=true (e.g. after adding/removing commands).
+        # Skipping sync on every restart avoids Discord rate limits during dev.
+        if os.environ.get("LANG2SQL_SYNC_COMMANDS", "").lower() == "true":
+            await self.tree.sync()
+            logger.info("slash commands synced")
 
     def _register_commands(self) -> None:
         tree = self.tree
@@ -135,6 +141,13 @@ async def define_metric(
         async def remember(interaction: discord.Interaction, text: str) -> None:
             await self._run(interaction, handlers.remember(to_identity(_interaction_context(interaction)), text))
 
+        @tree.command(name="enrich", description="LLM으로 DB 컬럼 메타데이터 자동 보강 (clear=True로 초기화)")
+        async def enrich(interaction: discord.Interaction, table: str = "", clear: bool = False) -> None:
+            await self._run(
+                interaction,
+                handlers.enrich(to_identity(_interaction_context(interaction)), table=table, clear=clear),
+            )
+
         @tree.command(name="semantic_show", description="Show definitions in effect here")
         async def semantic_show(interaction: discord.Interaction) -> None:
             await self._run(interaction, handlers.semantic_show(to_identity(_interaction_context(interaction))))
@@ -148,7 +161,10 @@ async def _run(self, interaction: discord.Interaction, coro) -> None:
         await interaction.response.defer(thinking=True)
         message = await coro
         content, file = _to_sendable(message)
-        await interaction.followup.send(content=content or "(empty)", file=file)
+        kwargs: dict = {"content": content or "(empty)"}
+        if file is not None:
+            kwargs["file"] = file
+        await interaction.followup.send(**kwargs)
 
     async def on_message(self, message: discord.Message) -> None:
         """Treat an @mention (or a reply inside a thread) as a free-form query."""
@@ -166,9 +182,16 @@ async def on_message(self, message: discord.Message) -> None:
             return
 
         identity = to_identity(_message_context(message))
-        out = await self._handlers.query(identity, text)
-        content, file = _to_sendable(out)
-        await message.channel.send(content=content or "(empty)", file=file)
+        try:
+            out = await self._handlers.query(identity, text)
+            content, file = _to_sendable(out)
+            if content and len(content) > 1900:
+                content = content[:1900] + "\n…(truncated)"
+            await message.channel.send(content=content or "(empty)", file=file)
+        except Exception as exc:
+            import traceback
+            traceback.print_exc()
+            await message.channel.send(content=f"❌ Error: {type(exc).__name__}: {exc}")
 
 
 def run() -> None:
@@ -182,6 +205,7 @@ def run() -> None:
         raise RuntimeError(
             f"{TOKEN_ENV} is not set; export your Discord bot token to run the bot."
         )
-    handlers = CommandHandlers(ContextConcierge())
+    data_path = os.environ.get("LANG2SQL_DATA_PATH", "lang2sql_data.db")
+    handlers = CommandHandlers(ContextConcierge(path=data_path))
     client = Lang2SQLBot(handlers)
     client.run(token)

src/lang2sql/frontends/discord/commands.py

@@ -22,6 +22,7 @@
 from ...adapters.db.dsn_builder import assemble
 from ...core.identity import Identity
 from ...core.ports.frontend import OutboundMessage
+from ...core.types import Role
 from ...harness.loop import agent_loop
 from ...tenancy.concierge import ContextConcierge
 from .render import render_answer
@@ -41,9 +42,39 @@ async def query(self, identity: Identity, text: str) -> OutboundMessage:
         thread/DM continues the conversation (tiebreaker #4).
         """
         ctx = await self._concierge.build_context(identity, user_text=text)
+        pre_loop_len = len(ctx.session.history())
         answer = await agent_loop(ctx, text)
+
+        history = ctx.session.history()
+        current_turn = history[pre_loop_len:]
+
+        call_id_to_sql: dict[str, str] = {
+            tc.id: tc.arguments["sql"]
+            for msg in current_turn
+            if msg.role == Role.ASSISTANT and msg.tool_calls
+            for tc in msg.tool_calls
+            if tc.name == "run_sql" and "sql" in tc.arguments
+        }
+
+        sql_queries: list[str] = []
+        sql_results: list[str] = []
+        for msg in current_turn:
+            if msg.role != Role.TOOL or msg.name != "run_sql" or not msg.content:
+                continue
+            sql = call_id_to_sql.get(msg.tool_call_id or "")
+            if sql and ("row(s):" in msg.content or "(0 rows)" in msg.content):
+                sql_queries.append(sql)
+                sql_results.append(msg.content)
+
+        ctx.session.compress()
         await self._concierge.store.save(identity.session_key(), ctx.session)
-        return render_answer(answer)
+
+        suffix = ""
+        if sql_queries:
+            suffix += "\n\n**SQL:**\n```sql\n" + "\n\n".join(sql_queries) + "\n```"
+        if sql_results:
+            suffix += "\n\n**결과:**\n```\n" + "\n\n".join(sql_results) + "\n```"
+        return render_answer(answer + suffix)
 
     async def define_metric(
         self,
@@ -149,6 +180,14 @@ async def register_db_for_guild(
             )
         )
 
+    async def enrich(self, identity: Identity, table: str = "", clear: bool = False) -> OutboundMessage:
+        """Run EnrichSchema tool: sample DB columns and LLM-infer descriptions."""
+        ctx = await self._concierge.build_context(identity)
+        result = await ctx.tools.dispatch(
+            "enrich_schema", {"table": table, "clear": clear}, ctx, "cmd:enrich"
+        )
+        return OutboundMessage(text=result.content)
+
     async def connect(self, identity: Identity, dsn: str) -> OutboundMessage:
         """V1 stub: stash a DB DSN keyed by guild/DM in the concierge kv store.
 

src/lang2sql/harness/context.py

@@ -9,8 +9,12 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 from ..core.identity import Identity
+
+if TYPE_CHECKING:
+    from ..adapters.storage.sqlite_store import SqliteStore
 from ..core.ports.audit import AuditPort
 from ..core.ports.explorer import ExplorerPort
 from ..core.ports.llm import LLMPort
@@ -30,4 +34,5 @@ class HarnessContext:
     safety: SafetyPipelinePort | None = None
     audit: AuditPort | None = None
     scope_resolver: ScopeResolverPort | None = None
+    store: SqliteStore | None = None
     max_turns: int = 8