Support custom conversation load/store paths and optional --continue

Author: uriyyo

clai/README.md

@@ -54,7 +54,7 @@ Either way, running `clai` will start an interactive session where you can chat
 ## Help
 
 ```
-usage: clai [-h] [-m [MODEL]] [-a AGENT] [-l] [-t [CODE_THEME]] [-c] [--no-stream] [--version] [prompt]
+usage: clai [-h] [-m [MODEL]] [-a AGENT] [-l] [-t [CODE_THEME]] [-c [CONTINUE_]] [--store STORE] [--no-stream] [--version] [prompt]
 
 Pydantic AI CLI v...
 
@@ -76,7 +76,9 @@ options:
   -l, --list-models     List all available models and exit
   -t [CODE_THEME], --code-theme [CODE_THEME]
                         Which colors to use for code, can be "dark", "light" or any theme from pygments.org/styles/. Defaults to "dark" which works well on dark terminals.
-  -c, --continue        Continue last conversation, if any, instead of starting a new one.
+  -c [CONTINUE_], --continue [CONTINUE_]
+                        Continue last conversation, if any, instead of starting a new one.
+  --store STORE         Store the last conversation to the specified path instead of the default location.
   --no-stream           Disable streaming from the model
   --version             Show version and exit
 ```

pydantic_ai_slim/pydantic_ai/_cli.py

@@ -153,10 +153,17 @@ def cli(  # noqa: C901
     parser.add_argument(
         '-c',
         '--continue',
+        nargs='?',
         dest='continue_',
-        action='store_true',
+        const=str(PYDANTIC_AI_HOME / LAST_CONVERSATION_FILENAME),
+        default=None,
         help='Continue last conversation, if any, instead of starting a new one.',
     )
+    parser.add_argument(
+        '--store',
+        help='Store the last conversation to the specified path instead of the default location.',
+        default=None,
+    )
     parser.add_argument('--no-stream', action='store_true', help='Disable streaming from the model')
     parser.add_argument('--version', action='store_true', help='Show version and exit')
 
@@ -216,40 +223,47 @@ def cli(  # noqa: C901
     else:
         code_theme = args.code_theme  # pragma: no cover
 
+    load_path: Path | None = None
+    if args.continue_:
+        load_path = Path(args.continue_)
+
+    store_path: Path = PYDANTIC_AI_HOME / LAST_CONVERSATION_FILENAME
+    if args.store:
+        store_path = Path(args.store)
+
     try:
-        history = load_last_conversation() if args.continue_ else None
+        history = load_conversation(load_path) if load_path else None
     except ValidationError:
         console.print(
-            '[red]Error loading last conversation, it is corrupted or invalid.\nStarting a new conversation.[/red]'
+            '[red]Error loading conversation, it is corrupted or invalid.\nStarting a new conversation.[/red]'
         )
         history = None
 
     if prompt := cast(str, args.prompt):
         try:
-            asyncio.run(ask_agent(agent, prompt, stream, console, code_theme, messages=history))
+            asyncio.run(ask_agent(agent, prompt, stream, console, code_theme, messages=history, store_path=store_path))
         except KeyboardInterrupt:
             pass
         return 0
 
     try:
-        return asyncio.run(run_chat(stream, agent, console, code_theme, prog_name, message_history=history))
+        return asyncio.run(
+            run_chat(stream, agent, console, code_theme, prog_name, message_history=history, store_path=store_path)
+        )
     except KeyboardInterrupt:  # pragma: no cover
         return 0
 
 
-def store_last_conversation(messages: list[ModelMessage], config_dir: Path | None = None) -> None:
-    last_conversation_path = (config_dir or PYDANTIC_AI_HOME) / LAST_CONVERSATION_FILENAME
-    last_conversation_path.parent.mkdir(parents=True, exist_ok=True)
-    last_conversation_path.write_bytes(ModelMessagesTypeAdapter.dump_json(messages))
-
+def store_conversation(messages: list[ModelMessage], store_path: Path) -> None:
+    store_path.parent.mkdir(parents=True, exist_ok=True)
+    store_path.write_bytes(ModelMessagesTypeAdapter.dump_json(messages))
 
-def load_last_conversation(config_dir: Path | None = None) -> list[ModelMessage] | None:
-    last_conversation_path = (config_dir or PYDANTIC_AI_HOME) / LAST_CONVERSATION_FILENAME
 
-    if not last_conversation_path.exists():
+def load_conversation(load_path: Path) -> list[ModelMessage] | None:
+    if not load_path.exists():
         return None
 
-    return ModelMessagesTypeAdapter.validate_json(last_conversation_path.read_bytes())
+    return ModelMessagesTypeAdapter.validate_json(load_path.read_bytes())
 
 
 async def run_chat(
@@ -261,6 +275,7 @@ async def run_chat(
     config_dir: Path | None = None,
     deps: AgentDepsT = None,
     message_history: Sequence[ModelMessage] | None = None,
+    store_path: Path | None = None,
 ) -> int:
     prompt_history_path = (config_dir or PYDANTIC_AI_HOME) / PROMPT_HISTORY_FILENAME
     prompt_history_path.parent.mkdir(parents=True, exist_ok=True)
@@ -287,7 +302,7 @@ async def run_chat(
                 return exit_value
         else:
             try:
-                messages = await ask_agent(agent, text, stream, console, code_theme, deps, messages, config_dir)
+                messages = await ask_agent(agent, text, stream, console, code_theme, deps, messages, store_path)
             except CancelledError:  # pragma: no cover
                 console.print('[dim]Interrupted[/dim]')
             except Exception as e:  # pragma: no cover
@@ -305,7 +320,7 @@ async def ask_agent(
     code_theme: str,
     deps: AgentDepsT = None,
     messages: Sequence[ModelMessage] | None = None,
-    config_dir: Path | None = None,
+    store_path: Path | None = None,
 ) -> list[ModelMessage]:
     status = Status('[dim]Working on it…[/dim]', console=console)
 
@@ -314,28 +329,28 @@ async def ask_agent(
             result = await agent.run(prompt, message_history=messages, deps=deps)
         content = str(result.output)
         console.print(Markdown(content, code_theme=code_theme))
-        result_messages = result.all_messages()
-        store_last_conversation(result_messages, config_dir)
+    else:
+        with status, ExitStack() as stack:
+            async with agent.iter(prompt, message_history=messages, deps=deps) as agent_run:
+                live = Live('', refresh_per_second=15, console=console, vertical_overflow='ellipsis')
+                async for node in agent_run:
+                    if Agent.is_model_request_node(node):
+                        async with node.stream(agent_run.ctx) as handle_stream:
+                            status.stop()  # stopping multiple times is idempotent
+                            stack.enter_context(live)  # entering multiple times is idempotent
 
-        return result_messages
+                            async for content in handle_stream.stream_output(debounce_by=None):
+                                live.update(Markdown(str(content), code_theme=code_theme))
 
-    with status, ExitStack() as stack:
-        async with agent.iter(prompt, message_history=messages, deps=deps) as agent_run:
-            live = Live('', refresh_per_second=15, console=console, vertical_overflow='ellipsis')
-            async for node in agent_run:
-                if Agent.is_model_request_node(node):
-                    async with node.stream(agent_run.ctx) as handle_stream:
-                        status.stop()  # stopping multiple times is idempotent
-                        stack.enter_context(live)  # entering multiple times is idempotent
+            assert agent_run.result is not None
+            result = agent_run.result
 
-                        async for content in handle_stream.stream_output(debounce_by=None):
-                            live.update(Markdown(str(content), code_theme=code_theme))
+    result_messages = result.all_messages()
 
-        assert agent_run.result is not None
-        result_messages = agent_run.result.all_messages()
-        store_last_conversation(result_messages, config_dir)
+    if store_path:
+        store_conversation(result_messages, store_path)
 
-        return result_messages
+    return result_messages
 
 
 class CustomAutoSuggest(AutoSuggestFromHistory):