From dc0faa5a79c13c9b65228153b57e42956c2a17b8 Mon Sep 17 00:00:00 2001 From: David Newell Date: Wed, 13 Mar 2024 15:57:22 +0000 Subject: [PATCH] chore: add input to clickhouse rows (#20901) --- ee/session_recordings/ai/embeddings_runner.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ee/session_recordings/ai/embeddings_runner.py b/ee/session_recordings/ai/embeddings_runner.py index 708b156c7981c..101c7175acb61 100644 --- a/ee/session_recordings/ai/embeddings_runner.py +++ b/ee/session_recordings/ai/embeddings_runner.py @@ -145,6 +145,7 @@ def run(self, items: List[Any], embeddings_preparation: type[EmbeddingPreparatio "session_id": session_id, "embeddings": embeddings, "source_type": source_type, + "input": input, } ) # we don't want to fail the whole batch if only a single recording fails @@ -198,7 +199,7 @@ def _num_tokens_for_input(self, string: str) -> int: def _flush_embeddings_to_clickhouse(self, embeddings: List[Dict[str, Any]], source_type: str) -> None: try: sync_execute( - "INSERT INTO session_replay_embeddings (session_id, team_id, embeddings, source_type) VALUES", + "INSERT INTO session_replay_embeddings (session_id, team_id, embeddings, source_type, input) VALUES", embeddings, ) SESSION_EMBEDDINGS_WRITTEN_TO_CLICKHOUSE.labels(source_type=source_type).inc(len(embeddings))