vibecore 0.5.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

vibecore/flow.py

@@ -4,7 +4,7 @@ import datetime
 import sys
 import threading
 from collections.abc import Callable, Coroutine
-from typing import Any, Generic, Protocol, TypeAlias, overload
+from typing import Any, Concatenate, Generic, TypeAlias, cast, overload
 
 from agents import (
     Agent,
@@ -12,7 +12,6 @@ from agents import (
     RunHooks,
     Runner,
     Session,
-    TContext,
     TResponseInputItem,
 )
 from agents.result import RunResultBase
@@ -20,7 +19,7 @@ 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
@@ -28,35 +27,39 @@ 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.
-
-        Args:
-            prompt: Optional prompt to display before getting input.
-
-        Returns:
-            The user's input string.
-        """
-        ...
+    pass
 
 
+TContext = TypeVar("TContext", default=None)
 TWorkflowReturn = TypeVar("TWorkflowReturn", default=RunResultBase)
-DecoratedCallable: TypeAlias = Callable[..., Coroutine[Any, Any, TWorkflowReturn]]
 
 
-class VibecoreRunnerBase(Generic[TWorkflowReturn]):
-    def __init__(self, vibecore: "Vibecore[TWorkflowReturn]") -> None:
+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:
-        raise NotImplementedError("session property implemented.")
-
-    async def user_input(self, prompt: str = "") -> str:
-        raise NotImplementedError("user_input method not implemented.")
+        return self._session
 
     async def print(self, message: str) -> None:
         print(message, file=sys.stderr)
@@ -86,46 +89,41 @@ class VibecoreRunnerBase(Generic[TWorkflowReturn]):
         return result
 
 
-class VibecoreSimpleRunner(VibecoreRunnerBase[TWorkflowReturn]):
-    def __init__(self, vibecore: "Vibecore[TWorkflowReturn]") -> None:
-        super().__init__(vibecore)
-
-        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,
-        )
-
-    @property
-    def session(self) -> Session:
-        return self._session
-
-
-class VibecoreCliRunner(VibecoreSimpleRunner[TWorkflowReturn]):
-    def __init__(self, vibecore: "Vibecore[TWorkflowReturn]") -> None:
-        super().__init__(vibecore)
+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:
+    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."
         )
-        return await self.vibecore.workflow_logic()
+        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
 
 
-class VibecoreStaticRunner(VibecoreSimpleRunner[TWorkflowReturn]):
-    def __init__(self, vibecore: "Vibecore[TWorkflowReturn]") -> None:
-        super().__init__(vibecore)
+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] = []
 
-    async def user_input(self, prompt: str = "") -> str:
-        assert self.inputs, "No more user inputs available."
-        return self.inputs.pop()
-
     async def print(self, message: str) -> None:
         # Capture printed messages instead of displaying them
         self.prints.append(message)
@@ -136,21 +134,29 @@ class VibecoreStaticRunner(VibecoreSimpleRunner[TWorkflowReturn]):
         assert self.vibecore.workflow_logic is not None, (
             "Workflow logic not defined. Please use the @vibecore.workflow() decorator."
         )
-        self.inputs.extend(inputs)
-        return await self.vibecore.workflow_logic()
+        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(VibecoreRunnerBase[TWorkflowReturn]):
-    def __init__(self, vibecore: "Vibecore[TWorkflowReturn]") -> None:
-        super().__init__(vibecore)
-        self.app = VibecoreApp(self.vibecore.context, self.vibecore.starting_agent, show_welcome=False)
-        self.app_ready_event = asyncio.Event()
 
-    @property
-    def session(self) -> Session:
-        return self.app.session
+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:
+    async def _user_input(self, prompt: str = "") -> str:
         if prompt:
             await self.print(prompt)
         self.app.query_one(MyTextArea).disabled = False
@@ -165,12 +171,12 @@ class VibecoreTextualRunner(VibecoreRunnerBase[TWorkflowReturn]):
 
     async def run_agent(
         self,
-        starting_agent: Agent[TContext],
+        starting_agent: Agent[AppAwareContext],
         input: str | list[TResponseInputItem],
         *,
-        context: TContext | None = None,
+        context: AppAwareContext | None = None,
         max_turns: int = DEFAULT_MAX_TURNS,
-        hooks: RunHooks[TContext] | None = None,
+        hooks: RunHooks[AppAwareContext] | None = None,
         run_config: RunConfig | None = None,
         previous_response_id: str | None = None,
         session: Session | None = None,
@@ -200,7 +206,6 @@ class VibecoreTextualRunner(VibecoreRunnerBase[TWorkflowReturn]):
         Args:
             app: App to run.
         """
-
         with self.app._context():
             try:
                 self.app._loop = asyncio.get_running_loop()
@@ -216,15 +221,19 @@ class VibecoreTextualRunner(VibecoreRunnerBase[TWorkflowReturn]):
         assert self.vibecore.workflow_logic is not None, (
             "Workflow logic not defined. Please use the @vibecore.workflow() decorator."
         )
-        try:
-            return await self.vibecore.workflow_logic()
-        except AppIsExiting:
-            raise
+        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, shutdown: bool = False) -> TWorkflowReturn:
-        self.app = VibecoreApp(self.vibecore.context, self.vibecore.starting_agent, show_welcome=False)
+    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
 
@@ -249,34 +258,29 @@ class VibecoreTextualRunner(VibecoreRunnerBase[TWorkflowReturn]):
                 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 RuntimeError("Unexpected state: both tasks completed")
+        raise AssertionError(f"Unexpected state: done={done}, pending={pending}")
 
 
-class Vibecore(Generic[TWorkflowReturn]):
-    def __init__(self, starting_agent: Agent[TContext], disable_user_input: bool = True) -> None:
-        self.context = VibecoreContext()
-        self.workflow_logic: Callable[..., Coroutine[Any, Any, TWorkflowReturn]] | None = None
-        self.starting_agent = starting_agent
-        self.disable_user_input = disable_user_input
-        self.runner: VibecoreRunnerBase[TWorkflowReturn] = VibecoreRunnerBase(self)
-
-    @property
-    def session(self) -> Session:
-        return self.runner.session
+WorkflowLogic: TypeAlias = Callable[
+    Concatenate[VibecoreRunner[TContext, TWorkflowReturn], str, ...],
+    Coroutine[Any, Any, TWorkflowReturn],
+]
 
-    async def user_input(self, prompt: str = "") -> str:
-        return await self.runner.user_input(prompt)
 
-    async def print(self, message: str) -> None:
-        return await self.runner.print(message)
+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[[DecoratedCallable[TWorkflowReturn]], DecoratedCallable[TWorkflowReturn]]:
+    def workflow(
+        self,
+    ) -> Callable[[WorkflowLogic[TContext, TWorkflowReturn]], WorkflowLogic[TContext, TWorkflowReturn]]:
         """Decorator to define the workflow logic for the app.
 
         Returns:
@@ -284,62 +288,80 @@ class Vibecore(Generic[TWorkflowReturn]):
         """
 
         def decorator(
-            func: DecoratedCallable[TWorkflowReturn],
-        ) -> DecoratedCallable[TWorkflowReturn]:
+            func: WorkflowLogic[TContext, TWorkflowReturn],
+        ) -> WorkflowLogic[TContext, TWorkflowReturn]:
             self.workflow_logic = func
             return func
 
         return decorator
 
-    async def run_agent(
+    @overload
+    async def run_textual(
         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,
+        inputs: str | None = None,
+        context: AppAwareContext | None = None,
         session: Session | None = None,
-    ) -> RunResultBase:
-        return await self.runner.run_agent(
-            starting_agent=starting_agent,
-            input=input,
-            context=context,
-            max_turns=max_turns,
-            hooks=hooks,
-            run_config=run_config,
-            previous_response_id=previous_response_id,
-            session=session,
-        )
+        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]
 
-    async def run_textual(self, shutdown: bool = False) -> TWorkflowReturn:
         if self.workflow_logic is None:
             raise ValueError("Workflow logic not defined. Please use the @vibecore.workflow() decorator.")
 
-        self.runner = VibecoreTextualRunner(self)
-        return await self.runner.run(shutdown=shutdown)
+        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) -> TWorkflowReturn:
+    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.")
 
-        self.runner = VibecoreCliRunner(self)
-        return await self.runner.run()
+        runner = VibecoreCliRunner(self, context=context, session=session)
+        return await runner.run()
 
     @overload
-    async def run(self, inputs: str) -> TWorkflowReturn: ...
+    async def run(
+        self, inputs: str, context: TContext | None = None, session: Session | None = None
+    ) -> TWorkflowReturn: ...
 
     @overload
-    async def run(self, inputs: list[str]) -> TWorkflowReturn: ...
+    async def run(
+        self, inputs: list[str], context: TContext | None = None, session: Session | None = None
+    ) -> TWorkflowReturn: ...
 
-    async def run(self, inputs: str | list[str]) -> 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.")
 
-        self.runner = VibecoreStaticRunner(self)
-        return await self.runner.run(inputs=inputs)
+        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,
@@ -43,10 +41,6 @@ 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_error(self, error: Exception) -> None:
         """Handle errors during streaming."""
         ...
@@ -205,10 +199,6 @@ class AgentStreamHandler:
                     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.