dbagent-cli 0.9.0__tar.gz → 0.9.1__tar.gz

{dbagent_cli-0.9.0 → dbagent_cli-0.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.9.0
+Version: 0.9.1
 Summary: Enterprise AI Database Copilot, Schema Graph Introspector & Autonomous Recovery Agent (CLI)
 Author: Santhosh Gupta
 Requires-Python: >=3.9

{dbagent_cli-0.9.0 → dbagent_cli-0.9.1}/dbagent/__init__.py

@@ -2,5 +2,5 @@
 DB-Agent: Universal Database Introspection and Script Generation AI Agent (CLI).
 """
 
-__version__ = "0.9.0"
+__version__ = "0.9.1"
 __author__ = "Santhosh Gupta"

{dbagent_cli-0.9.0 → dbagent_cli-0.9.1}/dbagent/agent/pipeline.py

@@ -25,6 +25,7 @@ from dbagent.sql.validator import SQLValidator
 from dbagent.execution.cost import QueryCostEstimator, QueryComplexity
 from dbagent.execution.cache import QueryResultCache
 from dbagent.semantic.dictionary import BusinessDictionary
+from dbagent.semantic.knowledge import KnowledgeBase, LearnedRule
 from dbagent.execution.recovery import RecoveryManager, RowSnapshotRecoveryStrategy, RecoveryStatus
 from dbagent.agent.confidence import ConfidenceEvaluator, ConfidenceScore
 from dbagent.agent.conversation import ConversationContext
@@ -123,6 +124,7 @@ class QueryPipeline:
         self.choice_callback = choice_callback
 
         self.dictionary = BusinessDictionary()
+        self.knowledge_base = KnowledgeBase()
         self.result_cache = QueryResultCache()
         self.audit_logger = AuditLogger()
         self.recovery_mgr = recovery_manager or RecoveryManager()
@@ -215,10 +217,25 @@ class QueryPipeline:
                 sub_schema, join_steps = self._get_fast_schema_context(user_prompt, exact, fuzzy)
                 result.schema = sub_schema
                 result.join_steps = join_steps
+
+                # Auto-Learn from user guidance/corrections
+                prompt_l = user_prompt.lower()
+                if any(w in prompt_l for w in ["mapped with", "mapped to", "observe the database", "the rule is", "remember that", "notice that"]):
+                    self.knowledge_base.learn(
+                        db_name=self.db_name,
+                        topic=f"Mapping for {', '.join(exact[:3]) if exact else 'Entities'}",
+                        rule_description=user_prompt,
+                        relevant_tables=exact + fuzzy[:4],
+                        source="user_guidance",
+                    )
+
+                kb_context = self.knowledge_base.format_prompt_context(self.db_name, user_prompt, exact + fuzzy)
                 conv_history = self.conversation.format_context_for_prompt()
+                analysis_prompt = f"{kb_context}\n\n{user_prompt}" if kb_context else user_prompt
+
                 result.message = self.generator.generate_analysis_response(
                     schema=sub_schema,
-                    user_prompt=user_prompt,
+                    user_prompt=analysis_prompt,
                     conversation_history=conv_history,
                     model=self.model,
                 )
@@ -384,13 +401,16 @@ class QueryPipeline:
             result.error = f"Schema retrieval failed: {str(e)}"
             return result
 
-        # --- 4. Semantic Business Dictionary Context Injection ---
+        # --- 4. Semantic Business Dictionary & Knowledge Base Injection ---
         dict_context = self.dictionary.format_prompt_context(user_prompt)
+        kb_context = self.knowledge_base.format_prompt_context(self.db_name, user_prompt, exact + fuzzy)
         conv_context = self.conversation.format_context_for_prompt()
 
         augmented_prompt = user_prompt
         if dict_context:
             augmented_prompt = f"{dict_context}\n\n{augmented_prompt}"
+        if kb_context:
+            augmented_prompt = f"{kb_context}\n\n{augmented_prompt}"
         if conv_context:
             augmented_prompt = f"{conv_context}\n\n{augmented_prompt}"
 

{dbagent_cli-0.9.0 → dbagent_cli-0.9.1}/dbagent/cli.py

@@ -633,6 +633,9 @@ def chat_command(
                     cmd_table.add_row(":fk <name>", "View foreign key relationships & dependency graph")
                     cmd_table.add_row(":path <t1> <t2>", "Discover shortest SQL join path between two tables")
                     cmd_table.add_row(":doctor", "Run database health, index, and query performance diagnostics")
+                    cmd_table.add_row(":learn <rule>", "Teach DB-Agent custom domain rules & mappings (:learn topic = rule)")
+                    cmd_table.add_row(":memory, :learned", "View all learned domain rules & relationship memory")
+                    cmd_table.add_row(":forget <id|topic>", "Remove a learned rule from persistent memory")
                     cmd_table.add_row(":mode <name>", "Switch safety mode (SAFE, CONFIRM, ADMIN, READ_ONLY)")
                     cmd_table.add_row(":history", "View session query history, execution timing & row counts")
                     cmd_table.add_row(":page <n>, :card", "Inspect wide multi-column query results")
@@ -776,6 +779,52 @@ def chat_command(
                         console.print("")
                     continue
 
+                if cmd_name in [":learn", ":remember"]:
+                    if not cmd_arg or "=" not in cmd_arg:
+                        print_warning("Usage: :learn <topic> = <description> (e.g. :learn site_user_map = ivis_site maps to ivis_user via site_group -> group)")
+                        continue
+                    parts = cmd_arg.split("=", 1)
+                    topic = parts[0].strip()
+                    desc = parts[1].strip()
+                    try:
+                        exact, fuzzy = connector.resolve_tables(desc)
+                    except Exception:
+                        exact, fuzzy = [], []
+                    learned_rule = pipeline.knowledge_base.learn(
+                        db_name=db_name,
+                        topic=topic,
+                        rule_description=desc,
+                        relevant_tables=exact + fuzzy[:4],
+                        source="user_explicit",
+                    )
+                    print_success(f"Learned & Persisted Rule: [bold cyan]{learned_rule.topic}[/bold cyan]")
+                    console.print(f"  [dim]• Relevant tables: {', '.join(learned_rule.relevant_tables) or 'general'}[/dim]")
+                    console.print(f"  [dim]• DB-Agent will automatically apply this rule in all future queries.[/dim]\n")
+                    continue
+
+                if cmd_name in [":memory", ":learned", ":knowledge"]:
+                    rules = pipeline.knowledge_base.list_all(db_name)
+                    if not rules:
+                        console.print(f"[dim]No learned rules stored for '{db_name}' yet. Use [bold]:learn <topic> = <desc>[/bold] to teach DB-Agent.[/dim]\n")
+                    else:
+                        console.print(f"\n[bold cyan]🧠 DB-Agent Learned Knowledge Base for '{db_name}':[/bold cyan]")
+                        for r in rules:
+                            tbl_str = f" [cyan]({', '.join(r.relevant_tables)})[/cyan]" if r.relevant_tables else ""
+                            console.print(f"  * [bold white]{r.topic}[/bold white]{tbl_str}:")
+                            console.print(f"    [dim]{r.rule_description}[/dim] [dim italic]({r.rule_id})[/dim italic]")
+                        console.print("\n[dim]To remove a rule: [bold]:forget <rule_id|topic>[/bold][/dim]\n")
+                    continue
+
+                if cmd_name in [":forget", ":unlearn"]:
+                    if not cmd_arg:
+                        print_warning("Usage: :forget <rule_id or topic>")
+                        continue
+                    if pipeline.knowledge_base.delete(db_name, cmd_arg):
+                        print_success(f"Removed rule '{cmd_arg}' from knowledge memory.")
+                    else:
+                        print_warning(f"No matching rule found for '{cmd_arg}'.")
+                    continue
+
                 if cmd_name == ":clear":
                     conversation_ctx.clear()
                     print_success("Conversation context cleared.")

{dbagent_cli-0.9.0 → dbagent_cli-0.9.1}/dbagent/schema/graph.py

@@ -176,6 +176,36 @@ class RelationshipGraph:
 
         return None
 
+    def find_all_paths(self, start_table: str, end_table: str, max_hops: int = 4) -> List[JoinPath]:
+        """
+        Finds all valid multi-hop join paths between start_table and end_table up to max_hops.
+        Discovers direct paths, junction tables, and indirect hierarchical relationships.
+        """
+        start = start_table.lower()
+        end = end_table.lower()
+
+        if start not in self.adj or end not in self.adj:
+            return []
+
+        all_paths: List[JoinPath] = []
+        queue: deque = deque([(start, [], {start})])  # (curr, list_of_steps, visited_set)
+
+        while queue:
+            curr, steps, visited = queue.popleft()
+
+            if curr == end and steps:
+                all_paths.append(JoinPath(start_table, end_table, steps))
+                continue
+
+            if len(steps) >= max_hops:
+                continue
+
+            for neighbor, step in self.adj.get(curr, []):
+                if neighbor not in visited:
+                    queue.append((neighbor, steps + [step], visited | {neighbor}))
+
+        return all_paths
+
     def find_connecting_path_for_tables(self, table_names: List[str]) -> List[JoinStep]:
         """
         Given a set of seed tables (e.g. ['users', 'events']), finds the minimal join tree / path

dbagent_cli-0.9.1/dbagent/semantic/knowledge.py

@@ -0,0 +1,146 @@
+"""
+Persistent Knowledge Base & Self-Learning Memory Engine.
+Stores domain rules, verified multi-hop relationship paths, and custom business logic per database in ~/.dbagent/knowledge/<db_name>.json.
+Automatically injects learned domain knowledge into LLM generation context.
+"""
+
+import json
+import time
+from pathlib import Path
+from typing import Dict, Any, List, Optional
+from dataclasses import dataclass, field, asdict
+
+
+@dataclass
+class LearnedRule:
+    """A verified domain rule or relationship path learned from user interaction."""
+    rule_id: str
+    topic: str
+    rule_description: str
+    relevant_tables: List[str] = field(default_factory=list)
+    created_at: str = ""
+    source: str = "user_feedback"
+
+
+class KnowledgeBase:
+    """
+    Persistent Knowledge Base & Self-Learning Memory for DB-Agent.
+    Stores and retrieves domain-specific relationship insights per database in ~/.dbagent/knowledge/<db_name>.json.
+    """
+
+    def __init__(self, storage_dir: Optional[Path] = None):
+        self.storage_dir = storage_dir or Path.home() / ".dbagent" / "knowledge"
+        self.storage_dir.mkdir(parents=True, exist_ok=True)
+
+    def _get_db_file(self, db_name: str) -> Path:
+        safe_name = "".join(c if c.isalnum() or c in ("-", "_") else "_" for c in db_name.lower())
+        return self.storage_dir / f"{safe_name}.json"
+
+    def learn(
+        self,
+        db_name: str,
+        topic: str,
+        rule_description: str,
+        relevant_tables: Optional[List[str]] = None,
+        source: str = "user_feedback",
+    ) -> LearnedRule:
+        """Persist a new learned domain rule or mapping path."""
+        db_file = self._get_db_file(db_name)
+        rules = self.list_all(db_name)
+
+        # Check if identical topic exists and update, else create new
+        existing_idx = None
+        for i, r in enumerate(rules):
+            if r.topic.lower().strip() == topic.lower().strip():
+                existing_idx = i
+                break
+
+        rule_id = f"rule_{len(rules) + 1}_{int(time.time())}"
+        new_rule = LearnedRule(
+            rule_id=rule_id,
+            topic=topic.strip(),
+            rule_description=rule_description.strip(),
+            relevant_tables=[t.lower() for t in (relevant_tables or [])],
+            created_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+            source=source,
+        )
+
+        if existing_idx is not None:
+            rules[existing_idx] = new_rule
+        else:
+            rules.append(new_rule)
+
+        with open(db_file, "w", encoding="utf-8") as f:
+            json.dump([asdict(r) for r in rules], f, indent=2)
+        return new_rule
+
+    def list_all(self, db_name: str) -> List[LearnedRule]:
+        """List all learned rules for a database."""
+        db_file = self._get_db_file(db_name)
+        if not db_file.exists():
+            return []
+        try:
+            with open(db_file, "r", encoding="utf-8") as f:
+                data = json.load(f)
+                return [LearnedRule(**item) for item in data]
+        except Exception:
+            return []
+
+    def delete(self, db_name: str, target: str) -> bool:
+        """Delete a rule by ID or topic."""
+        rules = self.list_all(db_name)
+        target_lower = target.lower().strip()
+        filtered = [r for r in rules if r.rule_id != target and r.topic.lower() != target_lower]
+        if len(filtered) < len(rules):
+            db_file = self._get_db_file(db_name)
+            with open(db_file, "w", encoding="utf-8") as f:
+                json.dump([asdict(r) for r in filtered], f, indent=2)
+            return True
+        return False
+
+    def clear(self, db_name: str) -> None:
+        """Clear all learned rules for a database."""
+        db_file = self._get_db_file(db_name)
+        if db_file.exists():
+            db_file.unlink()
+
+    def find_relevant(self, db_name: str, user_prompt: str, table_names: Optional[List[str]] = None) -> List[LearnedRule]:
+        """Finds learned rules relevant to the user prompt or active tables."""
+        all_rules = self.list_all(db_name)
+        if not all_rules:
+            return []
+
+        prompt_lower = user_prompt.lower()
+        active_tables = {t.lower() for t in (table_names or [])}
+        matched = []
+
+        for r in all_rules:
+            # 1. Match by topic in prompt
+            if r.topic.lower() in prompt_lower:
+                matched.append(r)
+                continue
+            # 2. Match by relevant tables
+            if any(t in active_tables or t in prompt_lower for t in r.relevant_tables):
+                matched.append(r)
+                continue
+            # 3. Keyword overlap
+            keywords = [w for w in r.topic.lower().split() if len(w) > 3]
+            if any(kw in prompt_lower for kw in keywords):
+                matched.append(r)
+
+        return matched
+
+    def format_prompt_context(self, db_name: str, user_prompt: str, table_names: Optional[List[str]] = None) -> str:
+        """Format learned domain knowledge for LLM prompt context injection."""
+        matched = self.find_relevant(db_name, user_prompt, table_names)
+        if not matched:
+            return ""
+
+        lines = [
+            "### 🧠 Persistent Database Domain Knowledge & Verified Mappings (Learned from User Feedback):"
+        ]
+        for m in matched:
+            tbl_str = f" [Tables: {', '.join(m.relevant_tables)}]" if m.relevant_tables else ""
+            lines.append(f"- **{m.topic}**{tbl_str}: {m.rule_description}")
+        lines.append("Always apply these verified mapping paths and domain rules when generating SQL or architectural explanations.\n")
+        return "\n".join(lines)

{dbagent_cli-0.9.0 → dbagent_cli-0.9.1}/dbagent_cli.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.9.0
+Version: 0.9.1
 Summary: Enterprise AI Database Copilot, Schema Graph Introspector & Autonomous Recovery Agent (CLI)
 Author: Santhosh Gupta
 Requires-Python: >=3.9

{dbagent_cli-0.9.0 → dbagent_cli-0.9.1}/dbagent_cli.egg-info/SOURCES.txt

@@ -49,6 +49,7 @@ dbagent/security/modes.py
 dbagent/security/rbac.py
 dbagent/semantic/__init__.py
 dbagent/semantic/dictionary.py
+dbagent/semantic/knowledge.py
 dbagent/semantic/templates.py
 dbagent/sql/__init__.py
 dbagent/sql/optimizer.py
@@ -71,6 +72,7 @@ tests/test_direct_sql.py
 tests/test_doctor_and_viz.py
 tests/test_generator.py
 tests/test_graph.py
+tests/test_knowledge.py
 tests/test_providers.py
 tests/test_recovery.py
 tests/test_safety.py

{dbagent_cli-0.9.0 → dbagent_cli-0.9.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dbagent-cli"
-version = "0.9.0"
+version = "0.9.1"
 description = "Enterprise AI Database Copilot, Schema Graph Introspector & Autonomous Recovery Agent (CLI)"
 readme = "README.md"
 requires-python = ">=3.9"

dbagent_cli-0.9.1/tests/test_knowledge.py

@@ -0,0 +1,80 @@
+"""
+Tests for Persistent KnowledgeBase & Multi-Hop find_all_paths.
+"""
+
+import pytest
+from pathlib import Path
+from dbagent.semantic.knowledge import KnowledgeBase, LearnedRule
+from dbagent.schema.graph import RelationshipGraph
+from dbagent.schema.models import DatabaseSchema, TableModel, ColumnModel, ForeignKeyModel
+
+
+@pytest.fixture
+def temp_kb(tmp_path):
+    return KnowledgeBase(storage_dir=tmp_path)
+
+
+def test_knowledge_base_learn_and_find(temp_kb):
+    rule = temp_kb.learn(
+        db_name="test_db",
+        topic="Site to User Group Mapping",
+        rule_description="ivis_site maps to ivis_user through ivis_site_group -> ivis_group -> ivis_user_group_map",
+        relevant_tables=["ivis_site", "ivis_site_group", "ivis_group", "ivis_user_group_map", "ivis_user"],
+        source="user_correction",
+    )
+    assert rule.rule_id.startswith("rule_")
+    assert rule.topic == "Site to User Group Mapping"
+
+    # List all
+    all_rules = temp_kb.list_all("test_db")
+    assert len(all_rules) == 1
+
+    # Find relevant by prompt
+    matched = temp_kb.find_relevant("test_db", "how to map sites to user")
+    assert len(matched) == 1
+    assert matched[0].topic == "Site to User Group Mapping"
+
+    # Context formatting
+    ctx = temp_kb.format_prompt_context("test_db", "how to map sites to user")
+    assert "Persistent Database Domain Knowledge" in ctx
+    assert "Site to User Group Mapping" in ctx
+
+    # Delete
+    deleted = temp_kb.delete("test_db", rule.rule_id)
+    assert deleted is True
+    assert len(temp_kb.list_all("test_db")) == 0
+
+
+def test_relationship_graph_find_all_paths():
+    schema = DatabaseSchema(
+        dialect_name="postgresql",
+        database_name="test_db",
+        tables=[
+            TableModel(
+                name="users",
+                columns=[ColumnModel(name="id", data_type="INT", is_primary_key=True)],
+                primary_key=["id"],
+            ),
+            TableModel(
+                name="user_groups",
+                columns=[
+                    ColumnModel(name="user_id", data_type="INT"),
+                    ColumnModel(name="group_id", data_type="INT"),
+                ],
+                foreign_keys=[
+                    ForeignKeyModel(name="fk_u", constrained_columns=["user_id"], referred_table="users", referred_columns=["id"]),
+                    ForeignKeyModel(name="fk_g", constrained_columns=["group_id"], referred_table="groups", referred_columns=["id"]),
+                ],
+            ),
+            TableModel(
+                name="groups",
+                columns=[ColumnModel(name="id", data_type="INT", is_primary_key=True)],
+                primary_key=["id"],
+            ),
+        ],
+    )
+
+    graph = RelationshipGraph(schema)
+    paths = graph.find_all_paths("users", "groups", max_hops=3)
+    assert len(paths) >= 1
+    assert paths[0].tables == ["users", "user_groups", "groups"]