taskweaver/llm/mock.py

import hashlib
import json
import os
import random
import time
from dataclasses import dataclass
from typing import Any, Dict, Generator, List, Literal, Optional

import yaml
from injector import inject

from taskweaver.llm.util import ChatMessageRoleType, ChatMessageType, format_chat_message

from .base import CompletionService, EmbeddingService, LLMServiceConfig

MockServiceModeType = Literal[
    "fixed",
    "record_only",
    "playback_only",
    "playback_or_record",
]


class MockApiServiceConfig(LLMServiceConfig):
    def _configure(self) -> None:
        self._set_name("mock")

        mock_mode = self._get_enum(
            "mode",
            options=["fixed", "record_only", "playback_only", "playback_or_record"],
            default="playback_or_record",
        )

        assert mock_mode in [
            "fixed",
            "record_only",
            "playback_only",
            "playback_or_record",
        ]
        self.mode: MockServiceModeType = mock_mode  # type: ignore

        self.fixed_chat_responses: str = self._get_str(
            "fixed_chat_responses",
            json.dumps(format_chat_message("assistant", "Hello!")),
        )

        self.fixed_embedding_responses: str = self._get_str(
            "fixed_embedding_responses",
            json.dumps([[0.0] * 64]),
        )

        self.cache_path: str = self._get_path(
            "cache_path",
            os.path.join(self.src.app_base_path, "cache", "mock.yaml"),
        )

        # split the chat completion response into chunks and delay each chunk by this amount
        # if negative, return the whole response at once
        self.playback_delay: float = self._get_float(
            "playback_delay",
            0.05,
        )


@dataclass
class MockCacheEntry:
    value: str
    query: str
    created_at: float
    last_accessed_at: float


class MockCacheStore:
    def __init__(self, path: str):
        self.path = path
        self.completion_store: Dict[str, MockCacheEntry] = {}
        self.embedding_store: Dict[str, MockCacheEntry] = {}

        if os.path.exists(self.path):
            self._init_from_disk()

    def get_completion(self, query: List[ChatMessageType]) -> Optional[ChatMessageType]:
        serialized_query = self._serialize_completion_query(query)
        serialized_value = self._get_from_store(self.completion_store, serialized_query)
        if serialized_value is None:
            return None
        return self._deserialize_completion_response(serialized_value)

    def get_embedding(self, query: str) -> Optional[List[float]]:
        serialized_query = self._serialize_embedding_query(query)
        serialized_value = self._get_from_store(self.embedding_store, serialized_query)
        if serialized_value is None:
            return None
        return self._deserialize_embedding_response(serialized_value)

    def _get_from_store(
        self,
        store: Dict[str, MockCacheEntry],
        query: str,
    ) -> Optional[str]:
        key = self._query_to_key(query)
        if key in store:
            entry = store[key]
            entry.last_accessed_at = time.time()
            return entry.value
        return None

    def set_completion(
        self,
        query: List[ChatMessageType],
        value: ChatMessageType,
    ) -> None:
        serialized_query = self._serialize_completion_query(query)
        serialized_value = self._serialize_completion_response(value)
        self._set_to_store(self.completion_store, serialized_query, serialized_value)

    def set_embedding(self, query: str, value: List[float]) -> None:
        serialized_query = self._serialize_embedding_query(query)
        serialized_value = self._serialize_embedding_response(value)
        self._set_to_store(self.embedding_store, serialized_query, serialized_value)

    def _set_to_store(
        self,
        store: Dict[str, MockCacheEntry],
        query: str,
        value: str,
    ) -> None:
        key = self._query_to_key(query)
        store[key] = MockCacheEntry(
            value=value,
            query=query,
            created_at=time.time(),
            last_accessed_at=time.time(),
        )
        self._save_to_disk()

    def _serialize_completion_query(self, query: List[ChatMessageType]) -> str:
        return "\n".join([self._serialize_completion_response(x) for x in query])

    def _serialize_completion_response(self, response: ChatMessageType) -> str:
        return f"{response['role']}:{response['content']}"

    def _deserialize_completion_response(self, response: str) -> ChatMessageType:
        segment = response.split(":", 1)
        role = segment[0] if len(segment) > 0 else "assistant"
        if role not in ["assistant", "user", "system"]:
            raise ValueError(f"Invalid role {role}")
        content = segment[1] if len(segment) > 1 else ""
        return format_chat_message(role, content)  # type: ignore

    def _serialize_embedding_query(self, query: str) -> str:
        return query

    def _serialize_embedding_response(self, response: List[float]) -> str:
        return ",".join([str(x) for x in response])

    def _deserialize_embedding_response(self, response: str) -> List[float]:
        return [float(x) for x in response.split(",")]

    def _query_to_key(self, query: str) -> str:
        return hashlib.md5(query.encode("utf-8")).hexdigest()

    def _init_from_disk(self):
        try:
            cache = yaml.safe_load(open(self.path, "r"))
        except Exception as e:
            print(f"Error loading cache file {self.path}: {e}")
            return

        try:
            completion_store = cache["completion_store"]
            for key, value in completion_store.items():
                try:
                    self.completion_store[key] = MockCacheEntry(**value)
                except Exception as e:
                    print(f"Error loading cache entry {key}: {e}")
        except Exception as e:
            print(f"Error loading completion store: {e}")

        try:
            embedding_store = cache["embedding_store"]
            for key, value in embedding_store.items():
                try:
                    self.embedding_store[key] = MockCacheEntry(**value)
                except Exception as e:
                    print(f"Error loading cache entry {key}: {e}")
        except Exception as e:
            print(f"Error loading embedding store: {e}")

    def _save_to_disk(self):
        # TODO: postpone immediate update and periodically save to disk
        try:
            yaml.safe_dump(
                {
                    "completion_store": {k: v.__dict__ for k, v in self.completion_store.items()},
                    "embedding_store": {k: v.__dict__ for k, v in self.embedding_store.items()},
                },
                open(self.path, "w"),
            )
        except Exception as e:
            print(f"Error saving cache file {self.path}: {e}")


class MockApiService(CompletionService, EmbeddingService):
    @inject
    def __init__(self, config: MockApiServiceConfig):
        self.config = config
        self.base_completion_service: Optional[CompletionService] = None
        self.base_embedding_service: Optional[EmbeddingService] = None
        self.cache = MockCacheStore(self.config.cache_path)

    def set_base_completion_service(
        self,
        base_completion_service: Optional[CompletionService],
    ):
        self.base_completion_service = base_completion_service

    def set_base_embedding_service(
        self,
        base_embedding_service: Optional[EmbeddingService],
    ):
        self.base_embedding_service = base_embedding_service

    def chat_completion(
        self,
        messages: List[ChatMessageType],
        use_backup_engine: bool = False,
        stream: bool = True,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        top_p: Optional[float] = None,
        stop: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Generator[ChatMessageType, None, None]:
        if self.config.mode == "fixed":
            return self._get_from_fixed_completion()

        cached_value = self.cache.get_completion(messages)

        # playback
        if cached_value is None:
            if self.config.mode == "playback_only":
                raise Exception("No cached value found")
        else:
            if self.config.mode != "record_only":
                return self._get_from_playback_completion(cached_value)

        # record
        def get_from_base() -> Generator[ChatMessageType, None, None]:
            if self.base_completion_service is None:
                raise Exception("base_completion_service is not set")
            new_value = format_chat_message("assistant", "")
            for chunk in self.base_completion_service.chat_completion(
                messages,
                use_backup_engine,
                stream,
                temperature,
                max_tokens,
                top_p,
                stop,
                **kwargs,
            ):
                new_value["role"] = chunk["role"]
                new_value["content"] += chunk["content"]
                yield chunk

            self.cache.set_completion(messages, new_value)

        return get_from_base()

    def get_embeddings(self, strings: List[str]) -> List[List[float]]:
        if self.config.mode == "fixed":
            return [self._get_from_fixed_embedding() for _ in strings]

        cached_values = [self.cache.get_embedding(x) for x in strings]

        cache_missed_values = [strings[i] for i, v in enumerate(cached_values) if v is None]

        if len(cache_missed_values) > 0:
            if self.config.mode == "playback_only":
                raise Exception("Not all cached values found")

        if self.base_embedding_service is None:
            raise Exception("base_embedding_service is not set")

        new_values = self.base_embedding_service.get_embeddings(cache_missed_values)

        cache_missed_values_index = 0
        result_values: List[List[float]] = []
        for i, v in enumerate(cached_values):
            if v is None:
                self.cache.set_embedding(
                    strings[i],
                    new_values[cache_missed_values_index],
                )
                result_values.append(new_values[cache_missed_values_index])
                cache_missed_values_index += 1
            else:
                result_values.append(v)

        return result_values

    def _get_from_fixed_completion(
        self,
    ) -> Generator[ChatMessageType, None, None]:
        fixed_responses: ChatMessageType = json.loads(
            self.config.fixed_chat_responses,
        )
        return self._get_from_playback_completion(fixed_responses)

    def _get_from_fixed_embedding(
        self,
    ) -> List[float]:
        fixed_responses: List[float] = json.loads(self.config.fixed_embedding_responses)
        return fixed_responses

    def _get_from_playback_completion(
        self,
        cached_value: ChatMessageType,
    ) -> Generator[ChatMessageType, None, None]:
        if self.config.playback_delay < 0:
            yield cached_value
            return

        role: ChatMessageRoleType = cached_value["role"]  # type: ignore
        content = cached_value["content"]
        cur_pos = 0
        while cur_pos < len(content):
            chunk_size = random.randint(3, 20)
            next_pos = min(cur_pos + chunk_size, len(content))
            yield format_chat_message(role, content[cur_pos:next_pos])
            cur_pos = next_pos
            time.sleep(self.config.playback_delay)