prompture 0.0.35__py3-none-any.whl → 0.0.40.dev1__py3-none-any.whl

prompture/async_conversation.py

@@ -4,9 +4,11 @@ from __future__ import annotations
 
 import json
 import logging
+import uuid
 from collections.abc import AsyncIterator
-from datetime import date, datetime
+from datetime import date, datetime, timezone
 from decimal import Decimal
+from pathlib import Path
 from typing import Any, Callable, Literal, Union
 
 from pydantic import BaseModel
@@ -15,6 +17,11 @@ from .async_driver import AsyncDriver
 from .callbacks import DriverCallbacks
 from .drivers.async_registry import get_async_driver_for_model
 from .field_definitions import get_registry_snapshot
+from .image import ImageInput, make_image
+from .persistence import load_from_file, save_to_file
+from .persona import Persona, get_persona
+from .serialization import export_conversation, import_conversation
+from .session import UsageSession
 from .tools import (
     clean_json_text,
     convert_value,
@@ -43,13 +50,33 @@ class AsyncConversation:
         *,
         driver: AsyncDriver | None = None,
         system_prompt: str | None = None,
+        persona: str | Persona | None = None,
         options: dict[str, Any] | None = None,
         callbacks: DriverCallbacks | None = None,
         tools: ToolRegistry | None = None,
         max_tool_rounds: int = 10,
+        conversation_id: str | None = None,
+        auto_save: str | Path | None = None,
+        tags: list[str] | None = None,
     ) -> None:
+        if system_prompt is not None and persona is not None:
+            raise ValueError("Cannot provide both 'system_prompt' and 'persona'. Use one or the other.")
+
+        # Resolve persona
+        resolved_persona: Persona | None = None
+        if persona is not None:
+            if isinstance(persona, str):
+                resolved_persona = get_persona(persona)
+                if resolved_persona is None:
+                    raise ValueError(f"Persona '{persona}' not found in registry.")
+            else:
+                resolved_persona = persona
+
         if model_name is None and driver is None:
-            raise ValueError("Either model_name or driver must be provided")
+            if resolved_persona is not None and resolved_persona.model_hint:
+                model_name = resolved_persona.model_hint
+            else:
+                raise ValueError("Either model_name or driver must be provided")
 
         if driver is not None:
             self._driver = driver
@@ -60,8 +87,15 @@ class AsyncConversation:
             self._driver.callbacks = callbacks
 
         self._model_name = model_name or ""
-        self._system_prompt = system_prompt
-        self._options = dict(options) if options else {}
+
+        # Apply persona: render system_prompt and merge settings
+        if resolved_persona is not None:
+            self._system_prompt = resolved_persona.render()
+            self._options = {**resolved_persona.settings, **(dict(options) if options else {})}
+        else:
+            self._system_prompt = system_prompt
+            self._options = dict(options) if options else {}
+
         self._messages: list[dict[str, Any]] = []
         self._usage = {
             "prompt_tokens": 0,
@@ -73,6 +107,14 @@ class AsyncConversation:
         self._tools = tools or ToolRegistry()
         self._max_tool_rounds = max_tool_rounds
 
+        # Persistence
+        self._conversation_id = conversation_id or str(uuid.uuid4())
+        self._auto_save = Path(auto_save) if auto_save else None
+        self._metadata: dict[str, Any] = {
+            "created_at": datetime.now(timezone.utc).isoformat(),
+            "tags": list(tags) if tags else [],
+        }
+
     # ------------------------------------------------------------------
     # Public helpers
     # ------------------------------------------------------------------
@@ -91,11 +133,12 @@ class AsyncConversation:
         """Reset message history (keeps system_prompt and driver)."""
         self._messages.clear()
 
-    def add_context(self, role: str, content: str) -> None:
+    def add_context(self, role: str, content: str, images: list[ImageInput] | None = None) -> None:
         """Seed the history with a user or assistant message."""
         if role not in ("user", "assistant"):
             raise ValueError("role must be 'user' or 'assistant'")
-        self._messages.append({"role": role, "content": content})
+        msg_content = self._build_content_with_images(content, images)
+        self._messages.append({"role": role, "content": msg_content})
 
     def register_tool(
         self,
@@ -112,17 +155,145 @@ class AsyncConversation:
         u = self._usage
         return f"Conversation: {u['total_tokens']:,} tokens across {u['turns']} turn(s) costing ${u['cost']:.4f}"
 
+    # ------------------------------------------------------------------
+    # Persistence properties
+    # ------------------------------------------------------------------
+
+    @property
+    def conversation_id(self) -> str:
+        """Unique identifier for this conversation."""
+        return self._conversation_id
+
+    @property
+    def tags(self) -> list[str]:
+        """Tags attached to this conversation."""
+        return self._metadata.get("tags", [])
+
+    @tags.setter
+    def tags(self, value: list[str]) -> None:
+        self._metadata["tags"] = list(value)
+
+    # ------------------------------------------------------------------
+    # Export / Import
+    # ------------------------------------------------------------------
+
+    def export(self, *, usage_session: UsageSession | None = None, strip_images: bool = False) -> dict[str, Any]:
+        """Export conversation state to a JSON-serializable dict."""
+        tools_metadata = (
+            [
+                {"name": td.name, "description": td.description, "parameters": td.parameters}
+                for td in self._tools.definitions
+            ]
+            if self._tools and self._tools.definitions
+            else None
+        )
+        return export_conversation(
+            model_name=self._model_name,
+            system_prompt=self._system_prompt,
+            options=self._options,
+            messages=self._messages,
+            usage=self._usage,
+            max_tool_rounds=self._max_tool_rounds,
+            tools_metadata=tools_metadata,
+            usage_session=usage_session,
+            metadata=self._metadata,
+            conversation_id=self._conversation_id,
+            strip_images=strip_images,
+        )
+
+    @classmethod
+    def from_export(
+        cls,
+        data: dict[str, Any],
+        *,
+        callbacks: DriverCallbacks | None = None,
+        tools: ToolRegistry | None = None,
+    ) -> AsyncConversation:
+        """Reconstruct an :class:`AsyncConversation` from an export dict.
+
+        The driver is reconstructed from the stored ``model_name`` using
+        :func:`get_async_driver_for_model`.  Callbacks and tool functions
+        must be re-attached by the caller.
+        """
+        imported = import_conversation(data)
+
+        model_name = imported.get("model_name") or ""
+        if not model_name:
+            raise ValueError("Cannot restore conversation: export has no model_name")
+        conv = cls(
+            model_name=model_name,
+            system_prompt=imported.get("system_prompt"),
+            options=imported.get("options", {}),
+            callbacks=callbacks,
+            tools=tools,
+            max_tool_rounds=imported.get("max_tool_rounds", 10),
+            conversation_id=imported.get("conversation_id"),
+            tags=imported.get("metadata", {}).get("tags", []),
+        )
+        conv._messages = imported.get("messages", [])
+        conv._usage = imported.get(
+            "usage",
+            {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0,
+                "cost": 0.0,
+                "turns": 0,
+            },
+        )
+        meta = imported.get("metadata", {})
+        if "created_at" in meta:
+            conv._metadata["created_at"] = meta["created_at"]
+        return conv
+
+    def save(self, path: str | Path, **kwargs: Any) -> None:
+        """Export and write to a JSON file."""
+        save_to_file(self.export(**kwargs), path)
+
+    @classmethod
+    def load(
+        cls,
+        path: str | Path,
+        *,
+        callbacks: DriverCallbacks | None = None,
+        tools: ToolRegistry | None = None,
+    ) -> AsyncConversation:
+        """Load a conversation from a JSON file."""
+        data = load_from_file(path)
+        return cls.from_export(data, callbacks=callbacks, tools=tools)
+
+    def _maybe_auto_save(self) -> None:
+        """Auto-save after each turn if configured."""
+        if self._auto_save is None:
+            return
+        try:
+            self.save(self._auto_save)
+        except Exception:
+            logger.debug("Auto-save failed for conversation %s", self._conversation_id, exc_info=True)
+
     # ------------------------------------------------------------------
     # Core methods
     # ------------------------------------------------------------------
 
-    def _build_messages(self, user_content: str) -> list[dict[str, Any]]:
+    @staticmethod
+    def _build_content_with_images(text: str, images: list[ImageInput] | None = None) -> str | list[dict[str, Any]]:
+        """Return plain string when no images, or a list of content blocks."""
+        if not images:
+            return text
+        blocks: list[dict[str, Any]] = [{"type": "text", "text": text}]
+        for img in images:
+            ic = make_image(img)
+            blocks.append({"type": "image", "source": ic})
+        return blocks
+
+    def _build_messages(self, user_content: str, images: list[ImageInput] | None = None) -> list[dict[str, Any]]:
         """Build the full messages array for an API call."""
         msgs: list[dict[str, Any]] = []
         if self._system_prompt:
             msgs.append({"role": "system", "content": self._system_prompt})
         msgs.extend(self._messages)
-        msgs.append({"role": "user", "content": user_content})
+        content = self._build_content_with_images(user_content, images)
+        msgs.append({"role": "user", "content": content})
         return msgs
 
     def _accumulate_usage(self, meta: dict[str, Any]) -> None:
@@ -131,11 +302,22 @@ class AsyncConversation:
         self._usage["total_tokens"] += meta.get("total_tokens", 0)
         self._usage["cost"] += meta.get("cost", 0.0)
         self._usage["turns"] += 1
+        self._maybe_auto_save()
+
+        from .ledger import _resolve_api_key_hash, record_model_usage
+
+        record_model_usage(
+            self._model_name,
+            api_key_hash=_resolve_api_key_hash(self._model_name),
+            tokens=meta.get("total_tokens", 0),
+            cost=meta.get("cost", 0.0),
+        )
 
     async def ask(
         self,
         content: str,
         options: dict[str, Any] | None = None,
+        images: list[ImageInput] | None = None,
     ) -> str:
         """Send a message and get a raw text response (async).
 
@@ -143,16 +325,17 @@ class AsyncConversation:
         dispatches to the async tool execution loop.
         """
         if self._tools and getattr(self._driver, "supports_tool_use", False):
-            return await self._ask_with_tools(content, options)
+            return await self._ask_with_tools(content, options, images=images)
 
         merged = {**self._options, **(options or {})}
-        messages = self._build_messages(content)
+        messages = self._build_messages(content, images=images)
         resp = await self._driver.generate_messages_with_hooks(messages, merged)
 
         text = resp.get("text", "")
         meta = resp.get("meta", {})
 
-        self._messages.append({"role": "user", "content": content})
+        user_content = self._build_content_with_images(content, images)
+        self._messages.append({"role": "user", "content": user_content})
         self._messages.append({"role": "assistant", "content": text})
         self._accumulate_usage(meta)
 
@@ -162,12 +345,14 @@ class AsyncConversation:
         self,
         content: str,
         options: dict[str, Any] | None = None,
+        images: list[ImageInput] | None = None,
     ) -> str:
         """Async tool-use loop: send -> check tool_calls -> execute -> re-send."""
         merged = {**self._options, **(options or {})}
         tool_defs = self._tools.to_openai_format()
 
-        self._messages.append({"role": "user", "content": content})
+        user_content = self._build_content_with_images(content, images)
+        self._messages.append({"role": "user", "content": user_content})
         msgs = self._build_messages_raw()
 
         for _round in range(self._max_tool_rounds):
@@ -228,6 +413,7 @@ class AsyncConversation:
         self,
         content: str,
         options: dict[str, Any] | None = None,
+        images: list[ImageInput] | None = None,
     ) -> AsyncIterator[str]:
         """Send a message and yield text chunks as they arrive (async).
 
@@ -235,13 +421,14 @@ class AsyncConversation:
         support streaming.
         """
         if not getattr(self._driver, "supports_streaming", False):
-            yield await self.ask(content, options)
+            yield await self.ask(content, options, images=images)
             return
 
         merged = {**self._options, **(options or {})}
-        messages = self._build_messages(content)
+        messages = self._build_messages(content, images=images)
 
-        self._messages.append({"role": "user", "content": content})
+        user_content = self._build_content_with_images(content, images)
+        self._messages.append({"role": "user", "content": user_content})
 
         full_text = ""
         async for chunk in self._driver.generate_messages_stream(messages, merged):
@@ -267,6 +454,7 @@ class AsyncConversation:
         options: dict[str, Any] | None = None,
         output_format: Literal["json", "toon"] = "json",
         json_mode: Literal["auto", "on", "off"] = "auto",
+        images: list[ImageInput] | None = None,
     ) -> dict[str, Any]:
         """Send a message with schema enforcement and get structured JSON back (async)."""
         merged = {**self._options, **(options or {})}
@@ -301,13 +489,14 @@ class AsyncConversation:
 
         full_user_content = f"{content}\n\n{instruct}"
 
-        messages = self._build_messages(full_user_content)
+        messages = self._build_messages(full_user_content, images=images)
         resp = await self._driver.generate_messages_with_hooks(messages, merged)
 
         text = resp.get("text", "")
         meta = resp.get("meta", {})
 
-        self._messages.append({"role": "user", "content": content})
+        user_content = self._build_content_with_images(content, images)
+        self._messages.append({"role": "user", "content": user_content})
 
         cleaned = clean_json_text(text)
         try:
@@ -361,6 +550,7 @@ class AsyncConversation:
         output_format: Literal["json", "toon"] = "json",
         options: dict[str, Any] | None = None,
         json_mode: Literal["auto", "on", "off"] = "auto",
+        images: list[ImageInput] | None = None,
     ) -> dict[str, Any]:
         """Extract structured information into a Pydantic model with conversation context (async)."""
         from .core import normalize_field_value
@@ -375,6 +565,7 @@ class AsyncConversation:
             options=options,
             output_format=output_format,
             json_mode=json_mode,
+            images=images,
         )
 
         json_object = result["json_object"]

prompture/async_core.py

@@ -35,6 +35,18 @@ from .tools import (
 logger = logging.getLogger("prompture.async_core")
 
 
+def _record_usage_to_ledger(model_name: str, meta: dict[str, Any]) -> None:
+    """Fire-and-forget ledger recording for standalone async core functions."""
+    from .ledger import _resolve_api_key_hash, record_model_usage
+
+    record_model_usage(
+        model_name,
+        api_key_hash=_resolve_api_key_hash(model_name),
+        tokens=meta.get("total_tokens", 0),
+        cost=meta.get("cost", 0.0),
+    )
+
+
 async def clean_json_text_with_ai(
     driver: AsyncDriver, text: str, model_name: str = "", options: dict[str, Any] | None = None
 ) -> str:
@@ -117,6 +129,8 @@ async def render_output(
         "model_name": model_name or getattr(driver, "model", ""),
     }
 
+    _record_usage_to_ledger(model_name, resp.get("meta", {}))
+
     return {"text": raw, "usage": usage, "output_format": output_format}
 
 
@@ -211,6 +225,8 @@ async def ask_for_json(
     raw = resp.get("text", "")
     cleaned = clean_json_text(raw)
 
+    _record_usage_to_ledger(model_name, resp.get("meta", {}))
+
     try:
         json_obj = json.loads(cleaned)
         json_string = cleaned

prompture/async_driver.py

@@ -35,6 +35,7 @@ class AsyncDriver:
     supports_messages: bool = False
     supports_tool_use: bool = False
     supports_streaming: bool = False
+    supports_vision: bool = False
 
     callbacks: DriverCallbacks | None = None
 
@@ -165,5 +166,67 @@ class AsyncDriver:
         except Exception:
             logger.exception("Callback %s raised an exception", event)
 
+    def _validate_model_capabilities(
+        self,
+        provider: str,
+        model: str,
+        *,
+        using_tool_use: bool = False,
+        using_json_schema: bool = False,
+        using_vision: bool = False,
+    ) -> None:
+        """Log warnings when the model may not support a requested feature.
+
+        Uses models.dev metadata as a secondary signal.  Warnings only — the
+        API is the final authority and models.dev data may be stale.
+        """
+        from .model_rates import get_model_capabilities
+
+        caps = get_model_capabilities(provider, model)
+        if caps is None:
+            return
+
+        if using_tool_use and caps.supports_tool_use is False:
+            logger.warning(
+                "Model %s/%s may not support tool use according to models.dev metadata",
+                provider,
+                model,
+            )
+        if using_json_schema and caps.supports_structured_output is False:
+            logger.warning(
+                "Model %s/%s may not support structured output / JSON schema according to models.dev metadata",
+                provider,
+                model,
+            )
+        if using_vision and caps.supports_vision is False:
+            logger.warning(
+                "Model %s/%s may not support vision/image inputs according to models.dev metadata",
+                provider,
+                model,
+            )
+
+    def _check_vision_support(self, messages: list[dict[str, Any]]) -> None:
+        """Raise if messages contain image blocks and the driver lacks vision support."""
+        if self.supports_vision:
+            return
+        for msg in messages:
+            content = msg.get("content")
+            if isinstance(content, list):
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "image":
+                        raise NotImplementedError(
+                            f"{self.__class__.__name__} does not support vision/image inputs. "
+                            "Use a vision-capable model."
+                        )
+
+    def _prepare_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """Transform universal message format into provider-specific wire format.
+
+        Vision-capable async drivers override this to convert the universal
+        image blocks into their provider-specific format.
+        """
+        self._check_vision_support(messages)
+        return messages
+
     # Re-export the static helper for convenience
     _flatten_messages = staticmethod(Driver._flatten_messages)