graphrag/tests/unit/llm/base/test_caching_llm.py
Chris Trevino f5c9c2bee0
Add History input to cache-key, cache data (#736)
* Update caching llm to use history inputs

* formatting

* linting

* update glean sections to have continuous history
2024-07-26 09:26:37 -07:00

74 lines
2.2 KiB
Python

# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License
"""Caching LLM Tests."""
import asyncio
from typing import Any, cast
from graphrag.llm import CompletionLLM, LLMOutput
from graphrag.llm.base.caching_llm import CachingLLM
from graphrag.llm.openai.openai_history_tracking_llm import OpenAIHistoryTrackingLLM
from graphrag.llm.types import LLMCache
class TestCache(LLMCache):
def __init__(self):
self.cache = {}
async def has(self, key: str) -> bool:
return key in self.cache
async def get(self, key: str) -> dict | None:
entry = self.cache.get(key)
return entry["result"] if entry else None
async def set(
self, key: str, value: str, debug_data: dict[str, Any] | None = None
) -> None:
self.cache[key] = {"result": value, **(debug_data or {})}
async def mock_responder(input: str, **kwargs: dict) -> LLMOutput:
await asyncio.sleep(0.0001)
return LLMOutput(output=f"response to [{input}]")
def throwing_responder(input: str, **kwargs: dict) -> LLMOutput:
raise ValueError
mock_responder_llm = cast(CompletionLLM, mock_responder)
throwing_llm = cast(CompletionLLM, throwing_responder)
async def test_caching_llm() -> None:
"""Test a composite LLM."""
llm = CachingLLM(
mock_responder_llm, llm_parameters={}, operation="test", cache=TestCache()
)
response = await llm("input 1")
assert response.output == "response to [input 1]"
llm.set_delegate(throwing_llm)
response = await llm("input 1")
assert response.output == "response to [input 1]"
async def test_composite_llm() -> None:
"""Test a composite LLM."""
caching = CachingLLM(
mock_responder_llm, llm_parameters={}, operation="test", cache=TestCache()
)
llm = OpenAIHistoryTrackingLLM(caching)
response = await llm("input 1")
history: list[dict] = cast(list[dict], response.history)
assert len(history) == 2
response = await llm("input 1")
history: list[dict] = cast(list[dict], response.history)
assert len(history) == 2
response = await llm("input 2", history=history)
history: list[dict] = cast(list[dict], response.history)
assert len(history) == 4