Fixing not using set_llm_session_ids from fh-llm-client (#792)

Author: jamesbraza

paperqa/docs.py

@@ -20,6 +20,7 @@
     LLMModel,
     LLMResult,
 )
+from llmclient.types import set_llm_session_ids
 from pydantic import (
     BaseModel,
     ConfigDict,
@@ -39,14 +40,7 @@
 from paperqa.prompts import CANNOT_ANSWER_PHRASE
 from paperqa.readers import read_doc
 from paperqa.settings import MaybeSettings, get_settings
-from paperqa.types import (
-    Doc,
-    DocDetails,
-    DocKey,
-    PQASession,
-    Text,
-    set_llm_session_ids,
-)
+from paperqa.types import Doc, DocDetails, DocKey, PQASession, Text
 from paperqa.utils import (
     gather_with_concurrency,
     get_loop,

paperqa/types.py

@@ -1,12 +1,10 @@
 from __future__ import annotations
 
-import contextvars
 import logging
 import os
 import re
 import warnings
 from collections.abc import Collection
-from contextlib import contextmanager
 from datetime import datetime
 from typing import Any, ClassVar, cast
 from uuid import UUID, uuid4
@@ -40,18 +38,6 @@
 DocKey = Any
 logger = logging.getLogger(__name__)
 
-# A context var that will be unique to threads/processes
-cvar_session_id = contextvars.ContextVar[UUID | None]("session_id", default=None)
-
-
-@contextmanager
-def set_llm_session_ids(session_id: UUID):
-    token = cvar_session_id.set(session_id)
-    try:
-        yield
-    finally:
-        cvar_session_id.reset(token)
-
 
 class Doc(Embeddable):
     docname: str

tests/test_paperqa.py

@@ -19,6 +19,7 @@
     HybridEmbeddingModel,
     LiteLLMEmbeddingModel,
     LLMModel,
+    LLMResult,
     SparseEmbeddingModel,
 )
 from pytest_subtests import SubTests
@@ -583,20 +584,19 @@ def test_query(docs_fixture) -> None:
     docs_fixture.query("Is XAI usable in chemistry?")
 
 
-def test_llmresult_callback(docs_fixture) -> None:
-    my_results = []
-
-    async def my_callback(result) -> None:
-        my_results.append(result)
+def test_llmresult_callback(docs_fixture: Docs) -> None:
+    my_results: list[LLMResult] = []
 
     settings = Settings.from_name("fast")
     summary_llm = settings.get_summary_llm()
-    summary_llm.llm_result_callback = my_callback
+    summary_llm.llm_result_callback = my_results.append
     docs_fixture.get_evidence(
         "What is XAI?", settings=settings, summary_llm_model=summary_llm
     )
     assert my_results
+    assert len(my_results) >= 1, "Expected the callback to append results"
     assert my_results[0].name
+    assert my_results[0].session_id
 
 
 def test_duplicate(stub_data_dir: Path) -> None: