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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.9.7
+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.7 → 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.7"
+__version__ = "0.9.8"
 __author__ = "Santhosh Gupta"

{dbagent_cli-0.9.7 → 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.7 → 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"
@@ -273,7 +283,10 @@ class QueryPipeline:
                     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.7 → 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.7 → 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.7 → 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.7
+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.7 → dbagent_cli-0.9.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dbagent-cli"
-version = "0.9.7"
+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.7 → 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
-
-