File size: 1,734 Bytes
7992c94 7b58f25 7992c94 7b58f25 7992c94 7b58f25 7992c94 7b58f25 7992c94 7b58f25 |
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 |
from pydantic import BaseModel
import json
from image_api import get_random_image
import base64
import requests
class TVCrewOutput(BaseModel):
image_keyword: str
class TVCrew:
def __init__(self, client, system_prompt: str):
self.model = "google/gemma-3-12b-it"
self.system_prompt = system_prompt
self.context = ""
self.client = client
def suggest_image(self) -> tuple[str, str] | tuple[None, None]:
"""
Finds a image related to conversation and returns the url and caption
"""
messages_tv_crew = [
{"role": "system", "content": self.system_prompt},
{"role": "user", "content": self.context},
]
response_tv_crew = self.client.chat.completions.parse(
model=self.model,
messages=messages_tv_crew,
response_format=TVCrewOutput,
)
response_tv_crew = response_tv_crew.choices[0].message.content
response_tv_crew_json = json.loads(response_tv_crew)
image_keyword = response_tv_crew_json["image_keyword"]
if image_keyword == "none":
return None, None
image_url, alt = get_random_image(image_keyword)
caption = None
if image_url:
caption = alt if len(alt) > 0 else image_keyword
self.add_context(f"TV Shows: {caption}\n")
img_bytes = requests.get(image_url).content
image_base64 = base64.b64encode(img_bytes).decode("utf-8")
return image_base64, caption
return None, None
def add_context(self, content):
"""Will be updated by main script"""
self.context += content
def clear_context(self):
self.context = ""
|