Fix LLMTimestampFileLoggingMiddleware (#213)

* Fix LLMTimestampFileLoggingMiddleware to create a single file for request-response pair; add path into parameters

* Add delay to MockLLM

* Update tests

* Add * in __init__

* Update docs

* Simplify syntax

council/llm/llm_function/llm_middleware.py

@@ -2,6 +2,7 @@
 
 import hashlib
 import json
+import os
 import time
 from collections import OrderedDict
 from enum import Enum
@@ -240,7 +241,7 @@ def __call__(self, llm: LLMBase, execute: ExecuteLLMRequest, request: LLMRequest
 
     def _log(self, content: str) -> None:
         if self.context_logger is None:
-            raise RuntimeError("Calling LLMLoggingMiddleware._log() outside of __call__()")
+            raise RuntimeError("Context logger not set - calling LLMLoggingMiddleware._log() outside of __call__()")
 
         self.context_logger.info(content)
 
@@ -269,20 +270,33 @@ class LLMTimestampFileLoggingMiddleware(LLMLoggingMiddlewareBase):
 
     def __init__(
         self,
-        prefix: str,
         strategy: LLMLoggingStrategy = LLMLoggingStrategy.Verbose,
+        *,
+        path: str = ".",
+        filename_prefix: Optional[str] = None,
         component_name: Optional[str] = None,
     ) -> None:
         super().__init__(strategy, component_name)
-        self.prefix = prefix
+        self.prefix = f"{filename_prefix}_" if filename_prefix is not None else ""
+        self.path = path
         self._lock = Lock()
+        self._current_filename: Optional[str] = None
 
-    def _log(self, content: str) -> None:
-        """Write content to a new log file with timestamp"""
+    def __call__(self, llm: LLMBase, execute: ExecuteLLMRequest, request: LLMRequest) -> LLMResponse:
         timestamp = time.strftime("%Y-%m-%d_%H-%M-%S")
-        filename = f"{self.prefix}_{timestamp}.log"
+        self._current_filename = os.path.join(self.path, f"{self.prefix}{timestamp}.log")
+        response = super().__call__(llm, execute, request)
+        self._current_filename = None
+        return response
+
+    def _log(self, content: str) -> None:
+        """Write content to the current log file"""
+        if self._current_filename is None:
+            raise RuntimeError(
+                "Current log filename not set - calling LLMTimestampFileLoggingMiddleware._log() outside of __call__()"
+            )
 
-        self.append_to_file(self._lock, file_path=filename, content=content)
+        self.append_to_file(self._lock, file_path=self._current_filename, content=content)
 
 
 class LLMRetryMiddleware:

council/mocks/mock_llm.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import time
 from typing import Any, Iterable, List, Optional, Protocol, Sequence
 
 from council import LLMContext
@@ -54,26 +55,29 @@ def from_spec(cls, spec: LLMConfigSpec) -> MockLLMConfiguration:
 
 
 class MockLLM(LLMBase[MockLLMConfiguration]):
-    def __init__(self, action: Optional[LLMMessagesToStr] = None, token_limit: int = -1) -> None:
+    def __init__(self, action: Optional[LLMMessagesToStr] = None, token_limit: int = -1, delay: float = 0.0) -> None:
         super().__init__(configuration=MockLLMConfiguration("mock"), token_counter=MockTokenCounter(token_limit))
         self._action = action
+        self._delay = delay
 
     def _post_chat_request(self, context: LLMContext, messages: Sequence[LLMMessage], **kwargs: Any) -> LLMResult:
+        if self._delay > 0:
+            time.sleep(self._delay)
         choices = self._action(messages) if self._action is not None else [f"{self.__class__.__name__}"]
         return LLMResult(choices=choices, consumptions=[Consumption.call(1, "mock_llm")])
 
     @staticmethod
-    def from_responses(responses: List[str]) -> MockLLM:
-        return MockLLM(action=(lambda x: responses))
+    def from_responses(responses: List[str], delay: float = 0.0) -> MockLLM:
+        return MockLLM(action=(lambda x: responses), delay=delay)
 
     @staticmethod
-    def from_response(response: str) -> MockLLM:
-        return MockLLM(action=(lambda x: [response]))
+    def from_response(response: str, delay: float = 0.0) -> MockLLM:
+        return MockLLM(action=(lambda x: [response]), delay=delay)
 
     @staticmethod
-    def from_multi_line_response(responses: Iterable[str]) -> MockLLM:
+    def from_multi_line_response(responses: Iterable[str], delay: float = 0.0) -> MockLLM:
         response = "\n".join(responses)
-        return MockLLM(action=(lambda x: [response]))
+        return MockLLM(action=(lambda x: [response]), delay=delay)
 
 
 class MockErrorLLM(LLMBase):

docs/source/reference/llm.md

@@ -95,7 +95,9 @@ Middleware components allow you to enhance LLM interactions by modifying request
 Core middlewares:
 
 - Caching: {class}`~council.llm.LLMCachingMiddleware`
-- Logging: {class}`~council.llm.LLMLoggingMiddleware` and {class}`~council.llm.LLMFileLoggingMiddleware`
+- Logging: 
+  - Context logger: {class}`~council.llm.LLMLoggingMiddleware`
+  - Files: {class}`~council.llm.LLMFileLoggingMiddleware` and {class}`~council.llm.LLMTimestampFileLoggingMiddleware`
 
 Middleware management:
 

tests/unit/llm/test_llm_middleware.py

@@ -69,25 +69,27 @@ def test_with_retry_no_error(self):
 class TestLlmTimestampFileLoggingMiddleware(unittest.TestCase):
     def setUp(self) -> None:
         self._llm = MockLLM.from_response("USD")
+        self._llm_with_delay = MockLLM.from_response("USD", delay=1.0)
         self._temp_dir = tempfile.mkdtemp()
 
+        self._log_prefix = "test_llm"
+
     def tearDown(self) -> None:
         shutil.rmtree(self._temp_dir)
 
     def test_creates_separate_log_files(self):
         messages = [LLMMessage.user_message("Give me an example of a currency")]
         request = LLMRequest.default(messages)
 
-        log_prefix = os.path.join(self._temp_dir, "test_llm")
         chain = LLMMiddlewareChain(self._llm)
-        chain.add_middleware(LLMTimestampFileLoggingMiddleware(prefix=log_prefix))
+        chain.add_middleware(LLMTimestampFileLoggingMiddleware(path=self._temp_dir, filename_prefix=self._log_prefix))
 
         chain.execute(request)
         time.sleep(1)  # ensure different timestamps
 
         chain.execute(request)
 
-        log_files = glob.glob(os.path.join(self._temp_dir, "test_llm_*.log"))
+        log_files = glob.glob(os.path.join(self._temp_dir, f"{self._log_prefix}_*.log"))
         self.assertEqual(2, len(log_files))
 
         for log_file in log_files:
@@ -96,3 +98,22 @@ def test_creates_separate_log_files(self):
                 self.assertIn("LLM input", content)
                 self.assertIn("LLM output", content)
                 self.assertIn("USD", content)
+
+    def test_with_delay(self):
+        messages = [LLMMessage.user_message("Give me an example of a currency")]
+        request = LLMRequest.default(messages)
+
+        chain = LLMMiddlewareChain(self._llm_with_delay)
+        chain.add_middleware(LLMTimestampFileLoggingMiddleware(path=self._temp_dir, filename_prefix=self._log_prefix))
+
+        chain.execute(request)
+
+        log_files = glob.glob(os.path.join(self._temp_dir, f"{self._log_prefix}_*.log"))
+        self.assertEqual(1, len(log_files))
+
+        for log_file in log_files:
+            with open(log_file, "r", encoding="utf-8") as f:
+                content = f.read()
+                self.assertIn("LLM input", content)
+                self.assertIn("LLM output", content)
+                self.assertIn("USD", content)