diff --git a/Sources/LangChain/llms/Ollama.swift b/Sources/LangChain/llms/Ollama.swift
index 05e56b3..955325c 100644
--- a/Sources/LangChain/llms/Ollama.swift
+++ b/Sources/LangChain/llms/Ollama.swift
@@ -100,6 +100,7 @@ public class Ollama: LLM {
             return LLMResult()
         }
         let llmResponse = try JSONDecoder().decode(GenerateResponse.self, from: data)
+        context = llmResponse.context
         return LLMResult(llm_output: llmResponse.response)
     }