| 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 = "" | |