dbagent-cli 0.9.6__tar.gz → 0.9.8__tar.gz

{dbagent_cli-0.9.6 → dbagent_cli-0.9.8}/PKG-INFO

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

{dbagent_cli-0.9.6 → dbagent_cli-0.9.8}/dbagent/__init__.py

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

{dbagent_cli-0.9.6 → dbagent_cli-0.9.8}/dbagent/agent/generator.py

@@ -247,6 +247,7 @@ class ScriptGenerator:
         user_prompt: str,
         conversation_history: Optional[str] = None,
         model: Optional[str] = None,
+        stream_callback: Optional[Callable[[str], None]] = None,
     ) -> str:
         """Generate a natural conversational response to the user's message with full session memory."""
         table_names = [t.name for t in schema.tables]
@@ -268,6 +269,7 @@ class ScriptGenerator:
             prompt=full_prompt,
             system_prompt=system_prompt,
             model=model,
+            stream_callback=stream_callback,
         )
 
     def generate_analysis_response(
@@ -276,6 +278,7 @@ class ScriptGenerator:
         user_prompt: str,
         conversation_history: Optional[str] = None,
         model: Optional[str] = None,
+        stream_callback: Optional[Callable[[str], None]] = None,
     ) -> str:
         """Generate a comprehensive architectural analysis report with ASCII diagrams, theory, and SQL queries."""
         # 1. Filter relevant tables including foreign keys and junction tables
@@ -301,6 +304,7 @@ class ScriptGenerator:
             prompt=full_prompt,
             system_prompt=system_prompt,
             model=model,
+            stream_callback=stream_callback,
         )
 
     # --- Full generation (for `generate` command) ---

{dbagent_cli-0.9.6 → dbagent_cli-0.9.8}/dbagent/agent/pipeline.py

@@ -56,6 +56,7 @@ class PipelineResult:
     # Conversational chat response
     is_chat: bool = False
     is_direct_sql: bool = False
+    was_streamed: bool = False
     message: str = ""
     # Query classification
     query_type: str = "unknown"  # read, write, ddl, chat, unknown
@@ -192,7 +193,13 @@ class QueryPipeline:
         retriever = ContextRetriever(targeted)
         return retriever.retrieve_context(user_prompt, exact_tables=exact, fuzzy_tables=fuzzy, max_tables=12)
 
-    def run(self, user_prompt: str, force: bool = False, bypass_cache: bool = False) -> PipelineResult:
+    def run(
+        self,
+        user_prompt: str,
+        force: bool = False,
+        bypass_cache: bool = False,
+        stream_callback: Optional[Callable[[str], None]] = None,
+    ) -> PipelineResult:
         """
         Executes end-to-end pipeline: Prompt -> Security/Semantics -> SQL -> Safe Exec -> Results.
         """
@@ -217,7 +224,10 @@ class QueryPipeline:
                     user_prompt=user_prompt,
                     conversation_history=conv_history,
                     model=self.model,
+                    stream_callback=stream_callback,
                 )
+                if stream_callback:
+                    result.was_streamed = True
                 result.is_chat = True
                 result.was_executed = True
                 result.query_type = "chat"
@@ -238,12 +248,18 @@ class QueryPipeline:
                     exact, fuzzy = [], []
 
                 # Inherit active tables from conversation context if follow-up prompt doesn't name tables explicitly
+                is_inherited = False
                 if not exact and not fuzzy:
                     inherited = self.conversation.get_active_tables()
                     if inherited:
                         exact = inherited[:8]
+                        is_inherited = True
+
+                effective_prompt = user_prompt
+                if is_inherited and exact:
+                    effective_prompt = f"Focus strictly on the architecture, linkages, and data flow between {', '.join(exact[:5])}. User request: {user_prompt}"
 
-                sub_schema, join_steps = self._get_fast_schema_context(user_prompt, exact, fuzzy)
+                sub_schema, join_steps = self._get_fast_schema_context(effective_prompt, exact, fuzzy)
                 result.schema = sub_schema
                 result.join_steps = join_steps
 
@@ -258,16 +274,19 @@ class QueryPipeline:
                         source="user_guidance",
                     )
 
-                kb_context = self.knowledge_base.format_prompt_context(self.db_name, user_prompt, exact + fuzzy)
+                kb_context = self.knowledge_base.format_prompt_context(self.db_name, effective_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
+                analysis_prompt = f"{kb_context}\n\n{effective_prompt}" if kb_context else effective_prompt
 
                 result.message = self.generator.generate_analysis_response(
                     schema=sub_schema,
                     user_prompt=analysis_prompt,
                     conversation_history=conv_history,
                     model=self.model,
+                    stream_callback=stream_callback,
                 )
+                if stream_callback:
+                    result.was_streamed = True
                 result.ground_truth = SchemaGroundTruthValidator.validate_response(result.message, sub_schema)
                 result.is_chat = True
                 result.was_executed = True

{dbagent_cli-0.9.6 → dbagent_cli-0.9.8}/dbagent/cli.py

@@ -339,7 +339,10 @@ def _display_pipeline_result(result: PipelineResult, show_sql: bool = True) -> N
 
     if result.is_chat and result.message:
         console.print("")
-        console.print(Markdown(result.message))
+        if not result.was_streamed:
+            console.print(Markdown(result.message))
+        if result.ground_truth:
+            console.print(result.ground_truth.format_summary())
         if result.execution_time_ms > 0:
             print_timing(result.execution_time_ms)
         console.print("")
@@ -527,8 +530,25 @@ def ask_command(
         choice_callback=_choice_callback,
     )
 
-    with DynamicSpinner():
-        result = pipeline.run(user_prompt=prompt, force=force, bypass_cache=no_cache)
+    spinner = DynamicSpinner.for_prompt(prompt, model_name=model or getattr(llm, "default_model", None), connector=connector)
+    stream_started = [False]
+
+    def _query_stream_callback(chunk: str) -> None:
+        if not stream_started[0]:
+            stream_started[0] = True
+            spinner.pause()
+            timing_str = f"⚡ {spinner.elapsed_ms / 1000:.2f}s" if spinner.elapsed_ms > 0 else ""
+            active_model = model or getattr(llm, "default_model", None)
+            model_display = f" ({active_model})" if active_model else ""
+            console.print(f"\n[bold cyan]● DB-Agent{model_display}[/bold cyan] [dim]{timing_str}[/dim]")
+        try:
+            sys.stdout.write(chunk)
+            sys.stdout.flush()
+        except Exception:
+            pass
+
+    with spinner:
+        result = pipeline.run(user_prompt=prompt, force=force, bypass_cache=no_cache, stream_callback=_query_stream_callback)
 
     _display_pipeline_result(result, show_sql=True)
     connector.close()
@@ -1094,14 +1114,30 @@ def chat_command(
                 console.print("[bold cyan]Goodbye![/bold cyan]")
                 break
 
-            # --- Natural language query via pipeline ---
-            # Update pipeline settings in-place (pipeline was created once before the loop)
+            # --- Natural language query via pipeline with live streaming ---
             pipeline.model = model
             pipeline.mode = current_mode
             pipeline.auto_execute = auto_execute
 
-            with DynamicSpinner.for_prompt(user_input, model_name=model or getattr(llm, "default_model", None), connector=connector):
-                result = pipeline.run(user_prompt=user_input)
+            spinner = DynamicSpinner.for_prompt(user_input, model_name=model or getattr(llm, "default_model", None), connector=connector)
+            stream_started = [False]
+
+            def _live_stream_callback(chunk: str) -> None:
+                if not stream_started[0]:
+                    stream_started[0] = True
+                    spinner.pause()
+                    timing_str = f"⚡ {spinner.elapsed_ms / 1000:.2f}s" if spinner.elapsed_ms > 0 else ""
+                    active_model = model or getattr(llm, "default_model", None)
+                    model_display = f" ({active_model})" if active_model else ""
+                    console.print(f"\n[bold cyan]● DB-Agent{model_display}[/bold cyan] [dim]{timing_str}[/dim]")
+                try:
+                    sys.stdout.write(chunk)
+                    sys.stdout.flush()
+                except Exception:
+                    pass
+
+            with spinner:
+                result = pipeline.run(user_prompt=user_input, stream_callback=_live_stream_callback)
 
             render_assistant_turn(
                 result,

{dbagent_cli-0.9.6 → dbagent_cli-0.9.8}/dbagent/connectors/relational.py

@@ -47,7 +47,7 @@ def _fuzzy_match_tables(table_names: List[str], user_prompt: str) -> Tuple[List[
     exact: List[str] = []
     fuzzy: List[str] = []
 
-    # Extract meaningful words from the prompt (3+ chars, skip SQL keywords)
+    # Extract meaningful words from the prompt (3+ chars, skip SQL keywords and diagram meta-terms)
     sql_keywords = {
         "select", "from", "where", "join", "left", "right", "inner", "outer",
         "order", "group", "having", "limit", "offset", "insert", "update",
@@ -57,6 +57,10 @@ def _fuzzy_match_tables(table_names: List[str], user_prompt: str) -> Tuple[List[
         "drop", "index", "primary", "foreign", "key", "null", "default",
         "hour", "minute", "day", "week", "month", "year", "ago", "today",
         "give", "tell", "want", "need", "please", "display", "fetch",
+        "flow", "diagram", "chart", "graph", "link", "links", "structure",
+        "design", "model", "blueprint", "mapping", "mappings", "overview",
+        "query", "code", "script", "report", "summary", "relation",
+        "relations", "relationship", "relationships", "architecture", "check",
     }
     prompt_words = set()
     for word in re.findall(r"\b[a-z_][a-z0-9_]*\b", prompt_lower):

{dbagent_cli-0.9.6 → dbagent_cli-0.9.8}/dbagent/ui/chat_screen.py

@@ -114,8 +114,15 @@ def render_assistant_turn(
         model_display = f" ({ai_model})" if ai_model else ""
         agent_title = f"[bold cyan]● DB-Agent{model_display}[/bold cyan] [dim]{timing_str}[/dim]"
     
-    # Conversational Chat Response
+    # Conversational Chat / Architectural Analysis Response
     if result.is_chat and result.message:
+        if result.was_streamed:
+            console.print("")
+            if result.ground_truth:
+                console.print(result.ground_truth.format_summary())
+            console.print("")
+            return
+
         console.print(f"\n{agent_title}")
         console.print(Markdown(result.message))
         if result.ground_truth:

{dbagent_cli-0.9.6 → dbagent_cli-0.9.8}/dbagent_cli.egg-info/PKG-INFO

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

{dbagent_cli-0.9.6 → dbagent_cli-0.9.8}/pyproject.toml

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

{dbagent_cli-0.9.6 → dbagent_cli-0.9.8}/tests/test_cli.py

@@ -481,6 +481,25 @@ def test_cli_changes_and_audit_commands():
     assert res_audit.exit_code == 0
 
 
+def test_pipeline_streaming_callback(sample_db: str):
+    from dbagent.connectors.factory import create_connector
+    from dbagent.llm.mock_provider import MockProvider
+    from dbagent.agent.pipeline import QueryPipeline
+
+    connector = create_connector(sample_db)
+    llm = MockProvider()
+    pipeline = QueryPipeline(connector=connector, llm=llm)
+
+    chunks = []
+    def callback(token: str):
+        chunks.append(token)
+
+    result = pipeline.run("hello assistant", stream_callback=callback)
+    assert result.was_streamed is True
+    assert result.is_chat is True
+    connector.close()
+
+
 def test_cli_recovery_commands():
     res_list = runner.invoke(app, ["recovery", "list"])
     assert res_list.exit_code == 0
@@ -491,5 +510,3 @@ def test_cli_recovery_commands():
 
     res_cleanup = runner.invoke(app, ["recovery", "cleanup"])
     assert res_cleanup.exit_code == 0
-
-