agentpool 2.1.9__py3-none-any.whl

agentpool/prompts/prompts.py

@@ -0,0 +1,581 @@
+"""Prompt models for AgentPool."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from dataclasses import field
+import inspect
+import os
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Self, assert_never
+
+from mcp.types import Prompt as MCPPrompt, PromptArgument
+from pydantic import ConfigDict, Field
+from pydantic_ai import BinaryContent, ImageUrl, SystemPromptPart, UserPromptPart
+from schemez import Schema
+from upathtools import UPath, to_upath
+
+from agentpool.log import get_logger
+from agentpool.mcp_server import MCPClient
+from agentpool.utils import importing
+from agentpool.utils.inspection import execute
+
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from fastmcp.prompts.prompt import FunctionPrompt, Prompt as FastMCPPrompt
+    from pydantic_ai import ModelRequestPart
+
+
+logger = get_logger(__name__)
+
+MessageContentType = Literal["text", "resource", "image_url", "image_base64"]
+# Our internal role type (could include more roles)
+MessageRole = Literal["system", "user", "assistant", "tool"]
+
+type CompletionFunction = Callable[[str], list[str]]
+"""Type for completion functions. Takes current value, returns possible completions."""
+
+
+class MessageContent(Schema):
+    """Content item in a message."""
+
+    type: MessageContentType
+    content: str  # The actual content (text/uri/url/base64)
+    alt_text: str | None = None  # For images or resource descriptions
+
+    def to_pydantic(self) -> UserPromptPart:
+        """Convert MessageContent to Pydantic model."""
+        if self.type == "text":
+            return UserPromptPart(self.content)
+        if self.type == "image_url":
+            return UserPromptPart([ImageUrl(self.content)])
+        if self.type == "image_base64":
+            bin_content = BinaryContent(self.content.encode(), media_type="image/jpeg")
+            return UserPromptPart([bin_content])
+        msg = "Unsupported content type"
+        raise ValueError(msg)
+
+
+class PromptParameter(Schema):
+    """Prompt argument with validation information."""
+
+    name: str
+    """Name of the argument as used in the prompt."""
+
+    description: str | None = None
+    """Human-readable description of the argument."""
+
+    required: bool = False
+    """Whether this argument must be provided when formatting the prompt."""
+
+    default: Any | None = None
+    """Default value if argument is optional."""
+
+    def to_mcp_argument(self) -> PromptArgument:
+        """Convert to MCP PromptArgument."""
+        from mcp.types import PromptArgument
+
+        return PromptArgument(name=self.name, description=self.description, required=self.required)
+
+
+class PromptMessage(Schema):
+    """A message in a prompt template."""
+
+    role: MessageRole
+    content: str | MessageContent | list[MessageContent] = ""
+
+    def get_text_content(self) -> str:
+        """Get text content of message."""
+        match self.content:
+            case str():
+                return self.content
+            case MessageContent() if self.content.type == "text":
+                return self.content.content
+            case list() if self.content:
+                # Join text content items with space
+                text_items = [
+                    item.content
+                    for item in self.content
+                    if isinstance(item, MessageContent) and item.type == "text"
+                ]
+                return " ".join(text_items) if text_items else ""
+            case _:
+                return ""
+
+    def to_pydantic_parts(self) -> list[ModelRequestPart]:
+        match self.role:
+            case "system":
+                return [SystemPromptPart(str(self.content))]
+            case "user":
+                match self.content:
+                    case str():
+                        return [UserPromptPart(self.content)]
+                    case MessageContent():
+                        return [self.content.to_pydantic()]
+                    case list():
+                        return [i.to_pydantic() for i in self.content]
+        return []
+
+
+class BasePrompt(Schema):
+    """Base class for all prompts."""
+
+    name: str | None = Field(None, exclude=True)
+    """Technical identifier (automatically set from config key during registration)."""
+
+    description: str
+    """Human-readable description of what this prompt does."""
+
+    title: str | None = None
+    """Title of the prompt."""
+
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    """Additional metadata for storing custom prompt information."""
+    # messages: list[PromptMessage]
+
+    async def format(self, arguments: dict[str, Any] | None = None) -> list[PromptMessage]:
+        """Format this prompt with given arguments.
+
+        Args:
+            arguments: Optional argument values
+
+        Returns:
+            List of formatted messages
+
+        Raises:
+            ValueError: If required arguments are missing
+        """
+        raise NotImplementedError
+
+    def to_mcp_prompt(self) -> MCPPrompt:
+        """Convert to MCP Prompt."""
+        raise NotImplementedError
+
+
+class StaticPrompt(BasePrompt):
+    """Static prompt defined by message list."""
+
+    messages: list[PromptMessage]
+    """List of messages that make up this prompt."""
+
+    type: Literal["text"] = Field("text", init=False)
+    """Discriminator field identifying this as a static text prompt."""
+
+    arguments: list[PromptParameter] = Field(default_factory=list)
+    """List of arguments that this prompt accepts."""
+
+    def validate_arguments(self, provided: dict[str, Any]) -> None:
+        """Validate that required arguments are provided."""
+        required = {arg.name for arg in self.arguments if arg.required}
+        missing = required - set(provided)
+        if missing:
+            msg = f"Missing required arguments: {', '.join(missing)}"
+            raise ValueError(msg)
+
+    def to_mcp_prompt(self) -> MCPPrompt:
+        """Convert to MCP Prompt."""
+        from mcp.types import Prompt as MCPPrompt
+
+        if self.name is None:
+            msg = "Prompt name not set. This should be set during registration."
+            raise ValueError(msg)
+        args = [arg.to_mcp_argument() for arg in self.arguments]
+        return MCPPrompt(name=self.name, description=self.description, arguments=args)
+
+    def to_fastmcp_prompt(self) -> FastMCPPrompt:
+        from fastmcp.prompts.prompt import (
+            Prompt as FastMCPPrompt,
+            PromptArgument as FastMCPArgument,
+        )
+
+        params = [
+            FastMCPArgument(name=p.name, description=p.description, required=p.required)
+            for p in self.arguments
+        ]
+        return FastMCPPrompt(
+            name=self.name or "",
+            title=self.title,
+            description=self.description,
+            arguments=params,
+            icons=None,
+            tags=set(),
+        )
+
+    async def format(self, arguments: dict[str, Any] | None = None) -> list[PromptMessage]:
+        """Format static prompt messages with arguments."""
+        args = arguments or {}
+        self.validate_arguments(args)
+
+        # Add default values for optional arguments
+        for arg in self.arguments:
+            if arg.name not in args and not arg.required:
+                args[arg.name] = arg.default if arg.default is not None else ""
+
+        # Format all messages
+        formatted_messages = []
+        for msg in self.messages:
+            match msg.content:
+                case str():
+                    content: MessageContent | list[MessageContent] = MessageContent(
+                        type="text", content=msg.content.format(**args)
+                    )
+                case MessageContent() if msg.content.type == "text":
+                    msg_content = msg.content.content.format(**args)
+                    content = MessageContent(type="text", content=msg_content)
+                case list():
+                    content = [
+                        MessageContent(
+                            type=item.type,
+                            content=item.content.format(**args)
+                            if item.type == "text"
+                            else item.content,
+                            alt_text=item.alt_text,
+                        )
+                        for item in msg.content
+                        if isinstance(item, MessageContent)
+                    ]
+                case _:
+                    content = msg.content
+
+            formatted_messages.append(PromptMessage(role=msg.role, content=content))
+
+        return formatted_messages
+
+
+class DynamicPrompt(BasePrompt):
+    """Dynamic prompt loaded from callable."""
+
+    import_path: str | Callable[..., str | Awaitable[str]]
+    """Dotted import path to the callable that generates the prompt."""
+
+    template: str | None = None
+    """Optional template string for formatting the callable's output."""
+
+    completions: dict[str, str] | None = None
+    """Optional mapping of argument names to completion functions."""
+
+    type: Literal["function"] = Field("function", init=False)
+    """Discriminator field identifying this as a function-based prompt."""
+
+    @property
+    def messages(self) -> list[PromptMessage]:
+        """Get the template messages for this prompt.
+
+        Note: These are template messages - actual content will be populated
+        during format() when the callable is executed.
+        """
+        template = self.template or "{result}"
+        sys_content = MessageContent(type="text", content=f"Content from {self.name}:")
+        user_content = MessageContent(type="text", content=template)
+        return [
+            PromptMessage(role="system", content=sys_content),
+            PromptMessage(role="user", content=user_content),
+        ]
+
+    def to_fastmcp_prompt(self) -> FunctionPrompt:
+        from fastmcp.prompts.prompt import FunctionPrompt
+
+        return FunctionPrompt.from_function(self.fn, title=self.title, description=self.description)
+
+    @property
+    def fn(self) -> Callable[..., str | Awaitable[str]]:
+        if isinstance(self.import_path, str):
+            return importing.import_callable(self.import_path)
+        return self.import_path
+
+    def get_completion_functions(self) -> dict[str, CompletionFunction]:
+        """Resolve completion function import paths and return a completion fn dict."""
+        completion_funcs: dict[str, CompletionFunction] = {}
+        if not self.completions:
+            return {}
+        for arg_name, import_path in self.completions.items():
+            try:
+                func = importing.import_callable(import_path)
+                completion_funcs[arg_name] = func
+            except ValueError:
+                msg = "Failed to import completion function for %s: %s"
+                logger.warning(msg, arg_name, import_path)
+        return completion_funcs
+
+    async def format(self, arguments: dict[str, Any] | None = None) -> list[PromptMessage]:
+        """Format this prompt with given arguments."""
+        args = arguments or {}
+        try:
+            result: Any = await execute(self.fn, **args)
+            template = self.template or "{result}"
+            msg = template.format(result=result)
+            content = MessageContent(type="text", content=msg)
+            msg = f"Content from {self.name}:"
+            sys_content = MessageContent(type="text", content=msg)
+            return [
+                PromptMessage(role="system", content=sys_content),
+                PromptMessage(role="user", content=content),
+            ]
+        except Exception as exc:
+            msg = f"Failed to execute prompt callable: {exc}"
+            raise ValueError(msg) from exc
+
+    @classmethod
+    def from_callable(
+        cls,
+        fn: Callable[..., Any] | str,
+        *,
+        name_override: str | None = None,
+        description_override: str | None = None,
+        template_override: str | None = None,
+        completions: Mapping[str, CompletionFunction] | None = None,
+    ) -> DynamicPrompt:
+        """Create a prompt from a callable.
+
+        Args:
+            fn: Function or import path to create prompt from
+            name_override: Optional override for prompt name
+            description_override: Optional override for prompt description
+            template_override: Optional override for message template
+            completions: Optional dict mapping argument names to completion functions
+
+        Returns:
+            DynamicPrompt instance
+
+        Raises:
+            ValueError: If callable cannot be imported or is invalid
+        """
+        from docstring_parser import parse as parse_docstring
+
+        completions = completions or {}
+        if isinstance(fn, str):
+            fn = importing.import_callable(fn)
+
+        name = name_override or getattr(fn, "__name__", "unknown")
+        if docstring := inspect.getdoc(fn):
+            parsed = parse_docstring(docstring)
+            description = description_override or parsed.short_description
+        else:
+            description = description_override or f"Prompt from {name}"
+
+        from agentpool.utils.inspection import get_fn_qualname
+
+        path = f"{fn.__module__}.{get_fn_qualname(fn)}"
+        return cls(
+            name=name,
+            description=description or "",
+            import_path=path,
+            template=template_override,
+            metadata={"source": "function", "import_path": path},
+        )
+
+    def to_mcp_prompt(self) -> MCPPrompt:
+        """Convert to MCP Prompt."""
+        from docstring_parser import parse as parse_docstring
+        from mcp.types import Prompt as MCPPrompt
+
+        sig = inspect.signature(self.fn)
+        # hints = get_type_hints(self.fn, include_extras=True, localns=locals())
+        if docstring := inspect.getdoc(self.fn):
+            parsed = parse_docstring(docstring)
+            # Create mapping of param names to descriptions
+            arg_docs = {
+                param.arg_name: param.description
+                for param in parsed.params
+                if param.arg_name and param.description
+            }
+        else:
+            arg_docs = {}
+
+        # Create arguments
+        arguments = []
+        for param_name, param in sig.parameters.items():
+            if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD):
+                continue
+
+            required = param.default == param.empty
+            arg = PromptParameter(
+                name=param_name,
+                description=arg_docs.get(param_name),
+                required=required,
+                default=None if param.default is param.empty else param.default,
+            )
+            arguments.append(arg)
+
+        if self.name is None:
+            msg = "Prompt name not set. This should be set during registration."
+            raise ValueError(msg)
+        args = [arg.to_mcp_argument() for arg in arguments]
+        return MCPPrompt(name=self.name, description=self.description, arguments=args)
+
+
+class MCPClientPrompt(BasePrompt):
+    """A prompt that can be rendered from an MCP server."""
+
+    client: MCPClient
+    """The client used to render the prompt."""
+
+    name: str
+    """The name of the prompt."""
+
+    arguments: list[dict[str, Any]] = field(default_factory=list)
+    """A list of arguments for the prompt."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    def to_mcp_prompt(self) -> MCPPrompt:
+        """Convert to MCP Prompt."""
+        args = [
+            PromptArgument(name=i["name"], description=i["description"], required=i["required"])
+            for i in self.arguments
+        ]
+        return MCPPrompt(name=self.name, description=self.description, arguments=args)
+
+    @classmethod
+    def from_fastmcp(cls, client: MCPClient, prompt: MCPPrompt) -> Self:
+        """Convert MCP prompt to our Prompt class."""
+        arguments = [
+            {
+                "name": arg.name,
+                "description": arg.description,
+                "required": arg.required or False,
+            }
+            for arg in prompt.arguments or []
+        ]
+
+        return cls(
+            name=prompt.name,
+            description=prompt.description or "",
+            arguments=arguments,
+            client=client,
+        )
+
+    def __repr__(self) -> str:
+        return f"Prompt(name={self.name!r}, description={self.description!r})"
+
+    async def get_components(
+        self, arguments: dict[str, str] | None = None
+    ) -> list[SystemPromptPart | UserPromptPart]:
+        """Get prompt as pydantic-ai message components.
+
+        Args:
+            arguments: Arguments to pass to the prompt template
+
+        Returns:
+            List of message parts ready for agent usage
+
+        Raises:
+            RuntimeError: If prompt fetch fails
+            ValueError: If prompt contains unsupported message types
+        """
+        from agentpool.mcp_server.conversions import content_block_as_text
+
+        try:
+            result = await self.client.get_prompt(self.name, arguments)
+        except Exception as e:
+            msg = f"Failed to get prompt {self.name!r}: {e}"
+            raise RuntimeError(msg) from e
+
+        # Convert MCP messages to pydantic-ai parts
+        parts: list[SystemPromptPart | UserPromptPart] = []
+        for message in result.messages:
+            # Extract text content from MCP message
+            text_content = content_block_as_text(message.content)
+            # Convert based on role
+            match message.role:
+                case "user":
+                    parts.append(UserPromptPart(content=text_content))
+                case "assistant":
+                    # Convert assistant messages to user parts for context
+                    parts.append(UserPromptPart(content=f"Assistant: {text_content}"))
+                case _ as unreachable:
+                    assert_never(unreachable)
+
+        if not parts:
+            msg = f"No supported message parts found in prompt {self.name!r}"
+            raise ValueError(msg)
+
+        return parts
+
+
+class FilePrompt(BasePrompt):
+    """Prompt loaded from a file.
+
+    This type of prompt loads its content from a file, allowing for longer or more
+    complex prompts to be managed in separate files. The file content is loaded
+    and parsed according to the specified format.
+    """
+
+    path: str | os.PathLike[str] | UPath
+    """Path to the file containing the prompt content."""
+
+    fmt: Literal["text", "markdown", "jinja2"] = Field("text", alias="format")
+    """Format of the file content (text, markdown, or jinja2 template)."""
+
+    type: Literal["file"] = Field("file", init=False)
+    """Discriminator field identifying this as a file-based prompt."""
+
+    watch: bool = False
+    """Whether to watch the file for changes and reload automatically."""
+
+    def to_mcp_prompt(self) -> MCPPrompt:
+        """Convert to MCP Prompt."""
+        return MCPPrompt(
+            name=self.name or to_upath(self.path).name,
+            description=self.description,
+        )
+
+    @property
+    def messages(self) -> list[PromptMessage]:
+        """Get messages from file content."""
+        content = to_upath(self.path).read_text("utf-8")
+        match self.fmt:
+            case "text":
+                # Simple text format - whole file as user message
+                msg = MessageContent(type="text", content=content)
+            case "markdown":
+                # TODO: Parse markdown sections into separate messages
+                msg = MessageContent(type="text", content=content)
+            case "jinja2":
+                # Raw template - will be formatted during format()
+                msg = MessageContent(type="text", content=content)
+            case _ as unreachable:
+                assert_never(unreachable)
+        return [PromptMessage(role="user", content=msg)]
+
+    async def format(self, arguments: dict[str, Any] | None = None) -> list[PromptMessage]:
+        """Format the file content with arguments."""
+        args = arguments or {}
+        content = to_upath(self.path).read_text("utf-8")
+
+        if self.fmt == "jinja2":
+            # Use jinja2 for template formatting
+            import jinja2
+
+            env = jinja2.Environment(autoescape=True, enable_async=True)
+            template = env.from_string(content)
+            content = await template.render_async(**args)
+        else:
+            # Use simple string formatting
+            try:
+                content = content.format(**args)
+            except KeyError as exc:
+                msg = f"Missing argument in template: {exc}"
+                raise ValueError(msg) from exc
+        msg_content = MessageContent(type="text", content=content)
+        return [PromptMessage(role="user", content=msg_content)]
+
+
+# Type to use in configuration
+PromptType = Annotated[StaticPrompt | DynamicPrompt | FilePrompt, Field(discriminator="type")]
+
+
+if __name__ == "__main__":
+
+    def prompt_fn() -> str:
+        return "hello"
+
+    async def main() -> None:
+        prompt = DynamicPrompt(import_path=prompt_fn, name="test", description="test")
+        result = await prompt.format()
+        print(result)
+
+    import anyio
+
+    anyio.run(main)

agentpool/queries/tree-sitter-language-pack/README.md

@@ -0,0 +1,7 @@
+These scm files are all adapted from the github repositories listed here:
+
+https://github.com/Goldziher/tree-sitter-language-pack/blob/main/sources/language_definitions.json
+
+See this URL for information on the licenses of each repo:
+
+https://github.com/Goldziher/tree-sitter-language-pack/

agentpool/queries/tree-sitter-language-pack/arduino-tags.scm

@@ -0,0 +1,5 @@
+(function_declarator
+  declarator: (identifier) @name.definition.function) @definition.function
+
+(call_expression
+  function: (identifier) @name.reference.call) @reference.call

agentpool/queries/tree-sitter-language-pack/c-tags.scm

@@ -0,0 +1,9 @@
+(struct_specifier name: (type_identifier) @name.definition.class body:(_)) @definition.class
+
+(declaration type: (union_specifier name: (type_identifier) @name.definition.class)) @definition.class
+
+(function_declarator declarator: (identifier) @name.definition.function) @definition.function
+
+(type_definition declarator: (type_identifier) @name.definition.type) @definition.type
+
+(enum_specifier name: (type_identifier) @name.definition.type) @definition.type

agentpool/queries/tree-sitter-language-pack/chatito-tags.scm

@@ -0,0 +1,16 @@
+; Definitions
+(intent_def
+  (intent) @name.definition.intent) @definition.intent
+
+(slot_def
+  (slot) @name.definition.slot) @definition.slot
+
+(alias_def
+  (alias) @name.definition.alias) @definition.alias
+
+; References
+(slot_ref
+  (slot) @name.reference.slot) @reference.slot
+
+(alias_ref
+  (alias) @name.reference.alias) @reference.alias

agentpool/queries/tree-sitter-language-pack/clojure-tags.scm

@@ -0,0 +1,7 @@
+(list_lit
+  meta: _*
+  . (sym_lit name: (sym_name) @ignore)
+  . (sym_lit name: (sym_name) @name.definition.method)
+  (#match? @ignore "^def.*"))
+
+(sym_lit name: (sym_name) @name.reference.call)