Spaces:
Running
Running
File size: 3,830 Bytes
46a8a46 | 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """poster state management"""
from typing import Dict, Any, Optional, List, Tuple, TypedDict
from dataclasses import dataclass
import time
@dataclass
class ModelConfig:
model_name: str
provider: str
temperature: float = 0.7
max_tokens: int = 4096
@dataclass
class TokenUsage:
input_text: int = 0
output_text: int = 0
input_vision: int = 0
output_vision: int = 0
def add_text(self, inp: int, out: int):
self.input_text += inp
self.output_text += out
def add_vision(self, inp: int, out: int):
self.input_vision += inp
self.output_vision += out
class PosterState(TypedDict):
# core paths
pdf_path: str
output_dir: str
poster_name: str
# model configs
text_model: ModelConfig
vision_model: ModelConfig
# processing results
images: Optional[Dict[str, Any]]
tables: Optional[Dict[str, Any]]
narrative: Optional[Dict[str, str]]
poster_plan: Optional[List[Dict[str, Any]]]
poster_width: int
poster_height: int
wireframe_layout: Optional[List[Dict[str, Any]]]
content_filled_layout: Optional[List[Dict[str, Any]]]
final_layout: Optional[List[Dict[str, Any]]]
narrative_content: Optional[Dict[str, Any]]
classified_visuals: Optional[Dict[str, Any]]
structured_sections: Optional[Dict[str, Any]]
story_board: Optional[Dict[str, Any]]
curated_content: Optional[Dict[str, Any]]
design_layout: Optional[List[Dict[str, Any]]]
section_title_design: Optional[Dict[str, Any]]
color_scheme: Optional[Dict[str, str]]
keywords: Optional[Dict[str, Any]]
styled_layout: Optional[List[Dict[str, Any]]]
# poster assets
url: str
logo_path: str
aff_logo_path: Optional[str]
# metadata
tokens: TokenUsage
current_agent: str
errors: List[str]
def create_state(pdf_path: str, text_model: str = "gpt-4.1-2025-04-14", vision_model: str = "gpt-4.1-2025-04-14", width: int = 84, height: int = 42, url: str = "", logo_path: str = "", aff_logo_path: str = "") -> PosterState:
"""create initial poster state"""
from pathlib import Path
poster_name = Path(pdf_path).parent.name or "test_poster"
output_dir = f"output/{poster_name}"
return PosterState(
pdf_path=pdf_path,
output_dir=output_dir,
poster_name=poster_name,
text_model=_get_model_config(text_model),
vision_model=_get_model_config(vision_model),
images=None,
tables=None,
narrative=None,
poster_plan=None,
poster_width=width,
poster_height=height,
wireframe_layout=None,
content_filled_layout=None,
final_layout=None,
narrative_content=None,
classified_visuals=None,
structured_sections=None,
story_board=None,
curated_content=None,
design_layout=None,
section_title_design=None,
color_scheme=None,
keywords=None,
styled_layout=None,
url=url,
logo_path=logo_path,
aff_logo_path=aff_logo_path,
tokens=TokenUsage(),
current_agent="init",
errors=[]
)
def _get_model_config(model_id: str) -> ModelConfig:
"""get model configuration"""
configs = {
"claude": ModelConfig("claude-sonnet-4-20250514", "anthropic"),
"claude-sonnet-4-20250514": ModelConfig("claude-sonnet-4-20250514", "anthropic"),
"gemini": ModelConfig("gemini-2.5-pro", "google"),
"gpt-4o-2024-08-06": ModelConfig("gpt-4o-2024-08-06", "openai"),
"gpt-4.1-2025-04-14": ModelConfig("gpt-4.1-2025-04-14", "openai"),
"gpt-4.1-mini-2025-04-14": ModelConfig("gpt-4.1-mini-2025-04-14", "openai")
}
return configs.get(model_id, configs["gpt-4.1-2025-04-14"]) |