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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.9.0
+Version: 0.9.2
 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.2}/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.2"
 __author__ = "Santhosh Gupta"

{dbagent_cli-0.9.0 → dbagent_cli-0.9.2}/dbagent/agent/conversation.py

@@ -54,9 +54,13 @@ class ConversationContext:
         if len(self.turns) > self.max_history:
             self.turns.pop(0)
 
-        if tables and not is_chat:
+        if tables:
             self.active_tables.update(t.lower() for t in tables)
 
+    def get_active_tables(self) -> List[str]:
+        """Returns currently active database table names discussed in this session."""
+        return list(self.active_tables)
+
     def get_last_turn(self) -> Optional[ConversationTurn]:
         if self.turns:
             return self.turns[-1]

{dbagent_cli-0.9.0 → dbagent_cli-0.9.2}/dbagent/agent/generator.py

@@ -106,6 +106,11 @@ Provide ready-to-run, dialect-accurate SQL JOIN queries answering the user's spe
 
 ### 4. 💡 Architectural Insights & Best Practices
 Mention indexing recommendations, cascade rules, or performance considerations for these join paths.
+
+### 🛑 Critical Operational Directives:
+1. NEVER ask passive questions like "Would you like me to inspect tables?" or "Should I run a query?". Always proactively inspect and deliver the full architectural blueprint immediately.
+2. NEVER give generic textbook definitions. Always ground every explanation in the real database schema provided above.
+3. Always include ALL 4 sections (Diagram, Theoretical Breakdown, SQL Queries, and Insights).
 """
 
 
@@ -142,37 +147,52 @@ class ScriptGenerator:
         prompt_lower = user_prompt.lower().strip()
         clean_prompt = prompt_lower.rstrip(".!?,:;~ ")
 
-        # Chat / Conversational indicators (greetings, feedback, praise, acknowledgements)
-        chat_triggers = [
-            "hello", "hi", "hey", "howdy", "greetings", "good morning", "good evening", "good afternoon",
-            "who are you", "what can you do", "what are you", "help me", "how to use",
-            "what is this", "tell me about this database", "how does this work",
-            "can you help", "thanks", "thank you", "thx", "bye", "goodbye", "see ya",
-            "good", "great", "nice", "cool", "awesome", "perfect", "ok", "okay",
-            "yes", "no", "sure", "got it", "understood", "nice work", "well done",
-            "sounds good", "looks good", "alright", "super", "neat", "excellent",
-            "amazing", "wonderful", "cool thanks", "great thanks", "thanks a lot",
-        ]
-        for trig in chat_triggers:
-            if clean_prompt == trig or prompt_lower.startswith(f"{trig} ") or prompt_lower.startswith(f"{trig},") or prompt_lower.startswith(f"{trig}!"):
-                return "chat"
-
-        # Architecture / Relationship / Linkage / Analysis indicators
+        # 1. Architecture / Relationship / Linkage / Analysis indicators (Highest Priority)
         analysis_keywords = [
             "linkage architecture", "linkage", "architecture", "relationship between",
             "how are", "how is", "how were", "how they were related", "how were they related",
-            "how they were mapped", "how to map", "mapped to", "er diagram", "diagram of",
+            "how they were mapped", "how to map", "where they mapped", "where is it mapped",
+            "how they were assigned", "how are they assigned", "how will be linked", "how is it linked",
+            "mapped to", "mapped with", "er diagram", "diagram of",
             "data model", "explain relationship", "explain schema", "schema analysis",
             "analyze the database", "analyze database", "table relationship",
             "foreign key relationship", "how are tables related", "which sites in",
             "how to know that which", "how to know which", "mapping architecture",
             "entity relationship", "theoritical information", "theoretical information",
+            "check with my previous context", "check with previous context", "previous context",
+            "complete architecture", "scope architecture", "tenant level to user level",
+            "tenant level", "user level",
         ]
         for kw in analysis_keywords:
             if kw in prompt_lower:
                 return "analysis"
 
-        # DDL indicators
+        # Contextual clarification for analysis (e.g. "i mean ivis_scope", "i mean table_x")
+        if re.search(r'\bi mean\b', prompt_lower) and not re.search(r'\b(insert|update|delete|drop)\b', prompt_lower):
+            return "analysis"
+
+        # 2. Pure Greeting / Feedback / Help triggers (Only when not asking a technical question)
+        pure_chat_greetings = ["hello", "hi", "hey", "howdy", "greetings", "good morning", "good evening", "good afternoon"]
+        pure_chat_feedback = [
+            "thanks", "thank you", "thx", "bye", "goodbye", "see ya",
+            "good", "great", "nice", "cool", "awesome", "perfect", "ok", "okay",
+            "yes", "no", "sure", "got it", "understood", "nice work", "well done",
+            "sounds good", "looks good", "alright", "super", "neat", "excellent",
+            "amazing", "wonderful", "cool thanks", "great thanks", "thanks a lot",
+        ]
+        pure_chat_help = [
+            "who are you", "what can you do", "what are you", "help me", "how to use",
+            "what is this", "tell me about this database", "how does this work", "can you help",
+        ]
+
+        if clean_prompt in pure_chat_greetings or clean_prompt in pure_chat_feedback or clean_prompt in pure_chat_help:
+            return "chat"
+
+        for trig in pure_chat_greetings + pure_chat_help:
+            if prompt_lower.startswith(f"{trig} ") or prompt_lower.startswith(f"{trig},") or prompt_lower.startswith(f"{trig}!"):
+                return "chat"
+
+        # 3. DDL indicators
         ddl_keywords = [
             "create table", "alter table", "drop table", "add column",
             "remove column", "rename column", "rename table", "modify column",
@@ -182,7 +202,7 @@ class ScriptGenerator:
             if kw in prompt_lower:
                 return "ddl"
 
-        # Write indicators (word-boundary matching to avoid false positives like "show dataset")
+        # 4. Write indicators (word-boundary matching to avoid false positives like "show dataset")
         write_phrases = [
             "insert", "add record", "add row", "add new",
             "update", "delete", "drop record",
@@ -191,14 +211,10 @@ class ScriptGenerator:
         for kw in write_phrases:
             if kw in prompt_lower:
                 return "write"
-        # Single-word write verbs need word boundary to avoid false positives
-        # e.g. "set" should not match "dataset", "change" should not match "show changes"
-        import re as _re
         write_verbs = ["set", "change", "modify", "remove"]
         for verb in write_verbs:
-            if _re.search(r'\b' + verb + r'\b', prompt_lower):
-                # Exclude when preceded by 'show', 'list', 'display', 'view', 'get'
-                read_context = _re.search(r'\b(?:show|list|display|view|get|count)\b.*\b' + verb, prompt_lower)
+            if re.search(r'\b' + verb + r'\b', prompt_lower):
+                read_context = re.search(r'\b(?:show|list|display|view|get|count)\b.*\b' + verb, prompt_lower)
                 if not read_context:
                     return "write"
 

{dbagent_cli-0.9.0 → dbagent_cli-0.9.2}/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()
@@ -212,13 +214,34 @@ class QueryPipeline:
                 except Exception:
                     exact, fuzzy = [], []
 
+                # Inherit active tables from conversation context if follow-up prompt doesn't name tables explicitly
+                if not exact and not fuzzy:
+                    inherited = self.conversation.get_active_tables()
+                    if inherited:
+                        exact = inherited[:4]
+
                 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,
                 )
@@ -227,7 +250,13 @@ class QueryPipeline:
                 result.was_executed = True
                 result.query_type = "analysis"
                 result.execution_time_ms = (time.time() - t_start) * 1000
-                self.conversation.add_turn(user_prompt=user_prompt, sql="", assistant_response=result.message, is_chat=True)
+                self.conversation.add_turn(
+                    user_prompt=user_prompt,
+                    sql="",
+                    assistant_response=result.message,
+                    tables=exact + fuzzy,
+                    is_chat=True,
+                )
                 return result
             except Exception as e:
                 result.error = f"Architecture analysis failed: {str(e)}"
@@ -384,13 +413,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.2}/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.2}/dbagent/connectors/base.py

@@ -44,14 +44,42 @@ class BaseConnector(ABC):
 
     def resolve_tables(self, user_prompt: str) -> Tuple[List[str], List[str]]:
         """
-        Resolve which tables the user is referring to.
+        Resolve which tables the user is referring to with prefix stripping, pluralization handling, and fuzzy matching.
         Returns (exact_matches, fuzzy_matches).
-        Default implementation uses get_table_names with simple substring matching.
         """
+        import re
         all_tables = self.get_table_names()
         prompt_lower = user_prompt.lower()
-        exact = [t for t in all_tables if t.lower() in prompt_lower]
-        return exact, []
+        exact = []
+        fuzzy = []
+
+        tokens = set(re.findall(r'\b[a-zA-Z0-9_]+\b', prompt_lower))
+
+        for t in all_tables:
+            t_lower = t.lower()
+            # 1. Exact full table name in prompt
+            if re.search(r'\b' + re.escape(t_lower) + r'\b', prompt_lower):
+                exact.append(t)
+                continue
+
+            # 2. Match without common prefixes (e.g. "ivis_scope" -> "scope", "scopes")
+            unprefixed = re.sub(r'^(?:ivis_|tbl_|table_|app_)', '', t_lower)
+            if unprefixed and len(unprefixed) >= 3:
+                if (unprefixed in tokens or 
+                    f"{unprefixed}s" in tokens or 
+                    f"{unprefixed}es" in tokens or 
+                    re.search(r'\b' + re.escape(unprefixed) + r'(?:s|es)?\b', prompt_lower)):
+                    if t not in exact:
+                        exact.append(t)
+                        continue
+
+            # 3. Fuzzy partial match
+            for tok in tokens:
+                if len(tok) >= 4 and (tok in t_lower or t_lower in tok):
+                    if t not in exact and t not in fuzzy:
+                        fuzzy.append(t)
+
+        return exact, fuzzy
 
     def refresh_cache(self) -> None:
         """Invalidate any cached metadata. Override in subclasses."""

{dbagent_cli-0.9.0 → dbagent_cli-0.9.2}/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.2/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.2}/dbagent_cli.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.9.0
+Version: 0.9.2
 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.2}/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.2}/pyproject.toml

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

dbagent_cli-0.9.2/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"]