import dspy
from pydantic import BaseModel, Field
from typing import List

# 1. Define the structure of your metadata
class DocMetadata(BaseModel):
    synopsis: str = Field(description="A one-sentence summary of the document.")
    tags: List[str] = Field(description="Relevant tags (NPCs, Locations, Items, Plot Points).")
    entities: List[str] = Field(description="Key names of people, places, or factions.")

class IngestionSignature(dspy.Signature):
    """
    You are an expert Dungeon Master's assistant.
    Analyze the provided notes and extract a concise synopsis and relevant metadata.
    """
    note: str = dspy.InputField(desc="The DM notes or session recap content.")
    # By using the Pydantic model as the type, DSPy handles the JSON formatting for you
    answer: DocMetadata = dspy.OutputField()

class IngestionAgent(dspy.Module):
    def __init__(self):
        super().__init__()
        # We use TypedPredictor to enforce the Pydantic schema
        # We use ChainOfThought because it helps 8B models "reason" through the tags 
        # before committing to the final JSON structure.
        self.process = dspy.TypedPredictor(IngestionSignature)

    def forward(self, note: str):
        # The .answer will now be a DocMetadata object, not a string!
        prediction = self.process(note=note)
        return prediction