diff --git a/samples/tools/profiler/profiler/llm.py b/samples/tools/profiler/profiler/llm.py index 9d897f471..a39db5d67 100644 --- a/samples/tools/profiler/profiler/llm.py +++ b/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 diff --git a/samples/tools/profiler/profiler/profiler.py b/samples/tools/profiler/profiler/profiler.py index bb39d4c01..fccd44e70 100644 --- a/samples/tools/profiler/profiler/profiler.py +++ b/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 = ( diff --git a/samples/tools/profiler/profiler/state.py b/samples/tools/profiler/profiler/state.py index 9671bec83..e8493c194 100644 --- a/samples/tools/profiler/profiler/state.py +++ b/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={