Discord link archiver bot with Gemini tags/summaries
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224"""AI integration: Google Gemini for classification, summaries, tagging, and OCR."""
import jsonimport loggingfrom typing import TypedDict, Optional, List, Dict
import google.generativeai as genai
logger = logging.getLogger("archivist.ai")
class ClassificationResult(TypedDict): genre: str tags: List[str] new_tags: List[str] metadata: Dict[str, str] error: Optional[str]
class SummaryResult(TypedDict): gloss: Optional[str] summary: Optional[str] error: Optional[str]
class ImageResult(TypedDict): alt_text: Optional[str] transcription: Optional[str] error: Optional[str]
async def classify_and_tag( text: str, tag_pool: List[str], api_key: str, model: str = "gemini-2.0-flash", system_prompt: str = "",) -> ClassificationResult: """Pass 1: determine genre, metadata, and tags.""" genai.configure(api_key=api_key)
prompt = ( f"{system_prompt}\n\n" f"Existing tag pool: {json.dumps(tag_pool)}\n\n" f"Content to analyze (first 30k chars):\n{text[:30000]}" )
try: model_instance = genai.GenerativeModel(model) response = await model_instance.generate_content_async( prompt, generation_config=genai.types.GenerationConfig( temperature=0.1, response_mime_type="application/json", ), )
data = json.loads(response.text) return { "genre": data.get("genre", "Unknown"), "tags": [t.lower().strip() for t in data.get("tags", []) if t], "new_tags": [t.lower().strip() for t in data.get("new_tags", []) if t], "metadata": data.get("metadata", {}) or {}, "error": None, } except Exception as e: logger.exception("Classification failed") return { "genre": "Unknown", "tags": [], "new_tags": [], "metadata": {}, "error": str(e), }
async def generate_summary( text: str, title: Optional[str], url: str, genre: str, metadata: Dict[str, str], api_key: str, model: str = "gemini-2.0-flash", system_prompt: str = "",) -> SummaryResult: """Pass 2: generate a genre-aware summary.""" genai.configure(api_key=api_key)
prompt = f"{system_prompt}\n\n" prompt += f"GENRE: {genre}\n" if metadata: prompt += f"METADATA: {json.dumps(metadata)}\n" prompt += f"URL: {url}\n" if title: prompt += f"TITLE: {title}\n" prompt += f"\nCONTENT:\n{text[:30000]}"
try: model_instance = genai.GenerativeModel(model) response = await model_instance.generate_content_async( prompt, generation_config=genai.types.GenerationConfig(temperature=0.3), )
response_text = response.text or "" gloss = None summary = None for line in response_text.splitlines(): line = line.strip() if line.lower().startswith("**gloss:**"): gloss = line.split(":**", 1)[1].strip() elif line.lower().startswith("**summary:**"): summary = line.split(":**", 1)[1].strip()
if not gloss and not summary: # Fallback: don't dump unbounded LLM output into Discord — trim. fallback = response_text.strip() summary = (fallback[:500] + "…") if len(fallback) > 500 else fallback
return {"gloss": gloss, "summary": summary, "error": None} except Exception as e: logger.exception("Summary failed") return {"gloss": None, "summary": None, "error": str(e)}
async def describe_image( image_bytes: bytes, mime_type: str, api_key: str, model: str = "gemini-2.0-flash", system_prompt: str = "",) -> ImageResult: """ Describe an image with Gemini's multimodal input: produce alt-text plus a verbatim transcription if any text is visible. The expected response shape is configured in `config/default.yaml` under `gemini.image_system_prompt`. """ genai.configure(api_key=api_key)
prompt = system_prompt or ( "Describe this image for a blind or low-vision reader, then " "transcribe any visible text.\n" "Respond in this exact format:\n" "**Alt Text:** [one-to-two sentence description]\n" "**Transcription:** [verbatim text, or omit the line entirely]" )
try: model_instance = genai.GenerativeModel(model) response = await model_instance.generate_content_async( [prompt, {"mime_type": mime_type, "data": image_bytes}], generation_config=genai.types.GenerationConfig(temperature=0.2), )
response_text = response.text or ""
alt_text: Optional[str] = None transcription_lines: List[str] = [] in_transcription = False
for line in response_text.splitlines(): stripped = line.strip() lower = stripped.lower() if lower.startswith("**alt text:**"): alt_text = stripped.split(":**", 1)[1].strip() or None in_transcription = False elif lower.startswith("**transcription:**"): rest = stripped.split(":**", 1)[1].strip() if rest: transcription_lines.append(rest) in_transcription = True elif in_transcription: transcription_lines.append(line)
transcription = ( "\n".join(transcription_lines).strip() or None if transcription_lines else None )
# Fallback: no structured shape at all — treat the whole response as alt-text, # capped so a runaway model can't blow the Discord 2000-char limit. if alt_text is None and transcription is None: fallback = response_text.strip() alt_text = ( (fallback[:500] + "…") if len(fallback) > 500 else (fallback or None) )
return { "alt_text": alt_text, "transcription": transcription, "error": None, } except Exception as e: logger.exception("Image description failed") return {"alt_text": None, "transcription": None, "error": str(e)}
async def generate_embedding( text: str, api_key: str, model: str = "text-embedding-004",) -> Optional[List[float]]: """ Generate a 768-dim embedding for a piece of text, used for semantic search. Return shape must match the `vector(768)` column in `archived_links`. """ genai.configure(api_key=api_key) try: result = genai.embed_content( model=model, content=text, task_type="retrieval_document", ) embedding = result.get("embedding") if isinstance(result, dict) else None if not embedding: logger.warning("Embedding response had no 'embedding' field") return None return embedding except Exception as e: logger.error(f"Embedding failed: {type(e).__name__}: {e}") return None