Make changes

Author: DylanRussell

util/opentelemetry-util-genai/src/opentelemetry/util/genai/_upload/completion_hook.py

@@ -19,10 +19,11 @@
 import logging
 import posixpath
 import threading
+from collections import OrderedDict
 from concurrent.futures import Future, ThreadPoolExecutor
 from contextlib import ExitStack
 from dataclasses import asdict, dataclass
-from functools import lru_cache, partial
+from functools import partial
 from os import environ
 from time import time
 from typing import Any, Callable, Final, Literal
@@ -78,6 +79,12 @@ class CompletionRefs:
 UploadData = dict[str, Callable[[], JsonEncodeable]]
 
 
+def is_system_instructions_hashable(
+    system_instruction: list[types.MessagePart],
+) -> bool:
+    return all(isinstance(x, types.Text) for x in system_instruction)
+
+
 class UploadCompletionHook(CompletionHook):
     """An completion hook using ``fsspec`` to upload to external storage
 
@@ -98,10 +105,13 @@ def __init__(
         base_path: str,
         max_size: int = 20,
         upload_format: Format | None = None,
+        lru_cache_max_size: int = 1024,
     ) -> None:
         self._max_size = max_size
         self._fs, base_path = fsspec.url_to_fs(base_path)
         self._base_path = self._fs.unstrip_protocol(base_path)
+        self.lru_dict = OrderedDict()
+        self.lru_cache_max_size = lru_cache_max_size
 
         if upload_format not in _FORMATS + (None,):
             raise ValueError(
@@ -159,7 +169,7 @@ def _calculate_ref_path(
         # TODO: experimental with using the trace_id and span_id, or fetching
         # gen_ai.response.id from the active span.
         system_instruction_hash = None
-        if all(isinstance(x, types.Text) for x in system_instruction):
+        if is_system_instructions_hashable(system_instruction):
             # Get a hash of the text.
             system_instruction_hash = hashlib.sha256(
                 "\n".join(x.content for x in system_instruction).encode(  # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue, reportUnknownArgumentType]
@@ -181,10 +191,18 @@ def _calculate_ref_path(
             ),
         )
 
-    @lru_cache(maxsize=512)
     def _file_exists(self, path: str) -> bool:
+        if path in self.lru_dict:
+            self.lru_dict.move_to_end(path)
+            return True
         # https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.exists
-        return self._fs.exists(path)
+        file_exists = self._fs.exists(path)
+        if not file_exists:
+            return False
+        self.lru_dict[path] = True
+        if len(self.lru_dict) > self.lru_cache_max_size:
+            self.lru_dict.popitem(last=False)
+        return True
 
     def _do_upload(
         self, path: str, json_encodeable: Callable[[], JsonEncodeable]
@@ -214,6 +232,11 @@ def _do_upload(
                 gen_ai_json_dump(message, file)
                 file.write("\n")
 
+        if "_system_instruction" in path:
+            self.lru_dict[path] = True
+            if len(self.lru_dict) > self.lru_cache_max_size:
+                self.lru_dict.popitem(last=False)
+
     def on_completion(
         self,
         *,

util/opentelemetry-util-genai/tests/test_upload.py

@@ -125,8 +125,7 @@ def setUp(self):
         self.mock_fs.exists.return_value = False
 
         self.hook = UploadCompletionHook(
-            base_path=BASE_PATH,
-            max_size=MAXSIZE,
+            base_path=BASE_PATH, max_size=MAXSIZE, lru_cache_max_size=5
         )
 
     def tearDown(self) -> None:
@@ -184,7 +183,7 @@ def test_system_insruction_is_hashed_to_avoid_reupload(self):
             log_record=record,
         )
         # Wait a bit for file upload to finish..
-        time.sleep(2)
+        time.sleep(0.5)
         self.mock_fs.exists.return_value = True
         self.hook.on_completion(
             inputs=[],
@@ -206,6 +205,39 @@ def test_system_insruction_is_hashed_to_avoid_reupload(self):
             f"{expected_hash}_system_instruction.json",
         )
 
+    def test_lru_cache_works(self):
+        record = LogRecord()
+        self.hook.on_completion(
+            inputs=[],
+            outputs=[],
+            system_instruction=FAKE_SYSTEM_INSTRUCTION,
+            log_record=record,
+        )
+        # Wait a bit for file upload to finish..
+        time.sleep(0.5)
+        assert record.attributes is not None
+        self.assertTrue(
+            self.hook._file_exists(
+                record.attributes["gen_ai.system_instructions_ref"]
+            )
+        )
+        # LRU cache has a size of 5. So only AFTER 5 uploads should the original file be removed from the cache.
+        for x in range(5):
+            self.assertTrue(
+                record.attributes["gen_ai.system_instructions_ref"]
+                in self.hook.lru_dict
+            )
+            self.hook.on_completion(
+                inputs=[],
+                outputs=[],
+                system_instruction=[types.Text(content=str(x))],
+            )
+        print(self.hook.lru_dict)
+        self.assertFalse(
+            record.attributes["gen_ai.system_instructions_ref"]
+            in self.hook.lru_dict
+        )
+
     def test_upload_when_inputs_outputs_empty(self):
         record = LogRecord()
         self.hook.on_completion(