updated typechecks

Author: AkhileshNegi

backend/app/core/storage_utils.py

@@ -8,10 +8,11 @@
 import json
 import logging
 import mimetypes
+from collections.abc import Mapping, Sequence
 from datetime import datetime
 from io import BytesIO
 from pathlib import Path
-from typing import Literal
+from typing import Any, Literal
 from urllib.parse import unquote, urlparse
 from uuid import UUID
 
@@ -109,7 +110,7 @@ def upload_to_object_store(
 
 def upload_jsonl_to_object_store(
     storage: CloudStorage,
-    results: list[dict],
+    results: Sequence[Mapping[str, Any]],
     filename: str,
     subdirectory: str,
     format: Literal["json", "jsonl"] = "jsonl",

backend/app/crud/evaluations/core.py

@@ -1,9 +1,10 @@
 import logging
-from typing import Any
+from collections.abc import Mapping, Sequence
+from typing import Any, cast
 from uuid import UUID
 
 from langfuse import Langfuse
-from sqlmodel import Session, select
+from sqlmodel import Session, col, select
 
 from app.core.cloud.storage import get_cloud_storage
 from app.core.db import engine
@@ -119,7 +120,7 @@ def list_evaluation_runs(
     project_id: int,
     limit: int = 50,
     offset: int = 0,
-) -> list[EvaluationRun]:
+) -> Sequence[EvaluationRun]:
     """
     List all evaluation runs for an organization and project.
 
@@ -138,7 +139,7 @@ def list_evaluation_runs(
         .where(EvaluationRun.organization_id == organization_id)
         .where(EvaluationRun.project_id == project_id)
         .where(EvaluationRun.type == EvaluationType.TEXT.value)
-        .order_by(EvaluationRun.inserted_at.desc())
+        .order_by(col(EvaluationRun.inserted_at).desc())
         .limit(limit)
         .offset(offset)
     )
@@ -301,7 +302,7 @@ def get_or_fetch_score(
         logger.info(
             f"[get_or_fetch_score] Returning existing score | evaluation_id={eval_run.id}"
         )
-        return eval_run.score
+        return cast(EvaluationScore, eval_run.score)
 
     logger.info(
         f"[get_or_fetch_score] Fetching score from Langfuse | "
@@ -339,7 +340,7 @@ def get_or_fetch_score(
     update_evaluation_run(
         session=session,
         eval_run=eval_run,
-        update=EvaluationRunUpdate(score=score),
+        update=EvaluationRunUpdate(score=cast(dict[str, Any], score)),
     )
 
     total_traces = len(score.get("traces", []))
@@ -356,7 +357,7 @@ def _upload_score_traces(
     session: Session,
     eval_run_id: int,
     project_id: int,
-    traces: list[dict[str, Any]],
+    traces: Sequence[Mapping[str, Any]],
 ) -> str | None:
     """Upload per-trace records to S3 for an evaluation run.
 
@@ -399,7 +400,7 @@ def persist_score_traces(
     eval_run_id: int,
     organization_id: int,
     project_id: int,
-    traces: list[dict[str, Any]],
+    traces: Sequence[Mapping[str, Any]],
 ) -> EvaluationRun | None:
     """Persist the Q&A trace skeleton to S3 and record the ``score_trace_url``
     pointer, WITHOUT touching the ``score`` column (keeps the run score-less
@@ -484,12 +485,13 @@ def save_score(
         # IF TRACES DATA IS STORED IN S3 URL THEN HERE WE ARE JUST STORING THE SUMMARY SCORE
         # TODO: Evaluate whether this behaviour is needed or completely discard the storing data in db
         if score_trace_url:
-            db_score = {"summary_scores": summary_score}
-            if score.get("overall") is not None:
-                db_score["overall"] = score["overall"]
+            db_score: dict[str, Any] = {"summary_scores": summary_score}
+            overall = score.get("overall")
+            if overall is not None:
+                db_score["overall"] = overall
         else:
             # fallback to store data in db if failed to store in s3
-            db_score = score
+            db_score = cast(dict[str, Any], score)
 
         update_evaluation_run(
             session=session,
@@ -539,6 +541,8 @@ def group_traces_by_question_id(
 
     for trace in traces:
         question_id = trace.get("question_id")
+        if question_id is None:
+            continue
         if question_id not in groups:
             groups[question_id] = []
         groups[question_id].append(trace)

backend/app/crud/evaluations/cron_utils.py

@@ -11,7 +11,7 @@
 from typing import Any
 
 from sqlalchemy import Integer
-from sqlmodel import Session, select
+from sqlmodel import Session, col, select
 
 from app.core.batch import (
     BatchJobState,
@@ -49,7 +49,7 @@ def fetch_processing_runs(
     statement = select(EvaluationRun).where(
         EvaluationRun.type == eval_type,
         EvaluationRun.status == "processing",
-        EvaluationRun.batch_job_id.is_not(None),
+        col(EvaluationRun.batch_job_id).is_not(None),
     )
     return list(session.exec(statement).all())
 

backend/app/crud/evaluations/dataset.py

@@ -14,7 +14,7 @@
 from fastapi import HTTPException
 from sqlalchemy import Integer, cast
 from sqlalchemy.exc import IntegrityError
-from sqlmodel import Session, select
+from sqlmodel import Session, col, select
 
 from app.core.cloud.storage import CloudStorage
 from app.core.config import settings
@@ -228,7 +228,7 @@ def list_datasets(
         )
 
     statement = (
-        statement.order_by(EvaluationDataset.inserted_at.desc())
+        statement.order_by(col(EvaluationDataset.inserted_at).desc())
         .limit(limit)
         .offset(offset)
     )

backend/app/crud/evaluations/fast.py

@@ -19,7 +19,7 @@
 
 import logging
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import openai
@@ -80,12 +80,14 @@
 from app.crud.evaluations.score import (
     COSINE_SCORE_COMMENT,
     COSINE_SCORE_NAME,
+    DEFAULT_CATEGORY,
     JUDGE_FAILED_REASON,
     UNSCOREABLE_EMBEDDING_FAILED,
     UNSCOREABLE_EMPTY_GROUND_TRUTH,
     UNSCOREABLE_EMPTY_OUTPUT,
     EvaluationScore,
     OverallSummary,
+    SummaryScore,
     TraceData,
     TraceScore,
     compute_overall_summary,
@@ -599,11 +601,13 @@ def _merge_response_chunks(
 
     results: list[dict[str, Any]] = []
     for chunk_index in sorted(chunk_job_by_index):
+        raw_output_url = chunk_job_by_index[chunk_index].raw_output_url
+        assert raw_output_url is not None  # guaranteed by the filter above
         results.extend(
             _load_unit_from_s3(
                 session=session,
                 project_id=eval_run.project_id,
-                url=chunk_job_by_index[chunk_index].raw_output_url,
+                url=raw_output_url,
             )
         )
 
@@ -881,7 +885,7 @@ def _attach_metric_scores(
     *,
     spec: JudgeMetricSpec,
     judge_results: dict[str, JudgeResult],
-    summary_scores: list[dict[str, Any]],
+    summary_scores: list[SummaryScore],
 ) -> None:
     """Append one metric's run-level summary score from the combined results.
 
@@ -959,7 +963,7 @@ def _stage3_score_and_trace(
     similarities: list[float] = []
     unscoreable: dict[str, str] = {}  # {ref: reason}
     write_items: list[dict[str, Any]] = []
-    summary_scores: list[dict[str, Any]] = []
+    summary_scores: list[SummaryScore] = []
     overall: OverallSummary | None = None
 
     if is_judge_run:
@@ -993,6 +997,7 @@ def _stage3_score_and_trace(
                 else:
                     unscoreable[ref] = UNSCOREABLE_EMBEDDING_FAILED
                 continue
+            assert embedding_pair is not None  # guaranteed by has_embeddings above
             cosine = calculate_cosine_similarity(
                 embedding_pair["output_embedding"],
                 embedding_pair["ground_truth_embedding"],
@@ -1171,7 +1176,7 @@ def _stage3_score_and_trace(
     traces: list[TraceData] = []
     for response in response_results:
         item_id = response["item_id"]
-        ref = item_id_to_ref.get(item_id, item_id)
+        ref = item_id_to_ref[item_id] if item_id in item_id_to_ref else item_id
         trace_scores: list[TraceScore] = []
         # v2 carries no cosine score or placeholder — only the judge scores below.
         if not is_judge_run:
@@ -1253,6 +1258,7 @@ def _stage3_score_and_trace(
                 "llm_answer": response.get("generated_output", ""),
                 "ground_truth_answer": response.get("ground_truth", ""),
                 "question_id": response.get("question_id"),
+                "category": response.get("category") or DEFAULT_CATEGORY,
                 "scores": trace_scores,
             }
         )
@@ -1409,7 +1415,7 @@ def run_fast_evaluation(
     )
     if saved is not None:
         eval_run = saved
-        eval_run.score = score
+        eval_run.score = cast(dict[str, Any], score)
 
     logger.info(
         f"[run_fast_evaluation] {log_prefix} Fast evaluation completed | "

backend/app/crud/evaluations/langfuse.py

@@ -13,6 +13,7 @@
 from typing import Any
 
 from langfuse import Langfuse
+from langfuse.api.commons.types.score_v1 import ScoreV1_Text
 
 from app.core.langfuse.langfuse import format_langfuse_error, set_trace_attributes
 from app.crud.evaluations.merge import compute_summary_scores
@@ -280,6 +281,12 @@ def update_traces_with_cosine_scores(
             comment = f"Cannot compute: {reason}"
         else:
             value = score_item.get("cosine_similarity")
+            if value is None:
+                logger.warning(
+                    "[update_traces_with_cosine_scores] "
+                    f"Score item missing cosine_similarity, skipping | trace_id={trace_id}"
+                )
+                continue
             comment = COSINE_SCORE_COMMENT
 
         try:
@@ -517,7 +524,7 @@ def _fetch_single_trace(trace_id: str) -> TraceData | None:
                 "question": "",
                 "llm_answer": "",
                 "ground_truth_answer": "",
-                "question_id": "",
+                "question_id": None,
                 "scores": [],
             }
 
@@ -552,7 +559,11 @@ def _fetch_single_trace(trace_id: str) -> TraceData | None:
             if trace.scores:
                 for score in trace.scores:
                     score_name = score.name
-                    score_value = score.value
+                    score_value = (
+                        score.string_value
+                        if isinstance(score, ScoreV1_Text)
+                        else score.value
+                    )
                     score_comment = score.comment
                     # Get data_type from Langfuse score, default to NUMERIC
                     data_type = getattr(score, "data_type", None) or "NUMERIC"

backend/app/crud/evaluations/merge.py

@@ -9,7 +9,7 @@
 import itertools
 import logging
 from collections import Counter
-from typing import Any
+from typing import cast
 
 import numpy as np
 
@@ -18,6 +18,7 @@
     COSINE_SCORE_NAME,
     DEFAULT_CATEGORY,
     EvaluationScore,
+    NumericSummaryScore,
     SummaryScore,
     TraceData,
     TraceScore,
@@ -126,9 +127,11 @@ def apply_cosine_breakdown(
     breakdown = summarize_unscoreable(unscoreable)
     for entry in summary_scores:
         if entry.get("name") == COSINE_SCORE_NAME:
+            # Cosine is always a NUMERIC score; only that shape carries these fields.
+            numeric_entry = cast(NumericSummaryScore, entry)
             if total_items is not None:
-                entry["total_items"] = total_items
-            entry["unscoreable"] = breakdown
+                numeric_entry["total_items"] = total_items
+            numeric_entry["unscoreable"] = breakdown
     return summary_scores
 
 
@@ -177,7 +180,7 @@ def _merge_single_trace(existing: TraceData, fresh: TraceData) -> TraceData:
     for fresh_score in fresh.get("scores", []):
         merged_scores_by_name[fresh_score["name"]] = fresh_score
 
-    merged: dict[str, Any] = {
+    merged: TraceData = {
         "trace_id": fresh.get("trace_id") or existing.get("trace_id", ""),
         "question": fresh.get("question") or existing.get("question", ""),
         "llm_answer": fresh.get("llm_answer") or existing.get("llm_answer", ""),
@@ -205,6 +208,7 @@ def _reconcile_trace(
     or ``updated``. Exactly one of ``existing``/``fresh`` may be None, never both.
     """
     if existing is None:
+        assert fresh is not None, "exactly one of existing/fresh must be set"
         return fresh, "added"
     if fresh is None:
         return existing, "reused"