chore: apply formatting and hash updates

Author: hammertoe

lib/db/postgres_client.py

@@ -3,7 +3,8 @@
 from __future__ import annotations
 
 from contextlib import contextmanager
-from typing import Any, Iterable
+from typing import Any
+from collections.abc import Iterable
 
 from psycopg_pool import ConnectionPool
 

lib/embeddings/google_client.py

@@ -32,10 +32,7 @@ def __init__(self):
                 raise ValueError("GOOGLE_API_KEY not set in environment")
             self.client = genai.Client(api_key=config.embedding.api_key)
         else:
-            if (
-                not config.embedding.vertex_project
-                or not config.embedding.vertex_location
-            ):
+            if not config.embedding.vertex_project or not config.embedding.vertex_location:
                 raise ValueError(
                     "VERTEX_PROJECT and VERTEX_LOCATION must be set when EMBEDDING_PROVIDER=vertex_ai"
                 )
@@ -68,9 +65,7 @@ def __init__(self):
 
         # De-dupe while preserving order.
         seen: set[str] = set()
-        self._model_candidates = [
-            m for m in candidates if not (m in seen or seen.add(m))
-        ]
+        self._model_candidates = [m for m in candidates if not (m in seen or seen.add(m))]
         self.model = self._model_candidates[0]
         self.dimensions = config.embedding.dimensions
         self.batch_size = config.embedding.batch_size
@@ -83,19 +78,15 @@ def _embed(self, *, text: str, task_type: str) -> Any:
             output_dimensionality=int(self.dimensions) if self.dimensions else None,
         )
 
-        return self.client.models.embed_content(
-            model=self.model, contents=text, config=cfg
-        )
+        return self.client.models.embed_content(model=self.model, contents=text, config=cfg)
 
     @retry(
         stop=stop_after_attempt(5),
         wait=wait_exponential(multiplier=1, min=2, max=10),
         retry=retry_if_exception_type(Exception),
         reraise=True,
     )
-    def generate_embedding(
-        self, text: str, task_type: str = "RETRIEVAL_DOCUMENT"
-    ) -> list[float]:
+    def generate_embedding(self, text: str, task_type: str = "RETRIEVAL_DOCUMENT") -> list[float]:
         """Generate embedding for a single text."""
 
         last_err: Exception | None = None
@@ -143,9 +134,7 @@ def generate_embeddings_batch(
                 f"({len(batch)} texts)..."
             )
 
-            batch_embeddings = [
-                self.generate_embedding(text, task_type) for text in batch
-            ]
+            batch_embeddings = [self.generate_embedding(text, task_type) for text in batch]
             all_embeddings.extend(batch_embeddings)
 
         return all_embeddings

lib/gemini_finish_reason.py

@@ -45,6 +45,4 @@ def raise_if_retryable_finish_reason(response: Any) -> None:
         return
 
     finish_reason_name = normalize_finish_reason_name(finish_reason)
-    raise RetryableFinishReasonError(
-        f"Retryable finish_reason encountered: {finish_reason_name}"
-    )
+    raise RetryableFinishReasonError(f"Retryable finish_reason encountered: {finish_reason_name}")

lib/google_client.py

@@ -85,6 +85,4 @@ def _safe_json_parse(self, response_text: str, context: str = "") -> dict[str, A
             return json.loads(response_text)
         except json.JSONDecodeError:
             preview = response_text[:200] if response_text else ""
-            raise ValueError(
-                f"Failed to parse JSON response ({context}). Preview: {preview}..."
-            )
+            raise ValueError(f"Failed to parse JSON response ({context}). Preview: {preview}...")

lib/knowledge_graph/base_kg_seeder.py

@@ -256,9 +256,7 @@ def _generate_embeddings_for_nodes(self, node_ids: Any) -> None:
         if not node_id_list:
             return
 
-        check_query = (
-            "SELECT id, label FROM kg_nodes WHERE id = ANY(%s) AND embedding IS NULL"
-        )
+        check_query = "SELECT id, label FROM kg_nodes WHERE id = ANY(%s) AND embedding IS NULL"
         rows = self.postgres.execute_query(check_query, (node_id_list,))
 
         labels_to_embed = {row[0]: row[1] for row in rows}
@@ -277,9 +275,7 @@ def _generate_embeddings_for_nodes(self, node_ids: Any) -> None:
                 print(f"Error generating embeddings batch: {e}")
                 return
 
-            update_rows = [
-                (vector_literal(vec), nid) for nid, vec in zip(node_ids, embeddings)
-            ]
+            update_rows = [(vector_literal(vec), nid) for nid, vec in zip(node_ids, embeddings)]
 
             update_query = """
                 UPDATE kg_nodes

lib/knowledge_graph/kg_extractor.py

@@ -227,8 +227,7 @@ def _parse_edges_from_llm_data(
                     target_ref=target_ref,
                     evidence=evidence,
                     utterance_ids=utterance_ids,
-                    earliest_timestamp=earliest_timestamp_str
-                    or window.earliest_timestamp,
+                    earliest_timestamp=earliest_timestamp_str or window.earliest_timestamp,
                     earliest_seconds=earliest_seconds or window.earliest_seconds,
                     confidence=float(edge_data.get("confidence", 0.5)),
                 )
@@ -295,9 +294,7 @@ def canonicalize_and_store(
     ) -> dict[str, Any]:
         """Canonicalize nodes and edges and store them in Postgres."""
 
-        def _normalize_speaker_ref(
-            ref: str, window_speaker_ids: list[str]
-        ) -> str | None:
+        def _normalize_speaker_ref(ref: str, window_speaker_ids: list[str]) -> str | None:
             ref = (ref or "").strip()
             if not ref:
                 return None
@@ -353,9 +350,7 @@ def _normalize_speaker_ref(
             speaker_nodes_data = []
             for speaker_id in speaker_ids_seen:
                 meta = speaker_meta.get(speaker_id, {})
-                label = (
-                    meta.get("full_name") or meta.get("normalized_name") or speaker_id
-                )
+                label = meta.get("full_name") or meta.get("normalized_name") or speaker_id
                 aliases = []
                 for candidate in (
                     meta.get("full_name"),
@@ -435,13 +430,11 @@ def _normalize_speaker_ref(
                 target_id = temp_to_canonical.get(target_ref, target_ref)
 
                 if not (
-                    edge.source_ref.startswith("speaker_")
-                    or edge.source_ref in temp_to_canonical
+                    edge.source_ref.startswith("speaker_") or edge.source_ref in temp_to_canonical
                 ):
                     stats["links_to_known"] += 1
                 if not (
-                    edge.target_ref.startswith("speaker_")
-                    or edge.target_ref in temp_to_canonical
+                    edge.target_ref.startswith("speaker_") or edge.target_ref in temp_to_canonical
                 ):
                     stats["links_to_known"] += 1
 
@@ -554,13 +547,9 @@ def _embed_new_nodes(self, node_ids: list[str]) -> None:
 
         ids = [x[0] for x in to_embed]
         texts = [x[1] for x in to_embed]
-        embeddings = self.embedding.generate_embeddings_batch(
-            texts, task_type="RETRIEVAL_DOCUMENT"
-        )
+        embeddings = self.embedding.generate_embeddings_batch(texts, task_type="RETRIEVAL_DOCUMENT")
 
-        update_rows = [
-            (vector_literal(vec), node_id) for node_id, vec in zip(ids, embeddings)
-        ]
+        update_rows = [(vector_literal(vec), node_id) for node_id, vec in zip(ids, embeddings)]
         self.postgres.execute_batch(
             """
             UPDATE kg_nodes

lib/knowledge_graph/kg_store.py

@@ -19,9 +19,7 @@ def canonicalize_and_store(
     *,
     postgres: PostgresClient,
     embedding: GoogleEmbeddingClient,
-    results: list[
-        tuple[Window, list[dict[str, Any]], list[dict[str, Any]], str, bool, str | None]
-    ],
+    results: list[tuple[Window, list[dict[str, Any]], list[dict[str, Any]], str, bool, str | None]],
     youtube_video_id: str,
     kg_run_id: str,
     extractor_model: str,
@@ -187,13 +185,11 @@ def _normalize_speaker_ref(ref: str, window_speaker_ids: list[str]) -> str | Non
             target_id = temp_to_canonical.get(target_ref, target_ref)
 
             if not (
-                edge["source_ref"].startswith("speaker_")
-                or edge["source_ref"] in temp_to_canonical
+                edge["source_ref"].startswith("speaker_") or edge["source_ref"] in temp_to_canonical
             ):
                 stats["links_to_known"] += 1
             if not (
-                edge["target_ref"].startswith("speaker_")
-                or edge["target_ref"] in temp_to_canonical
+                edge["target_ref"].startswith("speaker_") or edge["target_ref"] in temp_to_canonical
             ):
                 stats["links_to_known"] += 1
 
@@ -274,9 +270,7 @@ def _normalize_speaker_ref(ref: str, window_speaker_ids: list[str]) -> str | Non
         )
         existing_ids = {row[0] for row in existing_rows}
 
-        filtered_edges = [
-            e for e in edges_data if e[1] in existing_ids and e[3] in existing_ids
-        ]
+        filtered_edges = [e for e in edges_data if e[1] in existing_ids and e[3] in existing_ids]
         stats["edges_skipped_missing_nodes"] = len(edges_data) - len(filtered_edges)
         stats["edges"] = len(filtered_edges)
 
@@ -323,13 +317,9 @@ def _embed_new_nodes(
 
     ids = [x[0] for x in to_embed]
     texts = [x[1] for x in to_embed]
-    embeddings = embedding.generate_embeddings_batch(
-        texts, task_type="RETRIEVAL_DOCUMENT"
-    )
+    embeddings = embedding.generate_embeddings_batch(texts, task_type="RETRIEVAL_DOCUMENT")
 
-    update_rows = [
-        (vector_literal(vec), node_id) for node_id, vec in zip(ids, embeddings)
-    ]
+    update_rows = [(vector_literal(vec), node_id) for node_id, vec in zip(ids, embeddings)]
     postgres.execute_batch(
         """
         UPDATE kg_nodes

lib/knowledge_graph/model_compare.py

@@ -134,12 +134,8 @@ def canonicalize_edges(
     out: list[dict[str, Any]] = []
 
     for e in edges:
-        source_ref = normalize_speaker_ref(
-            str(e.get("source_ref", "")), window_speaker_ids
-        )
-        target_ref = normalize_speaker_ref(
-            str(e.get("target_ref", "")), window_speaker_ids
-        )
+        source_ref = normalize_speaker_ref(str(e.get("source_ref", "")), window_speaker_ids)
+        target_ref = normalize_speaker_ref(str(e.get("target_ref", "")), window_speaker_ids)
         if source_ref is None or target_ref is None:
             continue
 
@@ -269,9 +265,7 @@ def collect_signatures(
                 window_speaker_ids=r.window_speaker_ids,
             )
             for e in canon_edges:
-                sigs.add(
-                    edge_signature_strict(e) if strict else edge_signature_loose(e)
-                )
+                sigs.add(edge_signature_strict(e) if strict else edge_signature_loose(e))
         return sigs
 
     sigs_loose: dict[str, set[tuple]] = {}

lib/knowledge_graph/oss_kg_extractor.py

@@ -261,9 +261,7 @@ def extract_from_concept_window(
             data_pass1 = self._parse_json_response(raw_response_pass1)
 
             # Normalize pass1 output
-            normalize_utterance_ids_in_data(
-                data_pass1, youtube_video_id=youtube_video_id
-            )
+            normalize_utterance_ids_in_data(data_pass1, youtube_video_id=youtube_video_id)
             normalize_evidence_in_data(data_pass1, window_text=window.text)
 
             pass1_parse_success = True
@@ -331,9 +329,7 @@ def extract_from_concept_window(
                     data_pass2 = self._parse_json_response(raw_response_pass2)
 
                     # Normalize pass2 output
-                    normalize_utterance_ids_in_data(
-                        data_pass2, youtube_video_id=youtube_video_id
-                    )
+                    normalize_utterance_ids_in_data(data_pass2, youtube_video_id=youtube_video_id)
                     normalize_evidence_in_data(data_pass2, window_text=window.text)
 
                     # Merge additions
@@ -356,9 +352,7 @@ def extract_from_concept_window(
                     final_data = self._parse_json_response(raw_response_pass2)
 
                     # Normalize pass2 output
-                    normalize_utterance_ids_in_data(
-                        final_data, youtube_video_id=youtube_video_id
-                    )
+                    normalize_utterance_ids_in_data(final_data, youtube_video_id=youtube_video_id)
                     normalize_evidence_in_data(final_data, window_text=window.text)
 
                 pass2_error = None
@@ -406,8 +400,7 @@ def extract_from_concept_window(
                     "target_ref": target_ref,
                     "evidence": evidence,
                     "utterance_ids": utterance_ids,
-                    "earliest_timestamp": earliest_timestamp_str
-                    or window.earliest_timestamp,
+                    "earliest_timestamp": earliest_timestamp_str or window.earliest_timestamp,
                     "earliest_seconds": earliest_seconds or window.earliest_seconds,
                     "confidence": float(edge_data.get("confidence", 0.5)),
                 }

lib/knowledge_graph/oss_two_pass.py

@@ -13,9 +13,7 @@
 from lib.knowledge_graph.model_compare import normalize_speaker_ref
 
 
-def normalize_utterance_ids_in_data(
-    data: dict[str, Any], *, youtube_video_id: str
-) -> None:
+def normalize_utterance_ids_in_data(data: dict[str, Any], *, youtube_video_id: str) -> None:
     """Normalize utterance_ids to full "{youtube_video_id}:<seconds>" strings.
 
     Some models sometimes output bare seconds like "1851". The transcript windows
@@ -176,15 +174,11 @@ def validate_kg_llm_data(
     edges = data.get("edges")
     if not isinstance(nodes, list):
         issues.append(
-            ValidationIssue(
-                code="nodes_new_not_list", message="nodes_new must be a list"
-            )
+            ValidationIssue(code="nodes_new_not_list", message="nodes_new must be a list")
         )
         nodes = []
     if not isinstance(edges, list):
-        issues.append(
-            ValidationIssue(code="edges_not_list", message="edges must be a list")
-        )
+        issues.append(ValidationIssue(code="edges_not_list", message="edges must be a list"))
         edges = []
 
     for i, n in enumerate(nodes):
@@ -293,11 +287,7 @@ def validate_kg_llm_data(
                 )
             )
         else:
-            bad = [
-                str(uid)
-                for uid in utterance_ids
-                if str(uid) not in window_utterance_ids
-            ]
+            bad = [str(uid) for uid in utterance_ids if str(uid) not in window_utterance_ids]
             if bad:
                 issues.append(
                     ValidationIssue(
@@ -343,9 +333,7 @@ def should_run_second_pass(
     if mode == TwoPassMode.ALWAYS:
         return True, "always"
     if mode == TwoPassMode.ON_FAIL:
-        return (
-            not pass1_parse_success
-        ), "parse_fail" if not pass1_parse_success else None
+        return (not pass1_parse_success), "parse_fail" if not pass1_parse_success else None
     if mode == TwoPassMode.ON_LOW_EDGES:
         return (pass1_parse_success and edge_count < min_edges), (
             "low_edges" if pass1_parse_success and edge_count < min_edges else None
@@ -602,9 +590,7 @@ def merge_oss_additions(
     del_edges: list[Any] = del_edges_any if isinstance(del_edges_any, list) else []
 
     existing_ids = {
-        str(n.get("temp_id"))
-        for n in base_nodes
-        if isinstance(n, dict) and n.get("temp_id")
+        str(n.get("temp_id")) for n in base_nodes if isinstance(n, dict) and n.get("temp_id")
     }
 
     # Build remap for added nodes.