Agentic-Investor

"""Interfaces for prompt abstractions."""

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, ClassVar, Type, TypeVar
from pydantic import BaseModel, Field

# Define a type variable for generic model support
T = TypeVar("T", bound=BaseModel)


class BasePromptInput(BaseModel):
    """Base class for prompt input models."""

    model_config = {"extra": "forbid"}  # Equivalent to additionalProperties: false


class PromptContent(BaseModel):
    """Model for content in prompt responses."""

    type: str = Field(default="text", description="Content type identifier")

    # Common fields for all content types
    content_id: Optional[str] = Field(None, description="Optional content identifier")

    # Type-specific fields (using discriminated unions pattern)
    # Text content
    text: Optional[str] = Field(None, description="Text content when type='text'")

    # JSON content (for structured data)
    json_data: Optional[Dict[str, Any]] = Field(
        None, description="JSON data when type='json'"
    )

    # Model content (will be converted to json_data during serialization)
    model: Optional[Any] = Field(
        None, exclude=True, description="Pydantic model instance"
    )

    def model_post_init(self, __context: Any) -> None:
        """Post-initialization hook to handle model conversion."""
        if self.model and not self.json_data:
            # Convert model to json_data
            if isinstance(self.model, BaseModel):
                self.json_data = self.model.model_dump()
                if not self.type or self.type == "text":
                    self.type = "json"


class PromptResponse(BaseModel):
    """Model for prompt responses."""

    content: List[PromptContent]

    @classmethod
    def from_model(cls, model: BaseModel) -> "PromptResponse":
        """Create a PromptResponse from a Pydantic model.

        This makes it easier to return structured data directly.

        Args:
            model: A Pydantic model instance to convert

        Returns:
            A PromptResponse with the model data in JSON format
        """
        return cls(
            content=[
                PromptContent(type="json", json_data=model.model_dump(), model=model)
            ]
        )

    @classmethod
    def from_text(cls, text: str) -> "PromptResponse":
        """Create a PromptResponse from plain text.

        Args:
            text: The text content

        Returns:
            A PromptResponse with text content
        """
        return cls(content=[PromptContent(type="text", text=text)])


class Prompt(ABC):
    """Abstract base class for all prompts."""

    name: ClassVar[str]
    description: ClassVar[str]
    input_model: ClassVar[Type[BasePromptInput]]
    output_model: ClassVar[Optional[Type[BaseModel]]] = None

    @abstractmethod
    async def generate(self, input_data: BasePromptInput) -> PromptResponse:
        """Generate the prompt with given arguments."""
        pass

    def get_schema(self) -> Dict[str, Any]:
        """Get JSON schema for the prompt."""
        schema = {
            "name": self.name,
            "description": self.description,
            "input": self.input_model.model_json_schema(),
        }

        if self.output_model:
            schema["output"] = self.output_model.model_json_schema()

        return schema