Add caching

samples/tools/profiler/profiler/llm.py

@@ -1,5 +1,8 @@
 from typing import Protocol, runtime_checkable, List, Optional
 import os
+
+import hashlib
+from diskcache import Cache
 import openai
 
 from .message import OpenAIMessage
@@ -24,8 +27,9 @@ class OpenAIJSONService:
             except KeyError:
                 raise ValueError("Either set api_key arg or set OPENAI_API_KEY env")
         self._client = openai.Client(api_key=api_key)
+        self._cache_dir = ".cache"
 
-    def create(self, messages: List[OpenAIMessage]) -> OpenAIMessage:
+    def _create(self, messages: List[OpenAIMessage]) -> OpenAIMessage:
         response = self._client.chat.completions.create(
             model=self.MODEL,
             messages=[m.to_dict() for m in messages],
@@ -34,3 +38,17 @@ class OpenAIJSONService:
         )
         first_choice = response.choices[0]
         return OpenAIMessage(role=first_choice.message.role, content=first_choice.message.content)
+
+    def create(self, messages: List[OpenAIMessage]) -> OpenAIMessage:
+
+        all_messages = "".join([m.content for m in messages])
+        key = hashlib.sha256(all_messages.encode()).hexdigest()
+
+        with Cache(self._cache_dir) as cache:
+            content = cache.get(key, None)
+            if content is not None:
+                print("Cache hit")
+                return OpenAIMessage(role="assistant", content=content)
+            response = self._create(messages)
+            cache.set(key, response.content)
+            return response

samples/tools/profiler/profiler/profiler.py

@@ -94,7 +94,7 @@ class Profiler:
         state_space = self.state_space.filter_states(condition=source_or_role_in_tags)
 
         state_space_str = ""
-        for state in state_space:
+        for state in state_space.sorted_states():
             state_space_str += f"{state.name}: {state.description}" + "\n"
 
         prompt = (

samples/tools/profiler/profiler/state.py

@@ -45,6 +45,9 @@ class StateSpace:
         filtered_states = {state for state in self.states if condition(state)}
         return StateSpace(filtered_states)
 
+    def sorted_states(self) -> List[State]:
+        return sorted(self.states, key=lambda state: state.name)
+
 
 DEFAULT_STATE_SPACE = StateSpace(
     states={