Fix issue where cached JSON responses weren't being parsed (#8)

python/graphrag/graphrag/llm/openai/factories.py

@@ -16,6 +16,7 @@
     OnCacheActionFn,
 )
 
+from .json_parsing_llm import JsonParsingLLM
 from .openai_chat_llm import OpenAIChatLLM
 from .openai_completion_llm import OpenAICompletionLLM
 from .openai_configuration import OpenAIConfiguration
@@ -52,7 +53,8 @@ def create_openai_chat_llm(
     if cache is not None:
         result = _cached(result, config, operation, cache, on_cache_hit, on_cache_miss)
     result = OpenAIHistoryTrackingLLM(result)
-    return OpenAITokenReplacingLLM(result)
+    result = OpenAITokenReplacingLLM(result)
+    return JsonParsingLLM(result)
 
 
 def create_openai_completion_llm(

python/graphrag/graphrag/llm/openai/json_parsing_llm.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2024 Microsoft Corporation. All rights reserved.
+
+"""An LLM that unpacks cached JSON responses."""
+
+import json
+from typing import cast
+
+from typing_extensions import Unpack
+
+from graphrag.llm.types import (
+    LLM,
+    CompletionInput,
+    CompletionLLM,
+    CompletionOutput,
+    LLMInput,
+    LLMOutput,
+)
+
+
+class JsonParsingLLM(LLM[CompletionInput, CompletionOutput]):
+    """An OpenAI History-Tracking LLM."""
+
+    _delegate: CompletionLLM
+
+    def __init__(self, delegate: CompletionLLM):
+        self._delegate = delegate
+
+    async def __call__(
+        self,
+        input: CompletionInput,
+        **kwargs: Unpack[LLMInput],
+    ) -> LLMOutput[CompletionOutput]:
+        """Call the LLM with the input and kwargs."""
+        result = await self._delegate(input, **kwargs)
+        if kwargs.get("json") and result.json is None and result.output is not None:
+            result.json = cast(dict, json.loads(result.output))
+        return result