From d7709ae01d0e5de6c5f17621ecd0906acf8f012b Mon Sep 17 00:00:00 2001
From: Alonso Guevara <alonsog@microsoft.com>
Date: Wed, 21 May 2025 18:33:44 -0600
Subject: [PATCH 1/6] Add full_response to llm provider output

---
 .../language_model/providers/fnllm/models.py     | 11 +++++++++--
 graphrag/language_model/response/base.py         | 16 ++++++++++++++++
 tests/integration/language_model/test_factory.py |  8 +++++++-
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/graphrag/language_model/providers/fnllm/models.py b/graphrag/language_model/providers/fnllm/models.py
index fda91c96ba..806bef50d1 100644
--- a/graphrag/language_model/providers/fnllm/models.py
+++ b/graphrag/language_model/providers/fnllm/models.py
@@ -83,7 +83,11 @@ async def achat(
         else:
             response = await self.model(prompt, history=history, **kwargs)
         return BaseModelResponse(
-            output=BaseModelOutput(content=response.output.content),
+            output=BaseModelOutput(
+                content=response.output.content,
+                full_response=response.output.raw_model.to_dict(),
+            ),
+            raw_model_response=response.raw_model,
             parsed_response=response.parsed_json,
             history=response.history,
             cache_hit=response.cache_hit,
@@ -282,7 +286,10 @@ async def achat(
         else:
             response = await self.model(prompt, history=history, **kwargs)
         return BaseModelResponse(
-            output=BaseModelOutput(content=response.output.content),
+            output=BaseModelOutput(
+                content=response.output.content,
+                full_response=response.output.raw_model.to_dict(),
+            ),
             parsed_response=response.parsed_json,
             history=response.history,
             cache_hit=response.cache_hit,
diff --git a/graphrag/language_model/response/base.py b/graphrag/language_model/response/base.py
index abc17dda38..39e2603ea6 100644
--- a/graphrag/language_model/response/base.py
+++ b/graphrag/language_model/response/base.py
@@ -18,6 +18,11 @@ def content(self) -> str:
         """Return the textual content of the output."""
         ...
 
+    @property
+    def full_response(self) -> dict[str, Any] | None:
+        """Return the complete JSON response returned by the model."""
+        ...
+
 
 class ModelResponse(Protocol, Generic[T]):
     """Protocol for LLM response."""
@@ -37,12 +42,21 @@ def history(self) -> list:
         """Return the history of the response."""
         ...
 
+    @property
+    def raw_model_response(self) -> Any:
+        """Return the raw model response."""
+        ...
+
 
 class BaseModelOutput(BaseModel):
     """Base class for LLM output."""
 
     content: str = Field(..., description="The textual content of the output.")
     """The textual content of the output."""
+    full_response: dict[str, Any] | None = Field(
+        None, description="The complete JSON response returned by the LLM provider."
+    )
+    """The complete JSON response returned by the LLM provider."""
 
 
 class BaseModelResponse(BaseModel, Generic[T]):
@@ -52,6 +66,8 @@ class BaseModelResponse(BaseModel, Generic[T]):
     """"""
     parsed_response: T | None = None
     """Parsed response."""
+    raw_model_response: Any = None
+    """Raw model response."""
     history: list[Any] = Field(default_factory=list)
     """History of the response."""
     tool_calls: list = Field(default_factory=list)
diff --git a/tests/integration/language_model/test_factory.py b/tests/integration/language_model/test_factory.py
index e25e4e246a..af503265d6 100644
--- a/tests/integration/language_model/test_factory.py
+++ b/tests/integration/language_model/test_factory.py
@@ -33,7 +33,11 @@ async def achat(
         def chat(
             self, prompt: str, history: list | None = None, **kwargs: Any
         ) -> ModelResponse:
-            return BaseModelResponse(output=BaseModelOutput(content="content"))
+            return BaseModelResponse(
+                output=BaseModelOutput(
+                    content="content", full_response={"key": "value"}
+                )
+            )
 
         async def achat_stream(
             self, prompt: str, history: list | None = None, **kwargs: Any
@@ -49,9 +53,11 @@ def chat_stream(
     assert isinstance(model, CustomChatModel)
     response = await model.achat("prompt")
     assert response.output.content == "content"
+    assert response.output.full_response is None
 
     response = model.chat("prompt")
     assert response.output.content == "content"
+    assert response.output.full_response == {"key": "value"}
 
 
 async def test_create_custom_embedding_llm():

From cd08153655a1b9c85ac4c9f65feca34cb6a45c5b Mon Sep 17 00:00:00 2001
From: Alonso Guevara <alonsog@microsoft.com>
Date: Wed, 21 May 2025 18:37:51 -0600
Subject: [PATCH 2/6] Semver

---
 .semversioner/next-release/patch-20250522003454958473.json | 4 ++++
 1 file changed, 4 insertions(+)
 create mode 100644 .semversioner/next-release/patch-20250522003454958473.json

diff --git a/.semversioner/next-release/patch-20250522003454958473.json b/.semversioner/next-release/patch-20250522003454958473.json
new file mode 100644
index 0000000000..7a7feaf072
--- /dev/null
+++ b/.semversioner/next-release/patch-20250522003454958473.json
@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Add full llm response to LLM PRovider output"
+}

From 33e175c1e2358ff82b967b9bbc551d54e4502668 Mon Sep 17 00:00:00 2001
From: Alonso Guevara <alonsog@microsoft.com>
Date: Wed, 21 May 2025 18:46:31 -0600
Subject: [PATCH 3/6] Small leftover cleanup

---
 graphrag/language_model/providers/fnllm/models.py | 1 -
 graphrag/language_model/response/base.py          | 7 -------
 2 files changed, 8 deletions(-)

diff --git a/graphrag/language_model/providers/fnllm/models.py b/graphrag/language_model/providers/fnllm/models.py
index 806bef50d1..059b0412ca 100644
--- a/graphrag/language_model/providers/fnllm/models.py
+++ b/graphrag/language_model/providers/fnllm/models.py
@@ -87,7 +87,6 @@ async def achat(
                 content=response.output.content,
                 full_response=response.output.raw_model.to_dict(),
             ),
-            raw_model_response=response.raw_model,
             parsed_response=response.parsed_json,
             history=response.history,
             cache_hit=response.cache_hit,
diff --git a/graphrag/language_model/response/base.py b/graphrag/language_model/response/base.py
index 39e2603ea6..178259c4b7 100644
--- a/graphrag/language_model/response/base.py
+++ b/graphrag/language_model/response/base.py
@@ -42,11 +42,6 @@ def history(self) -> list:
         """Return the history of the response."""
         ...
 
-    @property
-    def raw_model_response(self) -> Any:
-        """Return the raw model response."""
-        ...
-
 
 class BaseModelOutput(BaseModel):
     """Base class for LLM output."""
@@ -66,8 +61,6 @@ class BaseModelResponse(BaseModel, Generic[T]):
     """"""
     parsed_response: T | None = None
     """Parsed response."""
-    raw_model_response: Any = None
-    """Raw model response."""
     history: list[Any] = Field(default_factory=list)
     """History of the response."""
     tool_calls: list = Field(default_factory=list)

From 9007ba1914b0dd66cedb7e36b518d5c92b83e742 Mon Sep 17 00:00:00 2001
From: Alonso Guevara <alonsog@microsoft.com>
Date: Wed, 21 May 2025 18:56:19 -0600
Subject: [PATCH 4/6] Add pyi to suppress Pyright errors. full_content is
 optional

---
 graphrag/language_model/response/base.pyi | 15 +++++++++++++++
 1 file changed, 15 insertions(+)
 create mode 100644 graphrag/language_model/response/base.pyi

diff --git a/graphrag/language_model/response/base.pyi b/graphrag/language_model/response/base.pyi
new file mode 100644
index 0000000000..e1503874e5
--- /dev/null
+++ b/graphrag/language_model/response/base.pyi
@@ -0,0 +1,15 @@
+# Copyright (c) 2025 Microsoft Corporation.
+# Licensed under the MIT License
+
+from typing import Any
+from pydantic import BaseModel
+
+class BaseModelOutput(BaseModel):
+    content: str
+    full_response: dict[str, Any] | None = None
+
+    def __init__(
+        self,
+        content: str,
+        full_response: dict[str, Any] | None = None,
+    ) -> None: ...
\ No newline at end of file

From a455e5a10a071e48864cc14429d47cf8ecab3bda Mon Sep 17 00:00:00 2001
From: Alonso Guevara <alonsog@microsoft.com>
Date: Wed, 21 May 2025 18:58:47 -0600
Subject: [PATCH 5/6] Format

---
 graphrag/language_model/response/base.pyi | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/graphrag/language_model/response/base.pyi b/graphrag/language_model/response/base.pyi
index e1503874e5..662bc3bfbb 100644
--- a/graphrag/language_model/response/base.pyi
+++ b/graphrag/language_model/response/base.pyi
@@ -2,6 +2,7 @@
 # Licensed under the MIT License
 
 from typing import Any
+
 from pydantic import BaseModel
 
 class BaseModelOutput(BaseModel):
@@ -12,4 +13,4 @@ class BaseModelOutput(BaseModel):
         self,
         content: str,
         full_response: dict[str, Any] | None = None,
-    ) -> None: ...
\ No newline at end of file
+    ) -> None: ...

From fed8f6aec8c01e50da4516369a2fdb8f1563b18d Mon Sep 17 00:00:00 2001
From: Alonso Guevara <alonsog@microsoft.com>
Date: Wed, 21 May 2025 19:13:29 -0600
Subject: [PATCH 6/6] Add missing stubs

---
 graphrag/language_model/response/base.pyi | 38 +++++++++++++++++++++--
 1 file changed, 36 insertions(+), 2 deletions(-)

diff --git a/graphrag/language_model/response/base.pyi b/graphrag/language_model/response/base.pyi
index 662bc3bfbb..7a33b0a304 100644
--- a/graphrag/language_model/response/base.pyi
+++ b/graphrag/language_model/response/base.pyi
@@ -1,16 +1,50 @@
 # Copyright (c) 2025 Microsoft Corporation.
 # Licensed under the MIT License
 
-from typing import Any
+from typing import Any, Generic, Protocol, TypeVar
 
 from pydantic import BaseModel
 
+_T = TypeVar("_T", bound=BaseModel, covariant=True)
+
+class ModelOutput(Protocol):
+    @property
+    def content(self) -> str: ...
+    @property
+    def full_response(self) -> dict[str, Any] | None: ...
+
+class ModelResponse(Protocol, Generic[_T]):
+    @property
+    def output(self) -> ModelOutput: ...
+    @property
+    def parsed_response(self) -> _T | None: ...
+    @property
+    def history(self) -> list[Any]: ...
+
 class BaseModelOutput(BaseModel):
     content: str
-    full_response: dict[str, Any] | None = None
+    full_response: dict[str, Any] | None
 
     def __init__(
         self,
         content: str,
         full_response: dict[str, Any] | None = None,
     ) -> None: ...
+
+class BaseModelResponse(BaseModel, Generic[_T]):
+    output: BaseModelOutput
+    parsed_response: _T | None
+    history: list[Any]
+    tool_calls: list[Any]
+    metrics: Any | None
+    cache_hit: bool | None
+
+    def __init__(
+        self,
+        output: BaseModelOutput,
+        parsed_response: _T | None = None,
+        history: list[Any] = ...,  # default provided by Pydantic
+        tool_calls: list[Any] = ...,  # default provided by Pydantic
+        metrics: Any | None = None,
+        cache_hit: bool | None = None,
+    ) -> None: ...