de-agentic 2.1.1__py3-none-any.whl

ai/__init__.py

@@ -0,0 +1 @@
+"""ai package for the simplified DE Agentic layout."""

ai/assistant.py

@@ -0,0 +1,124 @@
+"""The AI assistant: a thin, bounded-context layer over ``core.llm``.
+
+Two things make it more than a chat wrapper:
+
+* **Bounded memory.** The conversation is a deque capped at ``context_window``
+  turns, so a long session cannot silently grow the prompt.
+* **Grounded answers.** The schema digest is injected on every question, and
+  any SQL the model returns can be executed and fed back as the next turn -
+  so "what is the total revenue?" is answered with a real result, not a guess.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from typing import Any, Deque, Dict, List, Optional
+
+from core import llm
+from operations import query as query_ops
+from operations import profiling
+
+SYSTEM_PROMPT = (
+    "You are a concise data engineering assistant. Use only the tables and "
+    "columns listed in the schema - never invent others. Missing data is NULL: "
+    "test it with IS NULL / IS NOT NULL, never = '' or = 0. Reply with at most "
+    "4 short sentences, or one small SQL block with no extra commentary."
+)
+
+
+def extract_sql(text: str) -> str:
+    """Pull a SQL statement out of an answer (fenced block or bare)."""
+    import re
+
+    fenced = re.search(r"```(?:sql)?\s*(.+?)```", text or "", re.S | re.I)
+    if fenced:
+        return fenced.group(1).strip()
+    bare = re.search(r"\b(WITH|SELECT)\b.*", text or "", re.S | re.I)
+    return bare.group(0).strip().rstrip(";") if bare else ""
+
+
+class Assistant:
+    """A question/answer loop with bounded context and optional SQL execution."""
+
+    def __init__(self, settings: Optional[llm.Settings] = None,
+                 db: str = "sample.sqlite", context_window: int = 6,
+                 execute_sql: bool = True):
+        self.settings = settings or llm.resolve()
+        self.db = db
+        self.context_window = max(1, int(context_window))
+        self.execute_sql = execute_sql
+        self._history: Deque[Dict[str, str]] = deque(maxlen=self.context_window * 2)
+        self.last_execution: Optional[Dict[str, Any]] = None
+
+    # -- context -------------------------------------------------------------
+    @property
+    def history(self) -> List[Dict[str, str]]:
+        return list(self._history)
+
+    def clear(self) -> None:
+        self._history.clear()
+        self.last_execution = None
+
+    def _turns(self) -> List[Dict[str, str]]:
+        # deque has no slicing, so materialise before taking the last N turns.
+        return list(self._history)[-self.context_window:]
+
+    def schema_lines(self) -> str:
+        return "\n".join(profiling.profile_database(self.db)["lines"])
+
+    def build_prompt(self, question: str) -> str:
+        return (f"Database: {self.db}\n"
+                f"Schema:\n{self.schema_lines()}\n\n"
+                f"Question: {question}\n\n"
+                "Answer briefly and concretely.")
+
+    # -- asking --------------------------------------------------------------
+    def ask(self, question: str, stream: bool = False) -> str:
+        """Answer a question, optionally executing any SQL it returns."""
+        messages = [{"role": "system", "content": SYSTEM_PROMPT},
+                    {"role": "user", "content": self.build_prompt(question)}]
+        messages.extend(self._turns())
+        answer = llm.chat(self.settings, messages, stream=stream)
+        if stream:
+            return answer  # caller is rendering deltas; history is not updated
+        self._history.append({"role": "user", "content": question})
+        self._history.append({"role": "assistant", "content": str(answer)})
+        self.last_execution = None
+        if self.execute_sql:
+            self.last_execution = self._maybe_execute(str(answer))
+            if self.last_execution and self.last_execution.get("ok"):
+                self._history.append({
+                    "role": "user",
+                    "content": "Here is the result of running that SQL:\n"
+                               f"{self.last_execution['summary']}\n"
+                               "State the answer in one sentence.",
+                })
+                follow_up = llm.chat(self.settings, messages + [
+                    {"role": "user", "content": self._history[-1]["content"]}
+                ], stream=False)
+                self._history.append({"role": "assistant", "content": str(follow_up)})
+                return str(follow_up)
+        return str(answer)
+
+    # -- SQL execution -------------------------------------------------------
+    def _maybe_execute(self, answer: str) -> Optional[Dict[str, Any]]:
+        sql = extract_sql(answer)
+        if not sql:
+            return None
+        try:
+            result = query_ops.run_query(self.db, sql, max_rows=20)
+        except query_ops.QueryError as exc:
+            return {"ok": False, "sql": sql, "error": str(exc)}
+        rendered = "; ".join(
+            f"{column}={value}" for column, value in zip(result["columns"],
+                                                        result["rows"][0])
+        ) if result["rows"] else "(no rows)"
+        return {
+            "ok": True,
+            "sql": sql,
+            "row_count": result["row_count"],
+            "columns": result["columns"],
+            "rows": result["rows"],
+            "summary": f"{result['row_count']} row(s). First row: {rendered}"
+                       if result["rows"] else "0 rows returned.",
+        }

ai/tasks/__init__.py

@@ -0,0 +1 @@
+"""ai/tasks package for the simplified DE Agentic layout."""