vibecore 0.3.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

vibecore/flow.py

@@ -1,105 +1,367 @@
 import asyncio
+import contextlib
+import datetime
+import sys
 import threading
 from collections.abc import Callable, Coroutine
-from typing import Protocol
+from typing import Any, Concatenate, Generic, TypeAlias, cast, overload
 
-from agents import Agent
+from agents import (
+    Agent,
+    RunConfig,
+    RunHooks,
+    Runner,
+    Session,
+    TResponseInputItem,
+)
+from agents.result import RunResultBase
+from agents.run import DEFAULT_MAX_TURNS
 from textual.pilot import Pilot
+from typing_extensions import TypeVar
 
-from vibecore.context import VibecoreContext
+from vibecore.context import AppAwareContext
 from vibecore.main import AppIsExiting, VibecoreApp
+from vibecore.session import JSONLSession
+from vibecore.settings import settings
 from vibecore.widgets.core import MyTextArea
 from vibecore.widgets.messages import SystemMessage
 
 
-class UserInputFunc(Protocol):
-    """Protocol for user input function with optional prompt parameter."""
+class NoUserInputLeft(Exception):
+    """Raised when no more user inputs are available in static runner."""
 
-    async def __call__(self, prompt: str = "") -> str:
-        """Get user input with optional prompt message.
+    pass
 
-        Args:
-            prompt: Optional prompt to display before getting input.
 
-        Returns:
-            The user's input string.
-        """
-        ...
+TContext = TypeVar("TContext", default=None)
+TWorkflowReturn = TypeVar("TWorkflowReturn", default=RunResultBase)
+
+
+class VibecoreRunner(Generic[TContext, TWorkflowReturn]):
+    def __init__(
+        self,
+        vibecore: "Vibecore[TContext, TWorkflowReturn]",
+        context: TContext | None = None,
+        session: Session | None = None,
+    ) -> None:
+        self.vibecore = vibecore
+        self.context = context
+
+        if session is None:
+            session_id = f"chat-{datetime.datetime.now().strftime('%Y%m%d-%H%M%S')}"
+            self._session = JSONLSession(
+                session_id=session_id,
+                project_path=None,  # Will use current working directory
+                base_dir=settings.session.base_dir,
+            )
+        else:
+            self._session = session
+
+    @property
+    def session(self) -> Session:
+        return self._session
+
+    async def print(self, message: str) -> None:
+        print(message, file=sys.stderr)
+
+    async def run_agent(
+        self,
+        starting_agent: Agent[TContext],
+        input: str | list[TResponseInputItem],
+        *,
+        context: TContext | None = None,
+        max_turns: int = DEFAULT_MAX_TURNS,
+        hooks: RunHooks[TContext] | None = None,
+        run_config: RunConfig | None = None,
+        previous_response_id: str | None = None,
+        session: Session | None = None,
+    ) -> RunResultBase:
+        result = await Runner.run(
+            starting_agent=starting_agent,
+            input=input,  # Pass string directly when using session
+            context=context,
+            max_turns=max_turns,
+            hooks=hooks,
+            run_config=run_config,
+            previous_response_id=previous_response_id,
+            session=session,
+        )
+        return result
+
+
+class VibecoreCliRunner(VibecoreRunner[TContext, TWorkflowReturn]):
+    def __init__(
+        self,
+        vibecore: "Vibecore[TContext, TWorkflowReturn]",
+        context: TContext | None = None,
+        session: Session | None = None,
+    ) -> None:
+        super().__init__(vibecore, context=context, session=session)
+
+    async def _user_input(self, prompt: str = "") -> str:
+        return input(prompt)
+
+    async def run(self) -> TWorkflowReturn:
+        assert self.vibecore.workflow_logic is not None, (
+            "Workflow logic not defined. Please use the @vibecore.workflow() decorator."
+        )
+        result = None
+        while user_message := await self._user_input():
+            result = await self.vibecore.workflow_logic(self, user_message)
 
+        assert result, "No result available after inputs exhausted."
+        return result
 
-async def flow(
-    agent: Agent,
-    logic: Callable[[VibecoreApp, VibecoreContext, UserInputFunc], Coroutine],
-    headless: bool = False,
-    shutdown: bool = False,
-    disable_user_input: bool = True,
-):
-    ctx = VibecoreContext()
-    app = VibecoreApp(ctx, agent, show_welcome=False)
 
-    app_ready_event = asyncio.Event()
+class VibecoreStaticRunner(VibecoreRunner[TContext, TWorkflowReturn]):
+    def __init__(
+        self,
+        vibecore: "Vibecore[TContext, TWorkflowReturn]",
+        context: TContext | None = None,
+        session: Session | None = None,
+    ) -> None:
+        super().__init__(vibecore, context=context, session=session)
+        self.inputs: list[str] = []
+        self.prints: list[str] = []
 
-    def on_app_ready() -> None:
+    async def print(self, message: str) -> None:
+        # Capture printed messages instead of displaying them
+        self.prints.append(message)
+
+    async def run(self, inputs: list[str] | None = None) -> TWorkflowReturn:
+        if inputs is None:
+            inputs = []
+        assert self.vibecore.workflow_logic is not None, (
+            "Workflow logic not defined. Please use the @vibecore.workflow() decorator."
+        )
+        result = None
+        for user_message in inputs:
+            result = await self.vibecore.workflow_logic(self, user_message)
+
+        assert result, "No result available after inputs exhausted."
+        return result
+
+
+class VibecoreTextualRunner(VibecoreRunner[AppAwareContext, TWorkflowReturn]):
+    def __init__(
+        self,
+        vibecore: "Vibecore[AppAwareContext, TWorkflowReturn]",
+        context: AppAwareContext | None = None,
+        session: Session | None = None,
+    ) -> None:
+        super().__init__(vibecore, context=context, session=session)
+        self.app = VibecoreApp(
+            self,
+            show_welcome=False,
+        )
+        self.app_ready_event = asyncio.Event()
+
+    async def _user_input(self, prompt: str = "") -> str:
+        if prompt:
+            await self.print(prompt)
+        self.app.query_one(MyTextArea).disabled = False
+        self.app.query_one(MyTextArea).focus()
+        user_input = await self.app.wait_for_user_input()
+        if self.vibecore.disable_user_input:
+            self.app.query_one(MyTextArea).disabled = True
+        return user_input
+
+    async def print(self, message: str) -> None:
+        await self.app.add_message(SystemMessage(message))
+
+    async def run_agent(
+        self,
+        starting_agent: Agent[AppAwareContext],
+        input: str | list[TResponseInputItem],
+        *,
+        context: AppAwareContext | None = None,
+        max_turns: int = DEFAULT_MAX_TURNS,
+        hooks: RunHooks[AppAwareContext] | None = None,
+        run_config: RunConfig | None = None,
+        previous_response_id: str | None = None,
+        session: Session | None = None,
+    ) -> RunResultBase:
+        result = Runner.run_streamed(
+            starting_agent=starting_agent,
+            input=input,  # Pass string directly when using session
+            context=context,
+            max_turns=max_turns,
+            hooks=hooks,
+            run_config=run_config,
+            previous_response_id=previous_response_id,
+            session=session,
+        )
+
+        self.app.current_worker = self.app.handle_streamed_response(result)
+        await self.app.current_worker.wait()
+        return result
+
+    def on_app_ready(self) -> None:
         """Called when app is ready to process events."""
-        app_ready_event.set()
+        self.app_ready_event.set()
 
-    async def run_app(app: VibecoreApp) -> None:
+    async def _run_app(self) -> None:
         """Run the apps message loop.
 
         Args:
             app: App to run.
         """
-
-        with app._context():
+        with self.app._context():
             try:
-                app._loop = asyncio.get_running_loop()
-                app._thread_id = threading.get_ident()
-                await app._process_messages(
-                    ready_callback=on_app_ready,
-                    headless=headless,
+                self.app._loop = asyncio.get_running_loop()
+                self.app._thread_id = threading.get_ident()
+                await self.app._process_messages(
+                    ready_callback=self.on_app_ready,
+                    headless=False,
                 )
             finally:
-                app_ready_event.set()
+                self.app_ready_event.set()
 
-    async def user_input(prompt: str = "") -> str:
-        if prompt:
-            await app.add_message(SystemMessage(prompt))
-        app.query_one(MyTextArea).disabled = False
-        app.query_one(MyTextArea).focus()
-        user_input = await app.wait_for_user_input()
-        if disable_user_input:
-            app.query_one(MyTextArea).disabled = True
-        return user_input
+    async def _run_logic(self) -> TWorkflowReturn:
+        assert self.vibecore.workflow_logic is not None, (
+            "Workflow logic not defined. Please use the @vibecore.workflow() decorator."
+        )
+        while True:
+            user_message = await self._user_input()
+            result = await self.vibecore.workflow_logic(self, user_message)
+        assert result, "No result available after inputs exhausted."
+        return result
+
+    async def run(self, inputs: list[str] | None = None, shutdown: bool = False) -> TWorkflowReturn:
+        if inputs:
+            self.app.message_queue.extend(inputs)
+        app_task = asyncio.create_task(self._run_app(), name=f"run_app({self.app})")
+        await self.app_ready_event.wait()
+
+        await self.app.load_session_history(self.session)
+        pilot = Pilot(self.app)
+        logic_task: asyncio.Task[TWorkflowReturn] | None = None
+
+        await pilot._wait_for_screen()
+        if self.vibecore.disable_user_input:
+            self.app.query_one(MyTextArea).disabled = True
+        logic_task = asyncio.create_task(self._run_logic(), name="logic_task")
+        done, pending = await asyncio.wait([logic_task, app_task], return_when=asyncio.FIRST_COMPLETED)
 
-    async def run_logic(app: VibecoreApp, ctx: VibecoreContext, user_input: UserInputFunc) -> None:
-        try:
-            await logic(app, ctx, user_input)
-        except AppIsExiting:
-            return
-
-    app_task = asyncio.create_task(run_app(app), name=f"with_app({app})")
-    await app_ready_event.wait()
-    pilot = Pilot(app)
-    logic_task: asyncio.Task | None = None
-
-    await pilot._wait_for_screen()
-    if disable_user_input:
-        app.query_one(MyTextArea).disabled = True
-    logic_task = asyncio.create_task(run_logic(app, ctx, user_input), name="logic_task")
-    done, pending = await asyncio.wait([logic_task, app_task], return_when=asyncio.FIRST_COMPLETED)
-
-    # If app has exited and logic is still running, cancel logic
-    if app_task in done and logic_task in pending:
-        logic_task.cancel()
-    # If logic is finished and app is still running
-    elif logic_task in done and app_task in pending:
-        if shutdown:
-            if not headless:
+        # If app has exited and logic is still running, cancel logic
+        if app_task in done and logic_task in pending:
+            logic_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await logic_task
+            raise AppIsExiting()
+        # If logic is finished and app is still running
+        elif logic_task in done and app_task in pending:
+            result = logic_task.result()
+            if shutdown:
                 await pilot._wait_for_screen()
                 await asyncio.sleep(1.0)
-            app.exit()
-        else:
-            # Enable text input so users can interact freely
-            app.query_one(MyTextArea).disabled = False
-            # Wait until app is exited
-            await app_task
+                self.app.exit()
+                await app_task
+            else:
+                await self.print("Workflow complete. Press Ctrl-Q to exit.")
+                # Enable text input so users can interact freely
+                self.app.query_one(MyTextArea).disabled = False
+                # Wait until app is exited
+                await app_task
+            return result
+        raise AssertionError(f"Unexpected state: done={done}, pending={pending}")
+
+
+WorkflowLogic: TypeAlias = Callable[
+    Concatenate[VibecoreRunner[TContext, TWorkflowReturn], str, ...],
+    Coroutine[Any, Any, TWorkflowReturn],
+]
+
+
+class Vibecore(Generic[TContext, TWorkflowReturn]):
+    def __init__(self, disable_user_input: bool = True) -> None:
+        self.workflow_logic: WorkflowLogic[TContext, TWorkflowReturn] | None = None
+        self.disable_user_input = disable_user_input
+
+    def workflow(
+        self,
+    ) -> Callable[[WorkflowLogic[TContext, TWorkflowReturn]], WorkflowLogic[TContext, TWorkflowReturn]]:
+        """Decorator to define the workflow logic for the app.
+
+        Returns:
+            A decorator that wraps the workflow logic function.
+        """
+
+        def decorator(
+            func: WorkflowLogic[TContext, TWorkflowReturn],
+        ) -> WorkflowLogic[TContext, TWorkflowReturn]:
+            self.workflow_logic = func
+            return func
+
+        return decorator
+
+    @overload
+    async def run_textual(
+        self,
+        inputs: str | None = None,
+        context: AppAwareContext | None = None,
+        session: Session | None = None,
+        shutdown: bool = False,
+    ) -> TWorkflowReturn: ...
+    @overload
+    async def run_textual(
+        self,
+        inputs: list[str] | None = None,
+        context: AppAwareContext | None = None,
+        session: Session | None = None,
+        shutdown: bool = False,
+    ) -> TWorkflowReturn: ...
+
+    async def run_textual(
+        self,
+        inputs: str | list[str] | None = None,
+        context: AppAwareContext | None = None,
+        session: Session | None = None,
+        shutdown: bool = False,
+    ) -> TWorkflowReturn:
+        if isinstance(inputs, str):
+            inputs = [inputs]
+
+        if self.workflow_logic is None:
+            raise ValueError("Workflow logic not defined. Please use the @vibecore.workflow() decorator.")
+
+        assert isinstance(context, AppAwareContext) or context is None, (
+            "Textual runner requires AppAwareContext or None."
+        )
+        # Type checker needs help: after the assertion, we know context is AppAwareContext | None
+        # and this Vibecore instance can be treated as Vibecore[AppAwareContext, TWorkflowReturn]
+        runner = VibecoreTextualRunner(
+            cast("Vibecore[AppAwareContext, TWorkflowReturn]", self),
+            context=context,
+            session=session,
+        )
+        return await runner.run(inputs=inputs, shutdown=shutdown)
+
+    async def run_cli(self, context: TContext | None = None, session: Session | None = None) -> TWorkflowReturn:
+        if self.workflow_logic is None:
+            raise ValueError("Workflow logic not defined. Please use the @vibecore.workflow() decorator.")
+
+        runner = VibecoreCliRunner(self, context=context, session=session)
+        return await runner.run()
+
+    @overload
+    async def run(
+        self, inputs: str, context: TContext | None = None, session: Session | None = None
+    ) -> TWorkflowReturn: ...
+
+    @overload
+    async def run(
+        self, inputs: list[str], context: TContext | None = None, session: Session | None = None
+    ) -> TWorkflowReturn: ...
+
+    async def run(
+        self, inputs: str | list[str], context: TContext | None = None, session: Session | None = None
+    ) -> TWorkflowReturn:
+        if isinstance(inputs, str):
+            inputs = [inputs]
+
+        if self.workflow_logic is None:
+            raise ValueError("Workflow logic not defined. Please use the @vibecore.workflow() decorator.")
+
+        runner = VibecoreStaticRunner(self, context=context, session=session)
+        return await runner.run(inputs=inputs)

vibecore/handlers/stream_handler.py

@@ -4,8 +4,6 @@ import json
 from typing import Protocol
 
 from agents import (
-    Agent,
-    AgentUpdatedStreamEvent,
     MessageOutputItem,
     RawResponsesStreamEvent,
     RunItemStreamEvent,
@@ -15,15 +13,12 @@ from agents import (
     ToolCallOutputItem,
 )
 from openai.types.responses import (
-    ResponseCompletedEvent,
     ResponseFunctionToolCall,
     ResponseOutputItemAddedEvent,
     ResponseOutputItemDoneEvent,
     ResponseReasoningItem,
-    ResponseReasoningSummaryPartAddedEvent,
-    ResponseReasoningSummaryTextDeltaEvent,
-    ResponseReasoningSummaryTextDoneEvent,
     ResponseTextDeltaEvent,
+    ResponseTextDoneEvent,
 )
 
 from vibecore.widgets.messages import (
@@ -43,8 +38,8 @@ class MessageHandler(Protocol):
         """Add a message to the widget's message list."""
         ...
 
-    async def handle_agent_update(self, new_agent: Agent) -> None:
-        """Handle agent updates."""
+    async def handle_agent_message_update(self, message: BaseMessage) -> None:
+        """Message in the widget's message list is updated with new delta or status"""
         ...
 
     async def handle_agent_error(self, error: Exception) -> None:
@@ -82,7 +77,15 @@ class AgentStreamHandler:
             self.agent_message = AgentMessage(self.message_content, status=MessageStatus.EXECUTING)
             await self.message_handler.handle_agent_message(self.agent_message)
         else:
-            self.agent_message.update(self.message_content)
+            # When content is short, we update more frequently for better UX
+            content_is_short_and_semantically_should_update = len(self.message_content) < 1000 and (
+                self.message_content.endswith(".") or "\n" in delta
+            )
+            # Else when content is long, we update less frequently to avoid UI lag
+            should_update_bulk_delta = len(self.message_content) - len(self.agent_message.text) > 200
+            if content_is_short_and_semantically_should_update or should_update_bulk_delta:
+                self.agent_message.update(self.message_content)
+                await self.message_handler.handle_agent_message_update(self.agent_message)
 
     async def handle_tool_call(self, tool_name: str, arguments: str, call_id: str) -> None:
         """Create and display tool message when tool is invoked.
@@ -108,11 +111,19 @@ class AgentStreamHandler:
         if call_id in self.tool_messages:
             tool_message = self.tool_messages[call_id]
             tool_message.update(MessageStatus.SUCCESS, str(output))
+            await self.message_handler.handle_agent_message_update(tool_message)
 
     async def handle_message_complete(self) -> None:
         """Finalize agent message when complete."""
+        if self.tool_messages:
+            # Some tool messages such as transfer_to_* may still be in executing status
+            # since it never gets a tool output event. We mark them as success here.
+            for tool_message in self.tool_messages.values():
+                if tool_message.status == MessageStatus.EXECUTING:
+                    tool_message.status = MessageStatus.SUCCESS
         if self.agent_message:
             self.agent_message.update(self.message_content, status=MessageStatus.IDLE)
+            await self.message_handler.handle_agent_message_update(self.agent_message)
             self.agent_message = None
             self.message_content = ""
 
@@ -132,42 +143,25 @@ class AgentStreamHandler:
                         self.reasoning_messages[reasoning_id] = reasoning_message
                         await self.message_handler.handle_agent_message(reasoning_message)
 
-                    case ResponseReasoningSummaryPartAddedEvent() as e:
-                        reasoning_id = e.item_id
-                        reasoning_message = self.reasoning_messages[reasoning_id]
-                        assert reasoning_message, f"Reasoning message with ID {reasoning_id} not found"
-                        updated = reasoning_message.text + "\n\n" if reasoning_message.text != "Thinking..." else ""
-                        reasoning_message.update(updated, status=MessageStatus.EXECUTING)
-
-                    case ResponseReasoningSummaryTextDeltaEvent(item_id=reasoning_id, delta=delta):
-                        reasoning_message = self.reasoning_messages[reasoning_id]
-                        assert reasoning_message, f"Reasoning message with ID {reasoning_id} not found"
-                        updated = reasoning_message.text + delta
-                        reasoning_message.update(updated, status=MessageStatus.EXECUTING)
-
-                    case ResponseReasoningSummaryTextDoneEvent(item_id=reasoning_id) as e:
-                        pass
-
                     case ResponseOutputItemDoneEvent(item=ResponseReasoningItem() as item):
                         reasoning_id = item.id
                         reasoning_message = self.reasoning_messages[reasoning_id]
                         assert reasoning_message, f"Reasoning message with ID {reasoning_id} not found"
-                        reasoning_message.update(reasoning_message.text, status=MessageStatus.IDLE)
-
-                    case ResponseOutputItemAddedEvent(item=ResponseFunctionToolCall(name=tool_name, call_id=call_id)):
-                        # TODO(serialx): We should move all tool call lifecycle start to here
-                        #                ResponseOutputItemDoneEvent can have race conditions in case of task tool calls
-                        #                Because we stream sub-agent events to our method: handle_task_tool_event()
-                        # XXX(serialx): Just handle task tool call lifecycle here for now
-                        #               I'm deferring this because `arguments` is not available here... need to refactor
-                        if tool_name == "task":
-                            await self.handle_tool_call(tool_name, "{}", call_id)
+                        text = "\n\n".join(summary.text for summary in item.summary)
+                        reasoning_message.update(text, status=MessageStatus.IDLE)
+                        await self.message_handler.handle_agent_message_update(reasoning_message)
 
                     case ResponseTextDeltaEvent(delta=delta) if delta:
                         await self.handle_text_delta(delta)
 
-                    case ResponseOutputItemDoneEvent(
-                        item=ResponseFunctionToolCall(name=tool_name, arguments=arguments, call_id=call_id)
+                    case ResponseTextDoneEvent() as e:
+                        self.agent_message = AgentMessage(e.text, status=MessageStatus.IDLE)
+                        await self.message_handler.handle_agent_message(self.agent_message)
+
+                    case (
+                        ResponseOutputItemDoneEvent(
+                            item=ResponseFunctionToolCall(name=tool_name, arguments=arguments, call_id=call_id)
+                        ) as e
                     ):
                         # XXX(serialx): See above comments
                         if tool_name == "task":
@@ -182,33 +176,18 @@ class AgentStreamHandler:
                         else:
                             await self.handle_tool_call(tool_name, arguments, call_id)
 
-                    case ResponseCompletedEvent():
-                        # When in agent handoff or stop at tool situations, the tools should be in executing status.
-                        # We find all the executing status tool messages and mark them as success.
-                        for tool_message in self.tool_messages.values():
-                            if tool_message.status == MessageStatus.EXECUTING:
-                                tool_message.status = MessageStatus.SUCCESS
-
             case RunItemStreamEvent(item=item):
                 # log(f"RunItemStreamEvent item: {item.type}")
                 match item:
-                    case ToolCallItem():
-                        pass
-                    case ToolCallOutputItem(output=output, raw_item=raw_item):
+                    case ToolCallItem(call_output=ResponseFunctionToolCall() as call):
+                        await self.handle_tool_call(call.name, call.arguments, call.call_id)
+                    case ToolCallOutputItem(raw_item=raw_item, output=output):
                         # Find the corresponding tool message by call_id
-                        if (
-                            isinstance(raw_item, dict)
-                            and "call_id" in raw_item
-                            and raw_item["call_id"] in self.tool_messages
-                        ):
+                        if raw_item["type"] == "function_call_output" and raw_item["call_id"] in self.tool_messages:
                             await self.handle_tool_output(output, raw_item["call_id"])
                     case MessageOutputItem():
                         await self.handle_message_complete()
 
-            case AgentUpdatedStreamEvent(new_agent=new_agent):
-                # log(f"AgentUpdatedStreamEvent new_agent: {new_agent.name}")
-                await self.message_handler.handle_agent_update(new_agent)
-
     async def handle_task_tool_event(self, tool_name: str, tool_call_id: str, event: StreamEvent) -> None:
         """Handle streaming events from task tool sub-agents.