pydantic-ai-slim 0.0.33__tar.gz → 0.0.34__tar.gz

{pydantic_ai_slim-0.0.33 → pydantic_ai_slim-0.0.34}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.33
+Version: 0.0.34
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -29,11 +29,15 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.0.33
+Requires-Dist: pydantic-graph==0.0.34
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.49.0; extra == 'anthropic'
+Provides-Extra: cli
+Requires-Dist: argcomplete>=3.5.0; extra == 'cli'
+Requires-Dist: prompt-toolkit>=3; extra == 'cli'
+Requires-Dist: rich>=13; extra == 'cli'
 Provides-Extra: cohere
 Requires-Dist: cohere>=5.13.11; extra == 'cohere'
 Provides-Extra: duckduckgo

pydantic_ai_slim-0.0.34/pydantic_ai/_cli.py

@@ -0,0 +1,225 @@
+from __future__ import annotations as _annotations
+
+import argparse
+import asyncio
+import sys
+from collections.abc import Sequence
+from datetime import datetime, timezone
+from importlib.metadata import version
+from pathlib import Path
+from typing import cast
+
+from typing_inspection.introspection import get_literal_values
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.models import KnownModelName
+from pydantic_graph.nodes import End
+
+try:
+    import argcomplete
+    from prompt_toolkit import PromptSession
+    from prompt_toolkit.auto_suggest import AutoSuggestFromHistory, Suggestion
+    from prompt_toolkit.buffer import Buffer
+    from prompt_toolkit.document import Document
+    from prompt_toolkit.history import FileHistory
+    from rich.console import Console, ConsoleOptions, RenderResult
+    from rich.live import Live
+    from rich.markdown import CodeBlock, Markdown
+    from rich.status import Status
+    from rich.syntax import Syntax
+    from rich.text import Text
+except ImportError as _import_error:
+    raise ImportError(
+        'Please install `rich`, `prompt-toolkit` and `argcomplete` to use the PydanticAI CLI, '
+        "you can use the `cli` optional group — `pip install 'pydantic-ai-slim[cli]'`"
+    ) from _import_error
+
+from pydantic_ai.agent import Agent
+from pydantic_ai.messages import ModelMessage, PartDeltaEvent, TextPartDelta
+
+__version__ = version('pydantic-ai')
+
+
+class SimpleCodeBlock(CodeBlock):
+    def __rich_console__(self, console: Console, options: ConsoleOptions) -> RenderResult:  # pragma: no cover
+        code = str(self.text).rstrip()
+        yield Text(self.lexer_name, style='dim')
+        yield Syntax(code, self.lexer_name, theme=self.theme, background_color='default', word_wrap=True)
+        yield Text(f'/{self.lexer_name}', style='dim')
+
+
+Markdown.elements['fence'] = SimpleCodeBlock
+
+
+def cli(args_list: Sequence[str] | None = None) -> int:  # noqa: C901  # pragma: no cover
+    parser = argparse.ArgumentParser(
+        prog='pai',
+        description=f"""\
+PydanticAI CLI v{__version__}\n\n
+
+Special prompt:
+* `/exit` - exit the interactive mode
+* `/markdown` - show the last markdown output of the last question
+* `/multiline` - toggle multiline mode
+""",
+        formatter_class=argparse.RawTextHelpFormatter,
+    )
+    parser.add_argument('prompt', nargs='?', help='AI Prompt, if omitted fall into interactive mode')
+    parser.add_argument(
+        '--model',
+        nargs='?',
+        help='Model to use, it should be "<provider>:<model>" e.g. "openai:gpt-4o". If omitted it will default to "openai:gpt-4o"',
+        default='openai:gpt-4o',
+    ).completer = argcomplete.ChoicesCompleter(list(get_literal_values(KnownModelName)))  # type: ignore[reportPrivateUsage]
+    parser.add_argument('--no-stream', action='store_true', help='Whether to stream responses from OpenAI')
+    parser.add_argument('--version', action='store_true', help='Show version and exit')
+
+    argcomplete.autocomplete(parser)
+    args = parser.parse_args(args_list)
+
+    console = Console()
+    console.print(f'pai - PydanticAI CLI v{__version__}', style='green bold', highlight=False)
+    if args.version:
+        return 0
+
+    now_utc = datetime.now(timezone.utc)
+    tzname = now_utc.astimezone().tzinfo.tzname(now_utc)  # type: ignore
+    try:
+        agent = Agent(
+            model=args.model or 'openai:gpt-4o',
+            system_prompt=f"""\
+    Help the user by responding to their request, the output should be concise and always written in markdown.
+    The current date and time is {datetime.now()} {tzname}.
+    The user is running {sys.platform}.""",
+        )
+    except UserError:
+        console.print(f'[red]Invalid model "{args.model}"[/red]')
+        return 1
+
+    stream = not args.no_stream
+
+    if prompt := cast(str, args.prompt):
+        try:
+            asyncio.run(ask_agent(agent, prompt, stream, console))
+        except KeyboardInterrupt:
+            pass
+        return 0
+
+    history = Path.home() / '.pai-prompt-history.txt'
+    session = PromptSession(history=FileHistory(str(history)))  # type: ignore
+    multiline = False
+    messages: list[ModelMessage] = []
+
+    while True:
+        try:
+            auto_suggest = CustomAutoSuggest(['/markdown', '/multiline', '/exit'])
+            text = cast(str, session.prompt('pai ➤ ', auto_suggest=auto_suggest, multiline=multiline))
+        except (KeyboardInterrupt, EOFError):
+            return 0
+
+        if not text.strip():
+            continue
+
+        ident_prompt = text.lower().strip(' ').replace(' ', '-').lstrip(' ')
+        if ident_prompt == '/markdown':
+            try:
+                parts = messages[-1].parts
+            except IndexError:
+                console.print('[dim]No markdown output available.[/dim]')
+                continue
+            for part in parts:
+                if part.part_kind == 'text':
+                    last_content = part.content
+                    console.print('[dim]Last markdown output of last question:[/dim]\n')
+                    console.print(Syntax(last_content, lexer='markdown', background_color='default'))
+
+            continue
+        if ident_prompt == '/multiline':
+            multiline = not multiline
+            if multiline:
+                console.print(
+                    'Enabling multiline mode. '
+                    '[dim]Press [Meta+Enter] or [Esc] followed by [Enter] to accept input.[/dim]'
+                )
+            else:
+                console.print('Disabling multiline mode.')
+            continue
+        if ident_prompt == '/exit':
+            console.print('[dim]Exiting…[/dim]')
+            return 0
+
+        try:
+            messages = asyncio.run(ask_agent(agent, text, stream, console, messages))
+        except KeyboardInterrupt:
+            return 0
+
+
+async def ask_agent(
+    agent: Agent,
+    prompt: str,
+    stream: bool,
+    console: Console,
+    messages: list[ModelMessage] | None = None,
+) -> list[ModelMessage]:  # pragma: no cover
+    status: None | Status = Status('[dim]Working on it…[/dim]', console=console)
+    live = Live('', refresh_per_second=15, console=console)
+    status.start()
+
+    async with agent.iter(prompt, message_history=messages) as agent_run:
+        console.print('\nResponse:', style='green')
+
+        content: str = ''
+        interrupted = False
+        try:
+            node = agent_run.next_node
+            while not isinstance(node, End):
+                node = await agent_run.next(node)
+                if Agent.is_model_request_node(node):
+                    async with node.stream(agent_run.ctx) as handle_stream:
+                        # NOTE(Marcelo): It took me a lot of time to figure out how to stop `status` and start `live`
+                        # in a context manager, so I had to do it manually with `stop` and `start` methods.
+                        # PR welcome to simplify this code.
+                        if status is not None:
+                            status.stop()
+                            status = None
+                        if not live.is_started:
+                            live.start()
+                        async for event in handle_stream:
+                            if isinstance(event, PartDeltaEvent) and isinstance(event.delta, TextPartDelta):
+                                if stream:
+                                    content += event.delta.content_delta
+                                    live.update(Markdown(content))
+        except KeyboardInterrupt:
+            interrupted = True
+        finally:
+            live.stop()
+
+        if interrupted:
+            console.print('[dim]Interrupted[/dim]')
+
+        assert agent_run.result
+        if not stream:
+            content = agent_run.result.data
+            console.print(Markdown(content))
+        return agent_run.result.all_messages()
+
+
+class CustomAutoSuggest(AutoSuggestFromHistory):
+    def __init__(self, special_suggestions: list[str] | None = None):  # pragma: no cover
+        super().__init__()
+        self.special_suggestions = special_suggestions or []
+
+    def get_suggestion(self, buffer: Buffer, document: Document) -> Suggestion | None:  # pragma: no cover
+        # Get the suggestion from history
+        suggestion = super().get_suggestion(buffer, document)
+
+        # Check for custom suggestions
+        text = document.text_before_cursor.strip()
+        for special in self.special_suggestions:
+            if special.startswith(text):
+                return Suggestion(special[len(text) :])
+        return suggestion
+
+
+def app():  # pragma: no cover
+    sys.exit(cli())

{pydantic_ai_slim-0.0.33 → pydantic_ai_slim-0.0.34}/pydantic_ai/agent.py

@@ -6,7 +6,7 @@ from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
 from copy import deepcopy
 from types import FrameType
-from typing import Any, Callable, Generic, cast, final, overload
+from typing import Any, Callable, ClassVar, Generic, cast, final, overload
 
 from opentelemetry.trace import NoOpTracer, use_span
 from typing_extensions import TypeGuard, TypeVar, deprecated
@@ -25,7 +25,7 @@ from . import (
     result,
     usage as _usage,
 )
-from .models.instrumented import InstrumentedModel
+from .models.instrumented import InstrumentationSettings, InstrumentedModel
 from .result import FinalResult, ResultDataT, StreamedRunResult
 from .settings import ModelSettings, merge_model_settings
 from .tools import (
@@ -56,6 +56,7 @@ __all__ = (
     'CallToolsNode',
     'ModelRequestNode',
     'UserPromptNode',
+    'InstrumentationSettings',
 )
 
 
@@ -112,8 +113,10 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
     The type of the result data, used to validate the result data, defaults to `str`.
     """
 
-    instrument: bool
-    """Automatically instrument with OpenTelemetry. Will use Logfire if it's configured."""
+    instrument: InstrumentationSettings | bool | None
+    """Options to automatically instrument with OpenTelemetry."""
+
+    _instrument_default: ClassVar[InstrumentationSettings | bool] = False
 
     _deps_type: type[AgentDepsT] = dataclasses.field(repr=False)
     _result_tool_name: str = dataclasses.field(repr=False)
@@ -147,7 +150,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
-        instrument: bool = False,
+        instrument: InstrumentationSettings | bool | None = None,
     ):
         """Create an agent.
 
@@ -177,7 +180,12 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                 [override the model][pydantic_ai.Agent.override] for testing.
             end_strategy: Strategy for handling tool calls that are requested alongside a final result.
                 See [`EndStrategy`][pydantic_ai.agent.EndStrategy] for more information.
-            instrument: Automatically instrument with OpenTelemetry. Will use Logfire if it's configured.
+            instrument: Set to True to automatically instrument with OpenTelemetry,
+                which will use Logfire if it's configured.
+                Set to an instance of [`InstrumentationSettings`][pydantic_ai.agent.InstrumentationSettings] to customize.
+                If this isn't set, then the last value set by
+                [`Agent.instrument_all()`][pydantic_ai.Agent.instrument_all]
+                will be used, which defaults to False.
         """
         if model is None or defer_model_check:
             self.model = model
@@ -213,6 +221,11 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             else:
                 self._register_tool(Tool(tool))
 
+    @staticmethod
+    def instrument_all(instrument: InstrumentationSettings | bool = True) -> None:
+        """Set the instrumentation options for all agents where `instrument` is not set."""
+        Agent._instrument_default = instrument
+
     @overload
     async def run(
         self,
@@ -422,7 +435,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         usage_limits = usage_limits or _usage.UsageLimits()
 
         if isinstance(model_used, InstrumentedModel):
-            tracer = model_used.tracer
+            tracer = model_used.options.tracer
         else:
             tracer = NoOpTracer()
         agent_name = self.name or 'agent'
@@ -1119,8 +1132,15 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         else:
             raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.')
 
-        if self.instrument and not isinstance(model_, InstrumentedModel):
-            model_ = InstrumentedModel(model_)
+        instrument = self.instrument
+        if instrument is None:
+            instrument = self._instrument_default
+
+        if instrument and not isinstance(model_, InstrumentedModel):
+            if instrument is True:
+                instrument = InstrumentationSettings()
+
+            model_ = InstrumentedModel(model_, instrument)
 
         return model_
 

{pydantic_ai_slim-0.0.33 → pydantic_ai_slim-0.0.34}/pydantic_ai/messages.py

@@ -189,7 +189,10 @@ class ToolReturnPart:
             return {'return_value': tool_return_ta.dump_python(self.content, mode='json')}
 
     def otel_event(self) -> Event:
-        return Event('gen_ai.tool.message', body={'content': self.content, 'role': 'tool', 'id': self.tool_call_id})
+        return Event(
+            'gen_ai.tool.message',
+            body={'content': self.content, 'role': 'tool', 'id': self.tool_call_id, 'name': self.tool_name},
+        )
 
 
 error_details_ta = pydantic.TypeAdapter(list[pydantic_core.ErrorDetails], config=pydantic.ConfigDict(defer_build=True))
@@ -244,7 +247,13 @@ class RetryPromptPart:
             return Event('gen_ai.user.message', body={'content': self.model_response(), 'role': 'user'})
         else:
             return Event(
-                'gen_ai.tool.message', body={'content': self.model_response(), 'role': 'tool', 'id': self.tool_call_id}
+                'gen_ai.tool.message',
+                body={
+                    'content': self.model_response(),
+                    'role': 'tool',
+                    'id': self.tool_call_id,
+                    'name': self.tool_name,
+                },
             )
 
 

{pydantic_ai_slim-0.0.33 → pydantic_ai_slim-0.0.34}/pydantic_ai/models/instrumented.py

@@ -43,9 +43,16 @@ MODEL_SETTING_ATTRIBUTES: tuple[
 ANY_ADAPTER = TypeAdapter[Any](Any)
 
 
-@dataclass
-class InstrumentedModel(WrapperModel):
-    """Model which is instrumented with OpenTelemetry."""
+@dataclass(init=False)
+class InstrumentationSettings:
+    """Options for instrumenting models and agents with OpenTelemetry.
+
+    Used in:
+
+    - `Agent(instrument=...)`
+    - [`Agent.instrument_all()`][pydantic_ai.agent.Agent.instrument_all]
+    - `InstrumentedModel`
+    """
 
     tracer: Tracer = field(repr=False)
     event_logger: EventLogger = field(repr=False)
@@ -53,20 +60,47 @@ class InstrumentedModel(WrapperModel):
 
     def __init__(
         self,
-        wrapped: Model | KnownModelName,
+        *,
+        event_mode: Literal['attributes', 'logs'] = 'attributes',
         tracer_provider: TracerProvider | None = None,
         event_logger_provider: EventLoggerProvider | None = None,
-        event_mode: Literal['attributes', 'logs'] = 'attributes',
     ):
+        """Create instrumentation options.
+
+        Args:
+            event_mode: The mode for emitting events. If `'attributes'`, events are attached to the span as attributes.
+                If `'logs'`, events are emitted as OpenTelemetry log-based events.
+            tracer_provider: The OpenTelemetry tracer provider to use.
+                If not provided, the global tracer provider is used.
+                Calling `logfire.configure()` sets the global tracer provider, so most users don't need this.
+            event_logger_provider: The OpenTelemetry event logger provider to use.
+                If not provided, the global event logger provider is used.
+                Calling `logfire.configure()` sets the global event logger provider, so most users don't need this.
+                This is only used if `event_mode='logs'`.
+        """
         from pydantic_ai import __version__
 
-        super().__init__(wrapped)
         tracer_provider = tracer_provider or get_tracer_provider()
         event_logger_provider = event_logger_provider or get_event_logger_provider()
         self.tracer = tracer_provider.get_tracer('pydantic-ai', __version__)
         self.event_logger = event_logger_provider.get_event_logger('pydantic-ai', __version__)
         self.event_mode = event_mode
 
+
+@dataclass
+class InstrumentedModel(WrapperModel):
+    """Model which is instrumented with OpenTelemetry."""
+
+    options: InstrumentationSettings
+
+    def __init__(
+        self,
+        wrapped: Model | KnownModelName,
+        options: InstrumentationSettings | None = None,
+    ) -> None:
+        super().__init__(wrapped)
+        self.options = options or InstrumentationSettings()
+
     async def request(
         self,
         messages: list[ModelMessage],
@@ -123,7 +157,7 @@ class InstrumentedModel(WrapperModel):
                 if isinstance(value := model_settings.get(key), (float, int)):
                     attributes[f'gen_ai.request.{key}'] = value
 
-        with self.tracer.start_as_current_span(span_name, attributes=attributes) as span:
+        with self.options.tracer.start_as_current_span(span_name, attributes=attributes) as span:
 
             def finish(response: ModelResponse, usage: Usage):
                 if not span.is_recording():
@@ -156,9 +190,9 @@ class InstrumentedModel(WrapperModel):
     def _emit_events(self, system: str, span: Span, events: list[Event]) -> None:
         for event in events:
             event.attributes = {'gen_ai.system': system, **(event.attributes or {})}
-        if self.event_mode == 'logs':
+        if self.options.event_mode == 'logs':
             for event in events:
-                self.event_logger.emit(event)
+                self.options.event_logger.emit(event)
         else:
             attr_name = 'events'
             span.set_attributes(

{pydantic_ai_slim-0.0.33 → pydantic_ai_slim-0.0.34}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pydantic-ai-slim"
-version = "0.0.33"
+version = "0.0.34"
 description = "Agent Framework / shim to use Pydantic with LLMs, slim package"
 authors = [{ name = "Samuel Colvin", email = "samuel@pydantic.dev" }]
 license = "MIT"
@@ -36,7 +36,7 @@ dependencies = [
     "griffe>=1.3.2",
     "httpx>=0.27",
     "pydantic>=2.10",
-    "pydantic-graph==0.0.33",
+    "pydantic-graph==0.0.34",
     "exceptiongroup; python_version < '3.11'",
     "opentelemetry-api>=1.28.0",
     "typing-inspection>=0.4.0",
@@ -55,6 +55,8 @@ mistral = ["mistralai>=1.2.5"]
 # Tools
 duckduckgo = ["duckduckgo-search>=7.0.0"]
 tavily = ["tavily-python>=0.5.0"]
+# CLI
+cli = ["rich>=13", "prompt-toolkit>=3", "argcomplete>=3.5.0"]
 
 [dependency-groups]
 dev = [
@@ -71,6 +73,9 @@ dev = [
     "diff-cover>=9.2.0",
 ]
 
+[project.scripts]
+pai = "pydantic_ai._cli:app"
+
 [tool.hatch.build.targets.wheel]
 packages = ["pydantic_ai"]