Gemini-Image-Library-Agent/schemas.py at main · peterjakubowski/Gemini-Image-Library-Agent · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
# Google Gemini Image Library Agent
#
# Author: Peter Jakubowski
# Date: 2/1/2025
# Description: Database schema
#

from pydantic.dataclasses import dataclass as pydantic_dataclass
from pydantic import BaseModel, RootModel, Field, ConfigDict
from datetime import datetime

# =========================
# ==== Database Schema ====
# =========================


# Core schema
@pydantic_dataclass
class Assets:
    """
    The root table of our database keeps track of the assets in our system.

    """

    # Configure the BaseModel to ignore any extra attributes given at creation
    model_config = ConfigDict(extra='ignore')
    id: int = Field(default=None)
    image_path: str = Field(default=None)
    file_name: str | None = Field(default=None)
    file_type: str | None = Field(default=None)


# Metadata schema
@pydantic_dataclass
class AssetMetadata:
    """
    The metadata table in our database keeps record of select image metadata that is extracted from the original image.
    For simplicity, we're going to save our list of keywords as a comma separated string.

    """

    # Configure the BaseModel to ignore any extra attributes given at creation
    model_config = ConfigDict(extra='ignore')
    id: int = Field(default=None)
    capture_date: datetime | None = Field(default=None)
    description: str | None = Field(default=None)
    keywords: str | None = Field(default=None)
    creator: str | None = Field(default=None)
    person_in_image: str | None = Field(default=None)
    location: str | None = Field(default=None)
    city: str | None = Field(default=None)
    state: str | None = Field(default=None)


# Generative schema
@pydantic_dataclass
class GenerativeMetadata:
    """
    The generative table in our database keeps record of generated content created by the Gemini model.
    For simplicity, we're going to save our list of keywords as a comma separated string.

    """

    # Configure the BaseModel to ignore any extra attributes given at creation
    model_config = ConfigDict(extra='ignore')
    id: int = Field(default=None)
    description: str = Field(default=None)
    keywords: str = Field(default=None)
    style: str = Field(default=None)
    mood: str = Field(default=None)
    image_type: str = Field(default=None)
    subject: str = Field(default=None)
    context: str = Field(default=None)
    details: str = Field(default=None)
    lighting: str = Field(default=None)
    framing: str = Field(default=None)
    lens_and_camera: str = Field(default=None)


# Embedding schema
@pydantic_dataclass
class Embeddings:
    """
    The embedding table in our database keeps record of embeddings generated by the Gemini model.

    """

    # Configure the BaseModel to ignore any extra attributes given at creation
    model_config = ConfigDict(extra='ignore')
    id: int = Field(default=None)
    genai_description_vector: bytes = Field(default=None)

# =========================
# ==== Response Schema ====
# =========================


# Generative description model response schema
class DescriptionResponseSchema(BaseModel):
    description: str = Field(description='Long, detailed description of the image 100-500 words in length.')
    keywords: list[str] = Field(description='List of keywords describing the image.')
    style: str = Field(description="Describe the style of the image. Examples: Monochromatic, Natural, Realistic, Soft")
    mood: str = Field(description="Describe the mood of the image. Examples: Calm, Dark, Optimistic, Reflective")
    image_type: str = Field(description="Describe the image type. Examples: Photograph, Drawing, Hyper-real painting")
    subject: str = Field(description="Describe the subject of the image. Examples: Cat, Cactus, Small Terrier")
    context: str = Field(description="Describe the context of the image. Examples: Sitting by a window")
    details: str = Field(description="Describe the details in the image. Examples: Staring at a cupcake")
    lighting: str = Field(description="Describe the lighting seen in the image. Examples: Natural light, Studio light")
    framing: str = Field(description="Describe the framing of the image. Examples: Medium shot, Close-up, Wide")
    lens_and_camera: str = Field(description=("Describe the characteristics of the lens and camera used in the image. "
                                              "Examples: Shallow depth of field, Wide angle, Sharp, Vintage, Macro"))